using System;
using System.CodeDom;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using HEAL.Attic;
using HeuristicLab.Common;
using HeuristicLab.Core;
using HeuristicLab.Data;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.Parameters;
using HeuristicLab.Problems.DataAnalysis.Symbolic.Regression;

namespace HeuristicLab.Problems.DataAnalysis.Symbolic {
  [Item("Constraints Evaluator", "Calculates NMSE of a symbolic regression solution with respect to constraints.")]
  [StorableType("27473973-DD8D-4375-997D-942E2280AE8E")]
  public class SymbolicRegressionSingleObjectiveConstraintEvaluator : SymbolicRegressionSingleObjectiveEvaluator {
    #region Parameter/Properties

    private const string UseConstantOptimizationParameterName = "Use Constant Optimization";

    private const string ConstantOptimizationIterationsParameterName = "Constant Optimization Iterations";

    private const string UseSoftConstraintsParameterName = "Use Soft Constraints Evaluation";

    private const string BoundsEstimatorParameterName = "Bounds estimator";

    private const string MaximumPenaltyFactorParamterName = "Maximum Penalty Factor";


    public IFixedValueParameter<BoolValue> UseConstantOptimizationParameter =>
      (IFixedValueParameter<BoolValue>)Parameters[UseConstantOptimizationParameterName];

    public IFixedValueParameter<IntValue> ConstantOptimizationIterationsParameter =>
      (IFixedValueParameter<IntValue>)Parameters[ConstantOptimizationIterationsParameterName];

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

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

    public bool UseConstantOptimization {
      get => UseConstantOptimizationParameter.Value.Value;
      set => UseConstantOptimizationParameter.Value.Value = value;
    }

    public int ConstantOptimizationIterations {
      get => ConstantOptimizationIterationsParameter.Value.Value;
      set => ConstantOptimizationIterationsParameter.Value.Value = value;
    }

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

    public double PenaltyMultiplier {
      get => 1.0;//PenaltyMultiplierParameter.Value.Value;
      //set => PenaltyMultiplierParameter.Value.Value = value;
    }

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

    public double MaximumPenaltyFactor {
      get => MaximumPenaltyFactorParamter.Value.Value;
      set => MaximumPenaltyFactorParamter.Value.Value = value;
    }


    //Use false for maximization, because we try to minimize the NMSE
    public override bool Maximization => false;

    #endregion

    #region Constructors/Cloning

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

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

    public SymbolicRegressionSingleObjectiveConstraintEvaluator() {
      Parameters.Add(new FixedValueParameter<BoolValue>(UseConstantOptimizationParameterName,
        "Define whether constant optimization is active or not.", new BoolValue(false)));
      Parameters.Add(new FixedValueParameter<IntValue>(ConstantOptimizationIterationsParameterName,
        "Define how many constant optimization steps should be performed.", new IntValue(10)));
      Parameters.Add(new FixedValueParameter<BoolValue>(UseSoftConstraintsParameterName,
        "Define whether the constraints are penalized by soft or hard constraints.", new BoolValue(false)));
      Parameters.Add(new ValueParameter<IBoundsEstimator>(BoundsEstimatorParameterName,
        "Select the Boundsestimator.", new IABoundsEstimator()));
      Parameters.Add(new FixedValueParameter<DoubleValue>(MaximumPenaltyFactorParamterName,
        "Specify how hard constraints violations should be punished", new DoubleValue(1.5)));
    }

    [StorableHook(HookType.AfterDeserialization)]
    private void AfterDeserialization() {
      if (!Parameters.ContainsKey(UseConstantOptimizationParameterName)) {
        Parameters.Add(new FixedValueParameter<BoolValue>(UseConstantOptimizationParameterName,
          "Define whether constant optimization is active or not.", new BoolValue(false)));
      }

      if (!Parameters.ContainsKey(ConstantOptimizationIterationsParameterName)) {
        Parameters.Add(new FixedValueParameter<IntValue>(ConstantOptimizationIterationsParameterName,
          "Define how many constant optimization steps should be performed.", new IntValue(10)));
      }

      if (!Parameters.ContainsKey(UseSoftConstraintsParameterName)) {
        Parameters.Add(new FixedValueParameter<BoolValue>(UseSoftConstraintsParameterName,
          "Define whether the constraints are penalized by soft or hard constraints.", new BoolValue(false)));
      }

      if (!Parameters.ContainsKey(BoundsEstimatorParameterName))
        Parameters.Add(new ValueParameter<IBoundsEstimator>(BoundsEstimatorParameterName,
          "Select the Boundsestimator.", new IABoundsEstimator()));

      if (!Parameters.ContainsKey(MaximumPenaltyFactorParamterName)) {
        Parameters.Add(new FixedValueParameter<DoubleValue>(MaximumPenaltyFactorParamterName,
          "Specify how hard constraints violations should be punished", new DoubleValue(1.5)));
      }
    }

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

