using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.Contracts;
using System.Linq;
using HeuristicLab.Common;
using HeuristicLab.Core;
using HeuristicLab.Problems.DataAnalysis;

namespace GradientBoostedTrees {
  public class RegressionTreeBuilder {
    private readonly IRandom random;
    private readonly IRegressionProblemData problemData;

    private readonly int nCols;
    private readonly double[][] x; // all training data (original order from problemData)
    private double[] y; // training labels (original order from problemData)

    private Dictionary<string, double> sumImprovements; // for variable relevance calculation

    private readonly string[] allowedVariables; // all variables in shuffled order
    private Dictionary<string, int> varName2Index; // maps the variable names to column indexes 
    private int effectiveVars; // number of variables that are used from allowedVariables

    private int effectiveRows; // number of rows that are used from 
    private readonly int[][] sortedIdxAll;
    private readonly int[][] sortedIdx; // random selection from sortedIdxAll (for r < 1.0)



    // helper arrays which are allocated to maximal necessary size only once in the ctor
    private readonly int[] internalIdx, which, leftTmp, rightTmp;
    private readonly double[] outx;
    private readonly int[] outSortedIdx;

    private RegressionTreeModel.TreeNode[] tree; // tree is represented as a flat array of nodes
    private int curTreeNodeIdx; // the index where the next tree node is stored

    private readonly IList<RegressionTreeModel.TreeNode> nodeQueue; //TODO

    // prepare and allocate buffer variables in ctor
    public RegressionTreeBuilder(IRegressionProblemData problemData, IRandom random) {
      this.problemData = problemData;
      this.random = random;

      var rows = problemData.TrainingIndices.Count();

      this.nCols = problemData.AllowedInputVariables.Count();

      allowedVariables = problemData.AllowedInputVariables.ToArray();
      varName2Index = new Dictionary<string, int>(allowedVariables.Length);
      for (int i = 0; i < allowedVariables.Length; i++) varName2Index.Add(allowedVariables[i], i);

      sortedIdxAll = new int[nCols][];
      sortedIdx = new int[nCols][];
      sumImprovements = new Dictionary<string, double>();
      internalIdx = new int[rows];
      which = new int[rows];
      leftTmp = new int[rows];
      rightTmp = new int[rows];
      outx = new double[rows];
      outSortedIdx = new int[rows];
      nodeQueue = new List<RegressionTreeModel.TreeNode>();

      x = new double[nCols][];
      y = problemData.Dataset.GetDoubleValues(problemData.TargetVariable, problemData.TrainingIndices).ToArray();


      int col = 0;
      foreach (var inputVariable in problemData.AllowedInputVariables) {
        x[col] = problemData.Dataset.GetDoubleValues(inputVariable, problemData.TrainingIndices).ToArray();
        sortedIdxAll[col] = Enumerable.Range(0, rows).OrderBy(r => x[col][r]).ToArray();
        sortedIdx[col] = new int[rows];
        col++;
      }
    }

    // r and m work in the same way as for alglib random forest
    // r is fraction of rows to use for training
    // m is fraction of variables to use for training
    public IRegressionModel CreateRegressionTree(int maxDepth, double r = 0.5, double m = 0.5) {
      // subtract mean of y first
      var yAvg = y.Average();
      for (int i = 0; i < y.Length; i++) y[i] -= yAvg;

      var seLoss = new SquaredErrorLoss();
      var zeros = Enumerable.Repeat(0.0, y.Length);
      var ones = Enumerable.Repeat(1.0, y.Length);

      var model = CreateRegressionTreeForGradientBoosting(y, maxDepth, problemData.TrainingIndices.ToArray(), seLoss.GetLineSearchFunc(y, zeros, ones), r, m);
      return new GradientBoostedTreesModel(new[] { new ConstantRegressionModel(yAvg), model }, new[] { 1.0, 1.0 });
    }

