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

namespace HeuristicLab.Problems.DataAnalysis.Symbolic.Regression {
  [StorableType("676B237C-DD9C-4F24-B64F-D44B0FA1F6A6")]
  [Creatable(CreatableAttribute.Categories.DataAnalysisRegression, Priority = 120)]
  [Item(Name = "ConstrainedNLS", Description = "Non-linear Regression with non-linear constraints")]
  public class ConstrainedNLS : BasicAlgorithm {
    public static readonly string IterationsParameterName = "Iterations";
    public static readonly string SolverParameterName = "Solver";
    public static readonly string ModelStructureParameterName = "Model structure";


    public IFixedValueParameter<IntValue> IterationsParameter {
      get { return (IFixedValueParameter<IntValue>)Parameters[IterationsParameterName]; }
    }
    public IFixedValueParameter<StringValue> ModelStructureParameter {
      get { return (IFixedValueParameter<StringValue>)Parameters[ModelStructureParameterName]; }
    }
    public IConstrainedValueParameter<StringValue> SolverParameter { get { return (IConstrainedValueParameter<StringValue>)Parameters[SolverParameterName]; } }

    public IFixedValueParameter<DoubleValue> FuncToleranceRelParameter {
      get { return (IFixedValueParameter<DoubleValue>)Parameters["FuncToleranceRel"]; }
    }
    public IFixedValueParameter<DoubleValue> FuncToleranceAbsParameter {
      get { return (IFixedValueParameter<DoubleValue>)Parameters["FuncToleranceAbs"]; }
    }
    public IFixedValueParameter<DoubleValue> MaxTimeParameter {
      get { return (IFixedValueParameter<DoubleValue>)Parameters["MaxTime"]; }
    }
    public IFixedValueParameter<BoolValue> CheckGradientParameter {
      get { return (IFixedValueParameter<BoolValue>)Parameters["CheckGradient"]; }
    }
    public int Iterations { get { return IterationsParameter.Value.Value; } set { IterationsParameter.Value.Value = value; } }

    public StringValue Solver {
      get { return SolverParameter.Value; }
      set { throw new NotImplementedException(); }
    }

    public string ModelStructure {
      get { return ModelStructureParameter.Value.Value; }
      set { ModelStructureParameter.Value.Value = value; }
    }
    public bool CheckGradient {
      get { return CheckGradientParameter.Value.Value; }
      set { CheckGradientParameter.Value.Value = value; }
    }

    public double FuncToleranceRel { get { return FuncToleranceRelParameter.Value.Value; } set { FuncToleranceRelParameter.Value.Value = value; } }
    public double FuncToleranceAbs { get { return FuncToleranceAbsParameter.Value.Value; } set { FuncToleranceAbsParameter.Value.Value = value; } }
    public double MaxTime { get { return MaxTimeParameter.Value.Value; } set { MaxTimeParameter.Value.Value = value; } }

