/*
 * SPDX-License-Identifier: Apache-2.0
 *
 * The OpenSearch Contributors require contributions made to
 * this file be licensed under the Apache-2.0 license or a
 * compatible open source license.
 */

/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you under
 * the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*
 * Modifications Copyright OpenSearch Contributors. See
 * GitHub history for details.
 */

package org.opensearch.search.aggregations.metrics;

import org.opensearch.core.ParseField;
import org.opensearch.core.xcontent.ConstructingObjectParser;
import org.opensearch.core.xcontent.ObjectParser;
import org.opensearch.core.xcontent.ObjectParser.ValueType;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.search.aggregations.metrics.InternalExtendedStats.Fields;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

import static org.opensearch.core.xcontent.ConstructingObjectParser.constructorArg;
import static org.opensearch.core.xcontent.ConstructingObjectParser.optionalConstructorArg;

/**
 * An extended stats agg result parsed between nodes
 *
 * @opensearch.internal
 */
public class ParsedExtendedStats extends ParsedStats implements ExtendedStats {

    protected double sumOfSquares;
    protected double variance;
    protected double variancePopulation;
    protected double varianceSampling;
    protected double stdDeviation;
    protected double stdDeviationPopulation;
    protected double stdDeviationSampling;
    protected double stdDeviationBoundUpper;
    protected double stdDeviationBoundLower;
    protected double stdDeviationBoundUpperPopulation;
    protected double stdDeviationBoundLowerPopulation;
    protected double stdDeviationBoundUpperSampling;
    protected double stdDeviationBoundLowerSampling;

    protected double sum;
    protected double avg;

    @Override
    public String getType() {
        return ExtendedStatsAggregationBuilder.NAME;
    }

    @Override
    public double getSumOfSquares() {
        return sumOfSquares;
    }

    @Override
    public double getVariance() {
        return variance;
    }

    @Override
    public double getVariancePopulation() {
        return variancePopulation;
    }

    @Override
    public double getVarianceSampling() {
        return varianceSampling;
    }

    @Override
    public double getStdDeviation() {
        return stdDeviation;
    }

    @Override
    public double getStdDeviationPopulation() {
        return stdDeviationPopulation;
    }

    @Override
    public double getStdDeviationSampling() {
        return stdDeviationSampling;
    }

    private void setStdDeviationBounds(List<Double> bounds) {
        this.stdDeviationBoundUpper = bounds.get(0);
        this.stdDeviationBoundLower = bounds.get(1);
        this.stdDeviationBoundUpperPopulation = bounds.get(2) == null ? 0 : bounds.get(2);
        this.stdDeviationBoundLowerPopulation = bounds.get(3) == null ? 0 : bounds.get(3);
        this.stdDeviationBoundUpperSampling = bounds.get(4) == null ? 0 : bounds.get(4);
        this.stdDeviationBoundLowerSampling = bounds.get(5) == null ? 0 : bounds.get(5);
    }

    @Override
    public double getStdDeviationBound(Bounds bound) {
        switch (bound) {
            case UPPER:
                return stdDeviationBoundUpper;
            case UPPER_POPULATION:
                return stdDeviationBoundUpperPopulation;
            case UPPER_SAMPLING:
                return stdDeviationBoundUpperSampling;
            case LOWER:
                return stdDeviationBoundLower;
            case LOWER_POPULATION:
                return stdDeviationBoundLowerPopulation;
            case LOWER_SAMPLING:
                return stdDeviationBoundLowerSampling;
            default:
                throw new IllegalArgumentException("Unknown bounds type " + bound);
        }
    }