    // specific interface that allows to specify the target labels and the training rows which is necessary when this functionality is called by the gradient boosting routine
    public IRegressionModel CreateRegressionTreeForGradientBoosting(double[] y, int maxDepth, int[] idx, LineSearchFunc lineSearch, double r = 0.5, double m = 0.5) {
      Contract.Assert(maxDepth > 0);
      Contract.Assert(r > 0);
      Contract.Assert(r <= 1.0);
      Contract.Assert(y.Count() == this.y.Length);
      Contract.Assert(m > 0);
      Contract.Assert(m <= 1.0);

      this.y = y; // y is changed in gradient boosting

      // shuffle row idx
      HeuristicLab.Random.ListExtensions.ShuffleInPlace(idx, random);

      int nRows = idx.Count();

      // shuffle variable idx
      HeuristicLab.Random.ListExtensions.ShuffleInPlace(allowedVariables, random);

      effectiveRows = (int)Math.Ceiling(nRows * r);
      effectiveVars = (int)Math.Ceiling(nCols * m);

      Array.Clear(which, 0, which.Length);

      // mark selected rows
      for (int row = 0; row < effectiveRows; row++) {
        which[idx[row]] = 1;
        internalIdx[row] = idx[row];
      }

      for (int col = 0; col < nCols; col++) {
        int i = 0;
        for (int row = 0; row < nRows; row++) {
          if (which[sortedIdxAll[col][row]] > 0) {
            Trace.Assert(i < effectiveRows);
            sortedIdx[col][i] = sortedIdxAll[col][row];
            i++;
          }
        }
      }

      // prepare array for the tree nodes (a tree of maxDepth=1 has 1 node, a tree of maxDepth=d has 2^d - 1 nodes)      
      int numNodes = (int)Math.Pow(2, maxDepth) - 1;
      //this.tree = new RegressionTreeModel.TreeNode[numNodes];
      this.tree = Enumerable.Range(0, numNodes).Select(_=>new RegressionTreeModel.TreeNode()).ToArray();
      this.curTreeNodeIdx = 0;

      // start and end idx are inclusive
      CreateRegressionTreeForIdx(maxDepth, 0, effectiveRows - 1, lineSearch);
      return new RegressionTreeModel(tree);
    }

    // startIdx and endIdx are inclusive
    private void CreateRegressionTreeForIdx(int maxDepth, int startIdx, int endIdx, LineSearchFunc lineSearch) {
      Contract.Assert(endIdx - startIdx >= 0);
      Contract.Assert(startIdx >= 0);
      Contract.Assert(endIdx < internalIdx.Length);

      // TODO: stop when y is constant
      // TODO: use priority queue of nodes to be expanded (sorted by improvement) instead of the recursion to maximum depth
      if (maxDepth <= 1 || endIdx - startIdx == 0) {
        // max depth reached or only one element    
        tree[curTreeNodeIdx].varName = RegressionTreeModel.TreeNode.NO_VARIABLE;
        tree[curTreeNodeIdx].val = lineSearch(internalIdx, startIdx, endIdx);
        curTreeNodeIdx++;
      } else {
        int i, j;
        double threshold;
        string bestVariableName;
        FindBestVariableAndThreshold(startIdx, endIdx, out threshold, out bestVariableName);

        // if bestVariableName is NO_VARIABLE then no split was possible anymore
        if (bestVariableName == RegressionTreeModel.TreeNode.NO_VARIABLE) {
          // max depth reached or only one element    
          tree[curTreeNodeIdx].varName = RegressionTreeModel.TreeNode.NO_VARIABLE;
          tree[curTreeNodeIdx].val = lineSearch(internalIdx, startIdx, endIdx);
          curTreeNodeIdx++;
        } else {


          int bestVarIdx = varName2Index[bestVariableName];
          // split - two pass

          // store which index goes where
          for (int k = startIdx; k <= endIdx; k++) {
            if (x[bestVarIdx][internalIdx[k]] <= threshold)
              which[internalIdx[k]] = -1; // left partition
            else
              which[internalIdx[k]] = 1; // right partition
          }

          // partition sortedIdx for each variable
          for (int col = 0; col < nCols; col++) {
            i = 0;
            j = 0;
            int k;
            for (k = startIdx; k <= endIdx; k++) {
              Debug.Assert(Math.Abs(which[sortedIdx[col][k]]) == 1);

              if (which[sortedIdx[col][k]] < 0) {
                leftTmp[i++] = sortedIdx[col][k];
              } else {
                rightTmp[j++] = sortedIdx[col][k];
              }
            }
            Debug.Assert(i > 0); // at least on element in the left partition
            Debug.Assert(j > 0); // at least one element in the right partition
            Debug.Assert(i + j == endIdx - startIdx + 1);
            k = startIdx;
            for (int l = 0; l < i; l++) sortedIdx[col][k++] = leftTmp[l];
            for (int l = 0; l < j; l++) sortedIdx[col][k++] = rightTmp[l];
          }

          // partition row indices
          i = startIdx;
          j = endIdx;
          while (i <= j) {
            Debug.Assert(Math.Abs(which[internalIdx[i]]) == 1);
            Debug.Assert(Math.Abs(which[internalIdx[j]]) == 1);
            if (which[internalIdx[i]] < 0) i++;
            else if (which[internalIdx[j]] > 0) j--;
            else {
              Trace.Assert(which[internalIdx[i]] > 0);
              Trace.Assert(which[internalIdx[j]] < 0);
              // swap
              int tmp = internalIdx[i];
              internalIdx[i] = internalIdx[j];
              internalIdx[j] = tmp;
              i++;
              j--;
            }
          }
          Debug.Assert(j < i);
          Debug.Assert(i >= startIdx);
          Debug.Assert(j <= endIdx);

          var parentIdx = curTreeNodeIdx;
          tree[parentIdx].varName = bestVariableName;
          tree[parentIdx].val = threshold;
          curTreeNodeIdx++;

          // create left subtree
          tree[parentIdx].leftIdx = curTreeNodeIdx;
          CreateRegressionTreeForIdx(maxDepth - 1, startIdx, j, lineSearch);

          // create right subtree
          tree[parentIdx].rightIdx = curTreeNodeIdx;
          CreateRegressionTreeForIdx(maxDepth - 1, i, endIdx, lineSearch);
        }
      }
    }

