using System;
using System.Collections.Generic;
using System.Linq;
using HeuristicLab.Common;
using HeuristicLab.Encodings.RealVectorEncoding;
using HeuristicLab.Problems.DataAnalysis;

namespace HeuristicLab.Networks.IntegratedOptimization.SurrogateModeling {
  public static class ExpectedImprovementHelpers {
    public static IEnumerable<double> Evaluate(IEnumerable<RealVector> points, IRegressionSolution solution, bool calculateExpectedImprovement = true) {
      var model = solution.Model;
      var problemData = solution.ProblemData;
      var dataset = problemData.Dataset;

      var modifiableDataset = new ModifiableDataset(problemData.AllowedInputVariables, problemData.AllowedInputVariables.Select(x => new List<double>()));
      foreach (var point in points)
        modifiableDataset.AddRow(point.Select(x => (object)x));

      var targets = model.GetEstimatedValues(modifiableDataset, Enumerable.Range(0, modifiableDataset.Rows))
                         .ToArray();

      if (calculateExpectedImprovement) {
        var confModel = model as IConfidenceRegressionModel;
        if (confModel != null) {
          var minTarget = dataset.GetDoubleValues(problemData.TargetVariable).ToList().Min();
          var uncertainties = confModel.GetEstimatedVariances(modifiableDataset, Enumerable.Range(0, modifiableDataset.Rows))
                                       .Select(Math.Sqrt)
                                       .ToArray();
          for (int i = 0; i < modifiableDataset.Rows; i++)
            targets[i] = CalculateExpectedImprovement(minTarget, targets[i], uncertainties[i]);
        }
      }

      return targets;
    }

    private static double CalculateExpectedImprovement(double bestTarget, double estimatedTarget, double modelUncertainty) {
      if (modelUncertainty.IsAlmost(0.0)) return 0.0;

      var delta = bestTarget - estimatedTarget;
      var x = delta / modelUncertainty;
      var expImp = delta * alglib.normaldistribution(x) + modelUncertainty * Math.Exp(-0.5 * x * x) / Math.Sqrt(2 * Math.PI);

      return double.IsNaN(expImp) || double.IsInfinity(expImp) ? 0.0 : expImp;
    }
  }
}