    private void setStdDeviationBoundsAsString(List<String> boundsAsString) {
        this.valueAsString.put(Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_upper", boundsAsString.get(0));
        this.valueAsString.put(Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_lower", boundsAsString.get(1));
        if (boundsAsString.get(2) != null) {
            this.valueAsString.put(Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_upper_population", boundsAsString.get(2));
        }
        if (boundsAsString.get(3) != null) {
            this.valueAsString.put(Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_lower_population", boundsAsString.get(3));
        }
        if (boundsAsString.get(4) != null) {
            this.valueAsString.put(Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_upper_sampling", boundsAsString.get(4));
        }
        if (boundsAsString.get(5) != null) {
            this.valueAsString.put(Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_lower_sampling", boundsAsString.get(5));
        }
    }

    @Override
    public String getStdDeviationAsString() {
        return valueAsString.getOrDefault(Fields.STD_DEVIATION_AS_STRING, Double.toString(stdDeviation));
    }

    @Override
    public String getStdDeviationPopulationAsString() {
        return valueAsString.getOrDefault(Fields.STD_DEVIATION_POPULATION_AS_STRING, Double.toString(stdDeviationPopulation));
    }

    @Override
    public String getStdDeviationSamplingAsString() {
        return valueAsString.getOrDefault(Fields.STD_DEVIATION_SAMPLING_AS_STRING, Double.toString(stdDeviationSampling));
    }

    @Override
    public String getStdDeviationBoundAsString(Bounds bound) {
        switch (bound) {
            case UPPER:
                return valueAsString.getOrDefault(
                    Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_upper",
                    Double.toString(stdDeviationBoundUpper)
                );
            case UPPER_POPULATION:
                return valueAsString.getOrDefault(
                    Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_upper_population",
                    Double.toString(stdDeviationBoundUpperPopulation)
                );
            case UPPER_SAMPLING:
                return valueAsString.getOrDefault(
                    Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_upper_sampling",
                    Double.toString(stdDeviationBoundUpperSampling)
                );
            case LOWER:
                return valueAsString.getOrDefault(
                    Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_lower",
                    Double.toString(stdDeviationBoundLower)
                );
            case LOWER_POPULATION:
                return valueAsString.getOrDefault(
                    Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_lower_population",
                    Double.toString(stdDeviationBoundLowerPopulation)
                );
            case LOWER_SAMPLING:
                return valueAsString.getOrDefault(
                    Fields.STD_DEVIATION_BOUNDS_AS_STRING + "_lower_sampling",
                    Double.toString(stdDeviationBoundLowerSampling)
                );
            default:
                throw new IllegalArgumentException("Unknown bounds type " + bound);
        }
    }

    @Override
    public String getSumOfSquaresAsString() {
        return valueAsString.getOrDefault(Fields.SUM_OF_SQRS_AS_STRING, Double.toString(sumOfSquares));
    }

    @Override
    public String getVarianceAsString() {
        return valueAsString.getOrDefault(Fields.VARIANCE_AS_STRING, Double.toString(variance));
    }

    @Override
    public String getVariancePopulationAsString() {
        return valueAsString.getOrDefault(Fields.VARIANCE_POPULATION_AS_STRING, Double.toString(variancePopulation));
    }

    @Override
    public String getVarianceSamplingAsString() {
        return valueAsString.getOrDefault(Fields.VARIANCE_SAMPLING_AS_STRING, Double.toString(varianceSampling));
    }

    @Override
    protected XContentBuilder otherStatsToXContent(XContentBuilder builder, Params params) throws IOException {
        if (count != 0) {
            builder.field(Fields.SUM_OF_SQRS, sumOfSquares);
            builder.field(Fields.VARIANCE, getVariance());
            builder.field(Fields.VARIANCE_POPULATION, getVariancePopulation());
            builder.field(Fields.VARIANCE_SAMPLING, getVarianceSampling());
            builder.field(Fields.STD_DEVIATION, getStdDeviation());
            builder.field(Fields.STD_DEVIATION_POPULATION, getStdDeviationPopulation());
            builder.field(Fields.STD_DEVIATION_SAMPLING, getStdDeviationSampling());
            builder.startObject(Fields.STD_DEVIATION_BOUNDS);
            {
                builder.field(Fields.UPPER, getStdDeviationBound(Bounds.UPPER));
                builder.field(Fields.LOWER, getStdDeviationBound(Bounds.LOWER));
                builder.field(Fields.UPPER_POPULATION, getStdDeviationBound(Bounds.UPPER_POPULATION));
                builder.field(Fields.LOWER_POPULATION, getStdDeviationBound(Bounds.LOWER_POPULATION));
                builder.field(Fields.UPPER_SAMPLING, getStdDeviationBound(Bounds.UPPER_SAMPLING));
                builder.field(Fields.LOWER_SAMPLING, getStdDeviationBound(Bounds.LOWER_SAMPLING));
            }
            builder.endObject();
            if (valueAsString.containsKey(Fields.SUM_OF_SQRS_AS_STRING)) {
                builder.field(Fields.SUM_OF_SQRS_AS_STRING, getSumOfSquaresAsString());
                builder.field(Fields.VARIANCE_AS_STRING, getVarianceAsString());
                builder.field(Fields.VARIANCE_POPULATION_AS_STRING, getVariancePopulationAsString());
                builder.field(Fields.VARIANCE_SAMPLING_AS_STRING, getVarianceSamplingAsString());
                builder.field(Fields.STD_DEVIATION_AS_STRING, getStdDeviationAsString());
                builder.field(Fields.STD_DEVIATION_POPULATION_AS_STRING, getStdDeviationPopulationAsString());
                builder.field(Fields.STD_DEVIATION_SAMPLING_AS_STRING, getStdDeviationSamplingAsString());
                builder.startObject(Fields.STD_DEVIATION_BOUNDS_AS_STRING);
                {
                    builder.field(Fields.UPPER, getStdDeviationBoundAsString(Bounds.UPPER));
                    builder.field(Fields.LOWER, getStdDeviationBoundAsString(Bounds.LOWER));
                    builder.field(Fields.UPPER_POPULATION, getStdDeviationBoundAsString(Bounds.UPPER_POPULATION));
                    builder.field(Fields.LOWER_POPULATION, getStdDeviationBoundAsString(Bounds.LOWER_POPULATION));
                    builder.field(Fields.UPPER_SAMPLING, getStdDeviationBoundAsString(Bounds.UPPER_SAMPLING));
                    builder.field(Fields.LOWER_SAMPLING, getStdDeviationBoundAsString(Bounds.LOWER_SAMPLING));
                }
                builder.endObject();
            }
        } else {
            builder.nullField(Fields.SUM_OF_SQRS);
            builder.nullField(Fields.VARIANCE);
            builder.nullField(Fields.VARIANCE_POPULATION);
            builder.nullField(Fields.VARIANCE_SAMPLING);
            builder.nullField(Fields.STD_DEVIATION);
            builder.nullField(Fields.STD_DEVIATION_POPULATION);
            builder.nullField(Fields.STD_DEVIATION_SAMPLING);
            builder.startObject(Fields.STD_DEVIATION_BOUNDS);
            {
                builder.nullField(Fields.UPPER);
                builder.nullField(Fields.LOWER);
                builder.nullField(Fields.UPPER_POPULATION);
                builder.nullField(Fields.LOWER_POPULATION);
                builder.nullField(Fields.UPPER_SAMPLING);
                builder.nullField(Fields.LOWER_SAMPLING);
            }
            builder.endObject();
        }
        return builder;
    }

    private static final ObjectParser<ParsedExtendedStats, Void> PARSER = new ObjectParser<>(
        ParsedExtendedStats.class.getSimpleName(),
        true,
        ParsedExtendedStats::new
    );

    private static final ConstructingObjectParser<List<Double>, Void> STD_BOUNDS_PARSER = new ConstructingObjectParser<>(
        ParsedExtendedStats.class.getSimpleName() + "_STD_BOUNDS",
        true,
        args -> Arrays.stream(args).map(d -> (Double) d).collect(Collectors.toList())
    );

    private static final ConstructingObjectParser<List<String>, Void> STD_BOUNDS_AS_STRING_PARSER = new ConstructingObjectParser<>(
        ParsedExtendedStats.class.getSimpleName() + "_STD_BOUNDS_AS_STRING",
        true,
        args -> Arrays.stream(args).map(d -> (String) d).collect(Collectors.toList())
    );

    static {
        STD_BOUNDS_PARSER.declareField(
            constructorArg(),
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.UPPER),
            ValueType.DOUBLE_OR_NULL
        );
        STD_BOUNDS_PARSER.declareField(
            constructorArg(),
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.LOWER),
            ValueType.DOUBLE_OR_NULL
        );
        STD_BOUNDS_PARSER.declareField(
            optionalConstructorArg(),
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.UPPER_POPULATION),
            ValueType.DOUBLE_OR_NULL
        );
        STD_BOUNDS_PARSER.declareField(
            optionalConstructorArg(),
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.LOWER_POPULATION),
            ValueType.DOUBLE_OR_NULL
        );
        STD_BOUNDS_PARSER.declareField(
            optionalConstructorArg(),
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.UPPER_SAMPLING),
            ValueType.DOUBLE_OR_NULL
        );
        STD_BOUNDS_PARSER.declareField(
            optionalConstructorArg(),
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.LOWER_SAMPLING),
            ValueType.DOUBLE_OR_NULL
        );
        STD_BOUNDS_AS_STRING_PARSER.declareString(constructorArg(), new ParseField(Fields.UPPER));
        STD_BOUNDS_AS_STRING_PARSER.declareString(constructorArg(), new ParseField(Fields.LOWER));
        STD_BOUNDS_AS_STRING_PARSER.declareString(optionalConstructorArg(), new ParseField(Fields.UPPER_POPULATION));
        STD_BOUNDS_AS_STRING_PARSER.declareString(optionalConstructorArg(), new ParseField(Fields.LOWER_POPULATION));
        STD_BOUNDS_AS_STRING_PARSER.declareString(optionalConstructorArg(), new ParseField(Fields.UPPER_SAMPLING));
        STD_BOUNDS_AS_STRING_PARSER.declareString(optionalConstructorArg(), new ParseField(Fields.LOWER_SAMPLING));
        declareExtendedStatsFields(PARSER);
    }