    #endregion

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

      if (UseConstantOptimization) {
        SymbolicRegressionConstantOptimizationEvaluator.OptimizeConstants(interpreter, solution, problemData, rows,
          false, ConstantOptimizationIterations, true,
          estimationLimits.Lower, estimationLimits.Upper);
      } else {
        if (applyLinearScaling) {
          //Check for interval arithmetic grammar
          //remove scaling nodes for linear scaling evaluation
          var rootNode = new ProgramRootSymbol().CreateTreeNode();
          var startNode = new StartSymbol().CreateTreeNode();
          SymbolicExpressionTree newTree = null;
          foreach (var node in solution.IterateNodesPrefix())
            if (node.Symbol.Name == "Scaling") {
              for (var i = 0; i < node.SubtreeCount; ++i) startNode.AddSubtree(node.GetSubtree(i));
              rootNode.AddSubtree(startNode);
              newTree = new SymbolicExpressionTree(rootNode);
              break;
            }

          //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);
          //Set alpha and beta to the scaling nodes from ia grammar
          foreach (var node in solution.IterateNodesPrefix())
            if (node.Symbol.Name == "Offset") {
              node.RemoveSubtree(1);
              var alphaNode = new ConstantTreeNode(new Constant()) { Value = alpha };
              node.AddSubtree(alphaNode);
            } else if (node.Symbol.Name == "Scaling") {
              node.RemoveSubtree(1);
              var betaNode = new ConstantTreeNode(new Constant()) { Value = beta };
              node.AddSubtree(betaNode);
            }
        }
      }

      var quality = Calculate(interpreter, solution, estimationLimits.Lower, estimationLimits.Upper, problemData, rows,
        PenaltyMultiplier, BoundsEstimator, UseSoftConstraints, MaximumPenaltyFactor);
      QualityParameter.ActualValue = new DoubleValue(quality);

      return base.InstrumentedApply();
    }

    public static double Calculate(
      ISymbolicDataAnalysisExpressionTreeInterpreter interpreter,
      ISymbolicExpressionTree solution, double lowerEstimationLimit,
      double upperEstimationLimit,
      IRegressionProblemData problemData, IEnumerable<int> rows,
      double penaltyMultiplier, IBoundsEstimator estimator,
      bool useSoftConstraints, double maximumPenaltyFactor) {

      var estimatedValues = interpreter.GetSymbolicExpressionTreeValues(solution, problemData.Dataset, rows);
      var targetValues = problemData.Dataset.GetDoubleValues(problemData.TargetVariable, rows);
      var constraints = problemData.IntervalConstraints.EnabledConstraints;
      var intervalCollection = problemData.VariableRanges;

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

      if (errorState != OnlineCalculatorError.None) {
        return 1.0;
      }

      var constraintResults = IntervalUtil.IntervalConstraintsViolation(constraints, estimator, intervalCollection, solution);

      if (constraintResults.Any(x => double.IsNaN(x) || double.IsInfinity(x))) {
        return 1.0;
      }


      if (useSoftConstraints) {
        if (maximumPenaltyFactor < 0.0)
          throw new ArgumentException("The parameter 'Maximum Penalty Factor' has to be greater or equal 0.0!");

        var zip = constraints.Zip(constraintResults, (c, e) => new { Constraint = c, Error = e });
        double sum = 0.0;

        foreach (var x in zip) {
          if (x.Constraint.Weight <= 0)
            throw new ArgumentException("Constraint weights <= 0 are not allowed!");

          double e = x.Error / x.Constraint.Weight;

          e = double.IsNaN(e) ? 0.0 : e;
          e = e > 1.0 ? 1.0 : e;

          sum += Math.Sqrt(e) * maximumPenaltyFactor;
        }

        double avgError = sum / zip.Count();
        nmse = nmse * avgError + nmse;
        nmse = Math.Min(1.0, nmse);
      } else if (constraintResults.Any(x => x != 0.0)) {
        nmse = 1.0;
      }

      return nmse;
    }

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

      var nmse = Calculate(SymbolicDataAnalysisTreeInterpreterParameter.ActualValue, tree,
        EstimationLimitsParameter.ActualValue.Lower, EstimationLimitsParameter.ActualValue.Upper,
        problemData, rows, PenaltyMultiplier, BoundsEstimator, UseSoftConstraints, MaximumPenaltyFactor);

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

      return nmse;
    }
  }
}