    private void FindBestVariableAndThreshold(int startIdx, int endIdx, out double threshold, out string bestVar) {
      Contract.Assert(startIdx < endIdx + 1); // at least 2 elements

      int rows = endIdx - startIdx + 1;
      Contract.Assert(rows >= 2);

      double sumY = 0.0;
      for (int i = startIdx; i <= endIdx; i++) {
        sumY += y[internalIdx[i]];
      }

      double bestImprovement = 0.0;
      double bestThreshold = double.PositiveInfinity;
      bestVar = string.Empty;

      for (int col = 0; col < effectiveVars; col++) {
        // sort values for variable to prepare for threshold selection 
        var curVariable = allowedVariables[col];
        var curVariableIdx = varName2Index[curVariable];
        for (int i = startIdx; i <= endIdx; i++) {
          var sortedI = sortedIdx[curVariableIdx][i];
          outSortedIdx[i - startIdx] = sortedI;
          outx[i - startIdx] = x[curVariableIdx][sortedI];
        }

        double curImprovement;
        double curThreshold;
        FindBestThreshold(outx, outSortedIdx, rows, y, sumY, out curThreshold, out curImprovement);

        if (curImprovement > bestImprovement) {
          bestImprovement = curImprovement;
          bestThreshold = curThreshold;
          bestVar = allowedVariables[col];
        }
      }

      UpdateVariableRelevance(bestVar, sumY, bestImprovement, rows);

      threshold = bestThreshold;

      // Contract.Assert(bestImprovement > 0);
      // Contract.Assert(bestImprovement < double.PositiveInfinity);
      // Contract.Assert(bestVar != string.Empty);
      // Contract.Assert(allowedVariables.Contains(bestVar));
    }

    // assumption is that the Average(y) = 0
    private void UpdateVariableRelevance(string bestVar, double sumY, double bestImprovement, int rows) {
      if (string.IsNullOrEmpty(bestVar)) return;
      // update variable relevance
      double err = sumY * sumY / rows;
      double errAfterSplit = bestImprovement;

      double delta = (errAfterSplit - err); // relative reduction in squared error
      double v;
      if (!sumImprovements.TryGetValue(bestVar, out v)) {
        sumImprovements[bestVar] = delta;
      }
      sumImprovements[bestVar] = v + delta;
    }

    // x [0..N-1] contains rows sorted values in the range from [0..rows-1]
    // sortedIdx [0..N-1] contains the idx of the values in x in the original dataset in the range from [0..rows-1]
    // rows specifies the number of valid entries in x and sortedIdx
    // y [0..N-1] contains the target values in original sorting order
    // sumY is y.Sum()
    // 
    // the routine returns the best threshold (x[i] + x[i+1]) / 2 for i = [0 .. rows-2] by calculating the reduction in squared error
    // additionally the reduction in squared error is returned in bestImprovement
    // if all elements of x are equal the routing fails to produce a threshold
    private static void FindBestThreshold(double[] x, int[] sortedIdx, int rows, double[] y, double sumY, out double bestThreshold, out double bestImprovement) {
      Contract.Assert(rows >= 2);

      double sl = 0.0;
      double sr = sumY;
      double nl = 0.0;
      double nr = rows;

      bestImprovement = 0.0;
      bestThreshold = double.NegativeInfinity;
      // for all thresholds
      // if we have n rows there are n-1 possible splits
      for (int i = 0; i < rows - 1; i++) {
        sl += y[sortedIdx[i]];
        sr -= y[sortedIdx[i]];

        nl++;
        nr--;
        Debug.Assert(nl > 0);
        Debug.Assert(nr > 0);

        if (x[i] < x[i + 1]) { // don't try to split when two elements are equal
          double curQuality = sl * sl / nl + sr * sr / nr;
          // curQuality = nl*nr / (nl+nr) * Sqr(sl / nl - sr / nr) // greedy function approximation page 12 eqn (35)

          if (curQuality > bestImprovement) {
            bestThreshold = (x[i] + x[i + 1]) / 2.0;
            bestImprovement = curQuality;
          }
        }
      }

      // if all elements where the same then no split can be found
    }


    public IEnumerable<KeyValuePair<string, double>> GetVariableRelevance() {
      double scaling = 100 / sumImprovements.Max(t => t.Value);
      return
        sumImprovements
        .Select(t => new KeyValuePair<string, double>(t.Key, t.Value * scaling))
        .OrderByDescending(t => t.Value);
    }
  }
}