#region License Information
/* HeuristicLab
 * Copyright (C) Heuristic and Evolutionary Algorithms Laboratory (HEAL)
 *
 * This file is part of HeuristicLab.
 *
 * HeuristicLab is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * HeuristicLab is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with HeuristicLab. If not, see <http://www.gnu.org/licenses/>.
 */
#endregion

using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using HEAL.Attic;
using HeuristicLab.Common;
using HeuristicLab.Core;
using HeuristicLab.Data;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.Parameters;
using HeuristicLab.Random;

namespace HeuristicLab.Problems.DataAnalysis.Symbolic.Regression {
  [Item("NMSE Evaluator with shape-constraints (single-objective)", "Calculates NMSE of a symbolic regression solution and checks constraints. The fitness is a combination of NMSE and constraint violations.")]
  [StorableType("27473973-DD8D-4375-997D-942E2280AE8E")]
  public class NMSESingleObjectiveConstraintsEvaluator : SymbolicRegressionSingleObjectiveEvaluator {
    #region Parameter/Properties

    private const string OptimizeParametersParameterName = "OptimizeParameters";
    private const string ParameterOptimizationIterationsParameterName = "ParameterOptimizationIterations";
    private const string UseSoftConstraintsParameterName = "UseSoftConstraintsEvaluation";
    private const string BoundsEstimatorParameterName = "BoundsEstimator";
    private const string PenaltyFactorParameterName = "PenaltyFactor";


    public IFixedValueParameter<BoolValue> OptimizerParametersParameter =>
      (IFixedValueParameter<BoolValue>)Parameters[OptimizeParametersParameterName];

    public IFixedValueParameter<IntValue> ParameterOptimizationIterationsParameter =>
      (IFixedValueParameter<IntValue>)Parameters[ParameterOptimizationIterationsParameterName];

    public IFixedValueParameter<BoolValue> UseSoftConstraintsParameter =>
      (IFixedValueParameter<BoolValue>)Parameters[UseSoftConstraintsParameterName];

    public IValueParameter<IBoundsEstimator> BoundsEstimatorParameter =>
      (IValueParameter<IBoundsEstimator>)Parameters[BoundsEstimatorParameterName];
    public IFixedValueParameter<DoubleValue> PenaltyFactorParameter =>
      (IFixedValueParameter<DoubleValue>)Parameters[PenaltyFactorParameterName];



    public bool OptimizeParameters {
      get => OptimizerParametersParameter.Value.Value;
      set => OptimizerParametersParameter.Value.Value = value;
    }

    public int ParameterOptimizationIterations {
      get => ParameterOptimizationIterationsParameter.Value.Value;
      set => ParameterOptimizationIterationsParameter.Value.Value = value;
    }

    public bool UseSoftConstraints {
      get => UseSoftConstraintsParameter.Value.Value;
      set => UseSoftConstraintsParameter.Value.Value = value;
    }

    public IBoundsEstimator BoundsEstimator {
      get => BoundsEstimatorParameter.Value;
      set => BoundsEstimatorParameter.Value = value;
    }

    public double PenalityFactor {
      get => PenaltyFactorParameter.Value.Value;
      set => PenaltyFactorParameter.Value.Value = value;
    }

    public override bool Maximization => false; // NMSE is minimized

    #endregion

    #region Constructors/Cloning

    [StorableConstructor]
    protected NMSESingleObjectiveConstraintsEvaluator(StorableConstructorFlag _) : base(_) { }

    protected NMSESingleObjectiveConstraintsEvaluator(
      NMSESingleObjectiveConstraintsEvaluator original, Cloner cloner) : base(original, cloner) { }

    public NMSESingleObjectiveConstraintsEvaluator() {
      Parameters.Add(new FixedValueParameter<BoolValue>(OptimizeParametersParameterName,
        "Define whether optimization of parameters is active or not (default: false).", new BoolValue(false)));
      Parameters.Add(new FixedValueParameter<IntValue>(ParameterOptimizationIterationsParameterName,
        "Define how many parameter optimization steps should be performed (default: 10).", new IntValue(10)));
      Parameters.Add(new FixedValueParameter<BoolValue>(UseSoftConstraintsParameterName,
        "Define whether the constraints are penalized by soft or hard constraints (default: false).", new BoolValue(false)));
      Parameters.Add(new ValueParameter<IBoundsEstimator>(BoundsEstimatorParameterName,
        "The estimator which is used to estimate output ranges of models (default: interval arithmetic).", new IntervalArithBoundsEstimator()));
      Parameters.Add(new FixedValueParameter<DoubleValue>(PenaltyFactorParameterName,
        "Punishment factor for constraint violations for soft constraint handling (fitness = NMSE + penaltyFactor * avg(violations)) (default: 1.0)", new DoubleValue(1.0)));
    }

