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

namespace HeuristicLab.Problems.DataAnalysis.Symbolic.Regression.SingleObjective.Evaluators {
  [Item("Constraint Scaling NMSE Evaluator", "")]
  [StorableType("5B69083F-74EE-446C-A7D1-9DEAAA128AC6")]
  public class SymbolicRegressionSingleObjectiveConstraintScalingNmseEvaluator
    : SymbolicRegressionSingleObjectiveEvaluator {

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

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

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

    public SymbolicRegressionSingleObjectiveConstraintScalingNmseEvaluator() { }

    public override bool Maximization => false;


    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                        = false; 
      
     //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);
     var a = alpha;
     var b = beta;
     //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,
        applyLinearScaling);
      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, bool applyLinearScaling) {
      var estimatedValues =
        interpreter.GetSymbolicExpressionTreeValues(solution, problemData.Dataset, rows);
      var targetValues = problemData.Dataset.GetDoubleValues(problemData.TargetVariable, rows);
      var constraints = problemData.IntervalConstraints.EnabledConstraints;
      var variableRanges = problemData.VariableRanges.GetReadonlyDictionary();

      if (!SymbolicRegressionConstraintAnalyzer.ConstraintsSatisfied(constraints, variableRanges, solution))
        return 1.0;

      var boundedEstimatedValues = estimatedValues.LimitToRange(lowerEstimationLimit, upperEstimationLimit);
      var nmse = OnlineNormalizedMeanSquaredErrorCalculator.Calculate(targetValues, boundedEstimatedValues, out var errorState);
      if (errorState != OnlineCalculatorError.None) 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,
        ApplyLinearScalingParameter.ActualValue.Value);

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

      return nmse;
    }
  }
}