    public ConstrainedNLS() : base() {
      Problem = new RegressionProblem();

      Parameters.Add(new FixedValueParameter<StringValue>(ModelStructureParameterName, "The function for which the parameters must be fit (only numeric constants are tuned).", new StringValue("1.0 * x*x + 0.0")));
      Parameters.Add(new FixedValueParameter<IntValue>(IterationsParameterName, "Determines how many iterations should be calculated while optimizing the constant of a symbolic expression tree (0 indicates other or default stopping criterion).", new IntValue(10)));
      var validSolvers = new ItemSet<StringValue>(new[] {
        "MMA",
        "COBYLA",
        "CCSAQ",
        "ISRES",
        "DIRECT_G",
        "NLOPT_GN_DIRECT_L",
        "NLOPT_GN_DIRECT_L_RAND",
        "NLOPT_GN_ORIG_DIRECT",
        "NLOPT_GN_ORIG_DIRECT_L",
        "NLOPT_GD_STOGO",
        "NLOPT_GD_STOGO_RAND",
        "NLOPT_LD_LBFGS_NOCEDAL",
        "NLOPT_LD_LBFGS",
        "NLOPT_LN_PRAXIS",
        "NLOPT_LD_VAR1",
        "NLOPT_LD_VAR2",
        "NLOPT_LD_TNEWTON",
        "NLOPT_LD_TNEWTON_RESTART",
        "NLOPT_LD_TNEWTON_PRECOND",
        "NLOPT_LD_TNEWTON_PRECOND_RESTART",
        "NLOPT_GN_CRS2_LM",
        "NLOPT_GN_MLSL",
        "NLOPT_GD_MLSL",
        "NLOPT_GN_MLSL_LDS",
        "NLOPT_GD_MLSL_LDS",
        "NLOPT_LN_NEWUOA",
        "NLOPT_LN_NEWUOA_BOUND",
        "NLOPT_LN_NELDERMEAD",
        "NLOPT_LN_SBPLX",
        "NLOPT_LN_AUGLAG",
        "NLOPT_LD_AUGLAG",
        "NLOPT_LN_BOBYQA",
        "NLOPT_AUGLAG",
        "NLOPT_LD_SLSQP",
        "NLOPT_LD_CCSAQ",
        "NLOPT_GN_ESCH",
        "NLOPT_GN_AGS",
      }.Select(s => new StringValue(s).AsReadOnly()));
      Parameters.Add(new ConstrainedValueParameter<StringValue>(SolverParameterName, "The solver algorithm", validSolvers, validSolvers.First()));
      Parameters.Add(new FixedValueParameter<DoubleValue>("FuncToleranceRel", new DoubleValue(0)));
      Parameters.Add(new FixedValueParameter<DoubleValue>("FuncToleranceAbs", new DoubleValue(0)));
      Parameters.Add(new FixedValueParameter<DoubleValue>("MaxTime", new DoubleValue(10)));
      Parameters.Add(new FixedValueParameter<BoolValue>("CheckGradient", "Flag to indicate whether the gradient should be checked using numeric approximation", new BoolValue(false)));

      CheckGradientParameter.Hidden = true;
    }

    public ConstrainedNLS(ConstrainedNLS original, Cloner cloner) : base(original, cloner) {
    }

    [StorableHook(HookType.AfterDeserialization)]
    public void AfterDeserializationHook() {
      if (!Parameters.ContainsKey("CheckGradient")) {
        Parameters.Add(new FixedValueParameter<BoolValue>("CheckGradient", "Flag to indicate whether the gradient should be checked using numeric approximation", new BoolValue(false)));

        CheckGradientParameter.Hidden = true;
      }
    }

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

    public override bool SupportsPause => false;

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