    [StorableHook(HookType.AfterDeserialization)]
    private void AfterDeserialization() { }

    public override IDeepCloneable Clone(Cloner cloner) {
      return new NMSESingleObjectiveConstraintsEvaluator(this, cloner);
    }

    #endregion

    public override IOperation InstrumentedApply() {
      var rows = GenerateRowsToEvaluate();
      var tree = SymbolicExpressionTreeParameter.ActualValue;
      var problemData = ProblemDataParameter.ActualValue;
      var interpreter = SymbolicDataAnalysisTreeInterpreterParameter.ActualValue;
      var estimationLimits = EstimationLimitsParameter.ActualValue;
      var applyLinearScaling = ApplyLinearScalingParameter.ActualValue.Value;
      var random = RandomParameter.ActualValue;

      if (OptimizeParameters) {
        SymbolicRegressionParameterOptimizationEvaluator.OptimizeParameters(interpreter, tree, problemData, rows,
          false, ParameterOptimizationIterations, true,
          estimationLimits.Lower, estimationLimits.Upper);
      } else {
        if (applyLinearScaling) {
          var rootNode = new ProgramRootSymbol().CreateTreeNode();
          var startNode = new StartSymbol().CreateTreeNode();
          var offset = tree.Root.GetSubtree(0) //Start
                                .GetSubtree(0); //Offset
          var scaling = offset.GetSubtree(0);

          //Check if tree contains offset and scaling nodes
          if (!(offset.Symbol is Addition) || !(scaling.Symbol is Multiplication))
            throw new ArgumentException($"{ItemName} can only be used with LinearScalingGrammar.");

          var t = (ISymbolicExpressionTreeNode)scaling.GetSubtree(0).Clone();
          rootNode.AddSubtree(startNode);
          startNode.AddSubtree(t);
          var newTree = new SymbolicExpressionTree(rootNode);

          //calculate alpha and beta for scaling
          var estimatedValues = interpreter.GetSymbolicExpressionTreeValues(newTree, problemData.Dataset, rows);

          var targetValues = problemData.Dataset.GetDoubleValues(problemData.TargetVariable, rows);
          OnlineLinearScalingParameterCalculator.Calculate(estimatedValues, targetValues, out var alpha, out var beta,
            out var errorState);

          if (errorState == OnlineCalculatorError.None) {
            //Set alpha and beta to the scaling nodes from ia grammar
            var offsetParameter = offset.GetSubtree(1) as NumberTreeNode;
            offsetParameter.Value = alpha;
            var scalingParameter = scaling.GetSubtree(1) as NumberTreeNode;
            scalingParameter.Value = beta;
          }
        } // else: alpha and beta are evolved
      }

      var quality = Calculate(interpreter, tree, estimationLimits.Lower, estimationLimits.Upper, problemData, rows,
        BoundsEstimator, random, UseSoftConstraints, PenalityFactor);
      QualityParameter.ActualValue = new DoubleValue(quality);

      return base.InstrumentedApply();
    }

    public static double Calculate(
      ISymbolicDataAnalysisExpressionTreeInterpreter interpreter,
      ISymbolicExpressionTree tree,
      double lowerEstimationLimit, double upperEstimationLimit,
      IRegressionProblemData problemData, IEnumerable<int> rows,
      IBoundsEstimator estimator, IRandom random,
      bool useSoftConstraints = false, double penaltyFactor = 1.0) {

      var trainingEstimatedValues = interpreter.GetSymbolicExpressionTreeValues(tree, problemData.Dataset, rows);
      var targetValues = problemData.Dataset.GetDoubleValues(problemData.TargetVariable, rows);

      var trainingBoundedEstimatedValues = trainingEstimatedValues.LimitToRange(lowerEstimationLimit, upperEstimationLimit);
      var nmse = OnlineNormalizedMeanSquaredErrorCalculator.Calculate(targetValues, trainingBoundedEstimatedValues,
        out var errorState);

      if (errorState != OnlineCalculatorError.None)
        return double.MaxValue;

      var violations = Enumerable.Empty<double>();
      if (problemData is ShapeConstrainedRegressionProblemData scProbData) {
        violations = CalculateShapeConstraintsViolations(scProbData, tree, interpreter, estimator, random).Select(x => x.Item2);
      }

      if (violations.Any(x => double.IsNaN(x) || double.IsInfinity(x)))
        return double.MaxValue;

      if (useSoftConstraints) {
        if (penaltyFactor < 0.0)
          throw new ArgumentException("The parameter has to be >= 0.0.", nameof(penaltyFactor));

        return nmse + penaltyFactor * violations.Average();
      }
      return violations.Any(x => x > 0.0) ? 1.0 : nmse;
    }