    protected static void declareExtendedStatsFields(ObjectParser<? extends ParsedExtendedStats, Void> objectParser) {
        declareAggregationFields(objectParser);
        declareStatsFields(objectParser);
        objectParser.declareField(
            (agg, value) -> agg.sumOfSquares = value,
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.SUM_OF_SQRS),
            ValueType.DOUBLE_OR_NULL
        );
        objectParser.declareField(
            (agg, value) -> agg.variance = value,
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.VARIANCE),
            ValueType.DOUBLE_OR_NULL
        );
        objectParser.declareField(
            (agg, value) -> agg.variancePopulation = value,
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.VARIANCE_POPULATION),
            ValueType.DOUBLE_OR_NULL
        );
        objectParser.declareField(
            (agg, value) -> agg.varianceSampling = value,
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.VARIANCE_SAMPLING),
            ValueType.DOUBLE_OR_NULL
        );
        objectParser.declareField(
            (agg, value) -> agg.stdDeviation = value,
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.STD_DEVIATION),
            ValueType.DOUBLE_OR_NULL
        );
        objectParser.declareField(
            (agg, value) -> agg.stdDeviationPopulation = value,
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.STD_DEVIATION_POPULATION),
            ValueType.DOUBLE_OR_NULL
        );
        objectParser.declareField(
            (agg, value) -> agg.stdDeviationSampling = value,
            (parser, context) -> parseDouble(parser, 0),
            new ParseField(Fields.STD_DEVIATION_SAMPLING),
            ValueType.DOUBLE_OR_NULL
        );
        objectParser.declareObject(
            ParsedExtendedStats::setStdDeviationBounds,
            STD_BOUNDS_PARSER,
            new ParseField(Fields.STD_DEVIATION_BOUNDS)
        );
        objectParser.declareString(
            (agg, value) -> agg.valueAsString.put(Fields.SUM_OF_SQRS_AS_STRING, value),
            new ParseField(Fields.SUM_OF_SQRS_AS_STRING)
        );
        objectParser.declareString(
            (agg, value) -> agg.valueAsString.put(Fields.VARIANCE_AS_STRING, value),
            new ParseField(Fields.VARIANCE_AS_STRING)
        );
        objectParser.declareString(
            (agg, value) -> agg.valueAsString.put(Fields.VARIANCE_POPULATION_AS_STRING, value),
            new ParseField(Fields.VARIANCE_POPULATION_AS_STRING)
        );
        objectParser.declareString(
            (agg, value) -> agg.valueAsString.put(Fields.VARIANCE_SAMPLING_AS_STRING, value),
            new ParseField(Fields.VARIANCE_SAMPLING_AS_STRING)
        );
        objectParser.declareString(
            (agg, value) -> agg.valueAsString.put(Fields.STD_DEVIATION_AS_STRING, value),
            new ParseField(Fields.STD_DEVIATION_AS_STRING)
        );
        objectParser.declareString(
            (agg, value) -> agg.valueAsString.put(Fields.STD_DEVIATION_POPULATION_AS_STRING, value),
            new ParseField(Fields.STD_DEVIATION_POPULATION_AS_STRING)
        );
        objectParser.declareString(
            (agg, value) -> agg.valueAsString.put(Fields.STD_DEVIATION_SAMPLING_AS_STRING, value),
            new ParseField(Fields.STD_DEVIATION_SAMPLING_AS_STRING)
        );
        objectParser.declareObject(
            ParsedExtendedStats::setStdDeviationBoundsAsString,
            STD_BOUNDS_AS_STRING_PARSER,
            new ParseField(Fields.STD_DEVIATION_BOUNDS_AS_STRING)
        );
    }

    public static ParsedExtendedStats fromXContent(XContentParser parser, final String name) {
        ParsedExtendedStats parsedStats = PARSER.apply(parser, null);
        parsedStats.setName(name);
        return parsedStats;
    }
}