    protected override void Run(CancellationToken cancellationToken) {
      var parser = new InfixExpressionParser();
      var tree = parser.Parse(ModelStructure);
      var problem = (IRegressionProblem)Problem;


      #region prepare results
      var functionEvaluations = new IntValue(0);
      Results.AddOrUpdateResult("Evaluations", functionEvaluations);
      var bestError = new DoubleValue(double.MaxValue);
      var curError = new DoubleValue(double.MaxValue);
      Results.AddOrUpdateResult("Best error", bestError);
      Results.AddOrUpdateResult("Current error", curError);
      Results.AddOrUpdateResult("Tree", tree);
      var qualitiesTable = new DataTable("Qualities");
      var curQualityRow = new DataRow("Current Quality");
      var bestQualityRow = new DataRow("Best Quality");
      qualitiesTable.Rows.Add(bestQualityRow);
      qualitiesTable.Rows.Add(curQualityRow);
      Results.AddOrUpdateResult("Qualities", qualitiesTable);

      var constraintRows = new List<IndexedDataRow<int>>(); // for access via index
      var constraintsTable = new IndexedDataTable<int>("Constraints");
      Results.AddOrUpdateResult("Constraints", constraintsTable);
      foreach (var constraint in problem.ProblemData.IntervalConstraints.Constraints.Where(c => c.Enabled)) {
        if (constraint.Interval.LowerBound > double.NegativeInfinity) {
          var constraintRow = new IndexedDataRow<int>("-" + constraint.Expression + " < " + (-constraint.Interval.LowerBound));
          constraintRows.Add(constraintRow);
          constraintsTable.Rows.Add(constraintRow);
        }
        if (constraint.Interval.UpperBound < double.PositiveInfinity) {
          var constraintRow = new IndexedDataRow<int>(constraint.Expression + " < " + (constraint.Interval.UpperBound));
          constraintRows.Add(constraintRow);
          constraintsTable.Rows.Add(constraintRow);
        }
      }

      var parametersTable = new IndexedDataTable<int>("Parameters");

      #endregion

      var state = new ConstrainedNLSInternal(Solver.Value, tree, Iterations, problem.ProblemData, FuncToleranceRel, FuncToleranceAbs, MaxTime);
      if (CheckGradient) state.CheckGradient = true;
      int idx = 0;
      var formatter = new InfixExpressionFormatter();
      var constraintDescriptions = state.ConstraintDescriptions.ToArray();
      foreach (var constraintTree in state.constraintTrees) {
        // HACK to remove parameter nodes which occurr multiple times
        var reparsedTree = parser.Parse(formatter.Format(constraintTree));
        Results.AddOrUpdateResult($"{constraintDescriptions[idx++]}", reparsedTree);
      }

      // we use a listener model here to get state from the solver      

      state.FunctionEvaluated += State_FunctionEvaluated;
      state.ConstraintEvaluated += State_ConstraintEvaluated;

      state.Optimize(ConstrainedNLSInternal.OptimizationMode.UpdateParametersAndKeepLinearScaling);
      bestError.Value = state.BestError;
      curQualityRow.Values.Add(state.CurError);
      bestQualityRow.Values.Add(bestError.Value);

      Results.AddOrUpdateResult("Best solution", CreateSolution((ISymbolicExpressionTree)state.BestTree.Clone(), problem.ProblemData));
      var bestConstraintValues = new DoubleArray(state.BestConstraintValues);
      bestConstraintValues.ElementNames = constraintDescriptions;
      Results.AddOrUpdateResult("Best solution constraint values", bestConstraintValues);


      // local function
      void State_FunctionEvaluated() {
        if (cancellationToken.IsCancellationRequested) state.RequestStop();
        functionEvaluations.Value++;
        bestError.Value = state.BestError;
        curError.Value = state.CurError;
        curQualityRow.Values.Add(state.CurError);
        bestQualityRow.Values.Add(bestError.Value);

        // on the first call create the data rows
        if(!parametersTable.Rows.Any()) {
          for(int i=0;i<state.BestSolution.Length;i++) {
            parametersTable.Rows.Add(new IndexedDataRow<int>("p" + i));
          }
        }
        for (int i = 0; i < state.BestSolution.Length; i++) {
          parametersTable.Rows["p" + i].Values.Add(Tuple.Create(functionEvaluations.Value, state.BestSolution[i])); // TODO: remove access via string
        }
      }

      // local function
      void State_ConstraintEvaluated(int constraintIdx, double value) {
        constraintRows[constraintIdx].Values.Add(Tuple.Create(functionEvaluations.Value, value));
      }
    }

    private static ISymbolicRegressionSolution CreateSolution(ISymbolicExpressionTree tree, IRegressionProblemData problemData) {
      var model = new SymbolicRegressionModel(problemData.TargetVariable, tree, new SymbolicDataAnalysisExpressionTreeLinearInterpreter());
      // model.CreateRegressionSolution produces a new ProblemData and recalculates intervals ==> use SymbolicRegressionSolution.ctor instead
      var sol = new SymbolicRegressionSolution(model, (IRegressionProblemData)problemData.Clone());
      // NOTE: the solution has slightly different derivative values because simplification of derivatives can be done differently when parameter values are fixed.

      // var sol = model.CreateRegressionSolution((IRegressionProblemData)problemData.Clone());

      return sol;
    }
  }
}