    public static IEnumerable<Tuple<ShapeConstraint, double>> CalculateShapeConstraintsViolations(
      IShapeConstrainedRegressionProblemData problemData, ISymbolicExpressionTree tree, 
      ISymbolicDataAnalysisExpressionTreeInterpreter interpreter, IBoundsEstimator estimator,
      IRandom random) {
      IList<Tuple<ShapeConstraint, double>> violations = new List<Tuple<ShapeConstraint, double>>();
      
      var baseConstraints = problemData.ShapeConstraints.EnabledConstraints;
      var intervalCollection = problemData.VariableRanges;
      var extendedShapeConstraints = problemData.CheckedExtendedConstraints;
      var allEstimatedValues = interpreter.GetSymbolicExpressionTreeValues(tree, problemData.Dataset, problemData.AllIndices);

      foreach (var constraint in baseConstraints)
        violations.Add(Tuple.Create(constraint, IntervalUtil.GetConstraintViolation(constraint, estimator, intervalCollection, tree) * constraint.Weight));

      IDictionary<string, IList> dict = new Dictionary<string, IList>();
      foreach (var varName in problemData.Dataset.VariableNames) {
        if (varName != problemData.TargetVariable)
          dict.Add(varName, problemData.Dataset.GetDoubleValues(varName).ToList());
        else dict.Add(varName, allEstimatedValues.ToList());
      }
      var tmpDataset = new Dataset(dict.Keys, dict.Values);

      foreach (var extendedConstraint in extendedShapeConstraints) {
        var enabledConstraints = extendedConstraint.ShapeConstraints.EnabledConstraints;
        if (enabledConstraints.Any()) {
          var extendedConstraintExprValues = interpreter.GetSymbolicExpressionTreeValues(extendedConstraint.Tree, tmpDataset, problemData.AllIndices);
          var extendedConstraintExprInterval = new Interval(extendedConstraintExprValues.Min(), extendedConstraintExprValues.Max());

          foreach (var constraint in enabledConstraints) {
            if (constraint.Regions.Count > 0) {
              // adapt dataset
              foreach (var kvp in constraint.Regions.GetReadonlyDictionary()) {
                var lb = double.IsNegativeInfinity(kvp.Value.LowerBound) ? double.MinValue : kvp.Value.LowerBound;
                var ub = double.IsPositiveInfinity(kvp.Value.UpperBound) ? double.MaxValue : kvp.Value.UpperBound;

                var vals = Enumerable.Range(0, dict[kvp.Key].Count - 2)
                                     .Select(x => UniformDistributedRandom.NextDouble(random, lb, ub))
                                     .ToList();
                vals.Add(lb);
                vals.Add(ub);
                vals.Sort();
                dict[kvp.Key] = vals;
              }
              // calc again with new regions
              tmpDataset = new Dataset(dict.Keys, dict.Values);
              // calc target again
              allEstimatedValues = interpreter.GetSymbolicExpressionTreeValues(tree, tmpDataset, problemData.AllIndices);
              dict[problemData.TargetVariable] = allEstimatedValues.ToList();
              tmpDataset = new Dataset(dict.Keys, dict.Values);
              extendedConstraintExprValues = interpreter.GetSymbolicExpressionTreeValues(extendedConstraint.Tree, tmpDataset, problemData.AllIndices);
              extendedConstraintExprInterval = new Interval(extendedConstraintExprValues.Min(), extendedConstraintExprValues.Max());
            }
            violations.Add(Tuple.Create(constraint, IntervalUtil.GetIntervalError(constraint.Interval, extendedConstraintExprInterval, constraint.Threshold) * constraint.Weight));
          }
        }
      }
      return violations;
    }

    public override double Evaluate(
      IExecutionContext context, ISymbolicExpressionTree tree, IRegressionProblemData problemData,
      IEnumerable<int> rows) {
      SymbolicDataAnalysisTreeInterpreterParameter.ExecutionContext = context;
      EstimationLimitsParameter.ExecutionContext = context;
      ApplyLinearScalingParameter.ExecutionContext = context;
      RandomParameter.ExecutionContext = context;

      var nmse = Calculate(SymbolicDataAnalysisTreeInterpreterParameter.ActualValue, tree,
        EstimationLimitsParameter.ActualValue.Lower, EstimationLimitsParameter.ActualValue.Upper,
        problemData, rows, BoundsEstimator, RandomParameter.Value, UseSoftConstraints, PenalityFactor);

      SymbolicDataAnalysisTreeInterpreterParameter.ExecutionContext = null;
      EstimationLimitsParameter.ExecutionContext = null;
      ApplyLinearScalingParameter.ExecutionContext = null;
      RandomParameter.ExecutionContext = null;

      return nmse;
    }
  }
}