using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using HeuristicLab.Algorithms.Bandits.BanditPolicies;
using HeuristicLab.Algorithms.DataAnalysis;
using HeuristicLab.Common;
using HeuristicLab.Problems.DataAnalysis;
using HeuristicLab.Problems.GrammaticalOptimization;

namespace HeuristicLab.Algorithms.Bandits.GrammarPolicies {
  // this represents grammar policies that use one of the available bandit policies for state selection
  // any bandit policy can be used to select actions for states
  // a separate datastructure is used to store visited states and to prevent revisiting of states
  public sealed class GenericGrammarPolicy : IGrammarPolicy {
    private Dictionary<string, IBanditPolicyActionInfo> stateInfo; // stores the necessary information for bandit policies for each state (=canonical phrase)
    private HashSet<string> done;
    private readonly bool useCanonicalPhrases;
    private readonly IProblem problem;
    private readonly IBanditPolicy banditPolicy;
    public double[] OptimalPulls { get; private set; }

    public GenericGrammarPolicy(IProblem problem, IBanditPolicy banditPolicy, bool useCanonicalPhrases = false) {
      this.useCanonicalPhrases = useCanonicalPhrases;
      this.problem = problem;
      this.banditPolicy = banditPolicy;
      this.stateInfo = new Dictionary<string, IBanditPolicyActionInfo>();
      this.done = new HashSet<string>();
    }

    private IBanditPolicyActionInfo[] activeAfterStates; // don't allocate each time
    private int[] actionIndexMap; // don't allocate each time

    public bool TrySelect(System.Random random, string curState, IEnumerable<string> afterStates, out int selectedStateIdx) {
      //// only for debugging
      //if (done.Count == 30000) {
      //  foreach (var pair in stateInfo) {
      //    var state = pair.Key;
      //    var info = (DefaultPolicyActionInfo)pair.Value;
      //    if (info.Tries > 0) {
      //      Console.WriteLine("{0};{1};{2};{3};{4};{5}", state, info.Tries, info.Value, info.MaxReward,
      //        optimalSolutions.Contains(problem.CanonicalRepresentation(state)) ? 1 : 0,
      //        string.Join(";", GenerateFeaturesPoly10(state)));
      //    }
      //  }
      //  System.Environment.Exit(1);
      //}

      // fail if all states are done (corresponding state infos are disabled)
      if (afterStates.All(s => Done(s))) {
        // fail because all follow states have already been visited => also disable the current state (if we can be sure that it has been fully explored)
        MarkAsDone(curState);

        selectedStateIdx = -1;
        return false;
      }

      // determine active actions (not done yet) and create an array to map the selected index back to original actions
      if (activeAfterStates == null || activeAfterStates.Length < afterStates.Count()) {
        activeAfterStates = new IBanditPolicyActionInfo[afterStates.Count()];
        actionIndexMap = new int[afterStates.Count()];
      }
      var idx = 0; int originalIdx = 0;
      foreach (var afterState in afterStates) {
        if (!Done(afterState)) {
          activeAfterStates[idx] = GetStateInfo(afterState);
          actionIndexMap[idx] = originalIdx;
          idx++;
        }
        originalIdx++;
      }

      //// select terminals first
      //var terminalAfterstates = afterStates.Select((s, i) => new { s, i }).FirstOrDefault(t => !Done(t.s) && problem.Grammar.IsTerminal(t.s));
      //if (terminalAfterstates != null) {
      //  selectedStateIdx = terminalAfterstates.i;
      //  return true;
      //}

      if (valueApproximation == null) {
        // no approximation yet? --> use bandit
        selectedStateIdx = actionIndexMap[banditPolicy.SelectAction(random, activeAfterStates.Take(idx))];
      } else if (afterStates.Any(s => problem.Grammar.IsTerminal(s) && !Done(s))) {
        selectedStateIdx = SelectionMaxValueTerminalAction(random, afterStates);
      } else {
        // only internal states? --> use bandit
        selectedStateIdx = actionIndexMap[banditPolicy.SelectAction(random, activeAfterStates.Take(idx))];
      }
      return true;
    }

    private int SelectionMaxValueTerminalAction(System.Random random, IEnumerable<string> afterStates) {
      int idx = 0;
      var terminalStates = new List<string>();
      var originalIdx = new List<int>();
      foreach (var state in afterStates) {
        if (problem.Grammar.IsTerminal(state) && !Done(state)) {
          terminalStates.Add(state);
          originalIdx.Add(idx);
        }
        idx++;
      }

      return originalIdx[SelectionMaxValueAction(random, terminalStates)];
    }

    private IRegressionSolution valueApproximation;
    private int SelectionMaxValueAction(System.Random random, IEnumerable<string> afterStates) {

      // eps greedy
      //if (random.NextDouble() < 0.1) return Enumerable.Range(0, afterStates.Count()).SelectRandom(random);

      Dataset ds;
      string[] variablesNames;
      CreateDataset(afterStates, afterStates.Select<string, IBanditPolicyActionInfo>(_ => null), out ds, out variablesNames);

      var v = valueApproximation.Model.GetEstimatedValues(ds, Enumerable.Range(0, ds.Rows)).ToArray();

      //boltzmann exploration
      //double beta = 100;
      //var w = v.Select(vi => Math.Exp(beta * vi));
      //
      //return Enumerable.Range(0, v.Length).SampleProportional(random, w);

      return Enumerable.Range(0, v.Length).MaxItems(i => v[i]).SelectRandom(random);
    }

    private void UpdateValueApproximation() {
      Dataset ds;
      string[] variableNames;
      CreateDataset(stateInfo.Keys, stateInfo.Values, out ds, out variableNames);
      var problemData = new RegressionProblemData(ds, variableNames.Skip(1), variableNames.First());
      //problemData.TestPartition.Start = problemData.TestPartition.End; // all data are training data
      valueApproximation = GradientBoostedTreesAlgorithmStatic.TrainGbm(problemData, new SquaredErrorLoss(), 50, 0.1,
        0.5, 0.5, 100);
      Console.WriteLine(valueApproximation.TrainingRSquared);
      Console.WriteLine(valueApproximation.TestRSquared);
    }

    private void CreateDataset(IEnumerable<string> states, IEnumerable<IBanditPolicyActionInfo> infos, out Dataset ds, out string[] variableNames) {
      variableNames = new string[] { "maxValue" }.Concat(GenerateFeaturesPoly10("E").Select((_, i) => "f" + i)).ToArray();

      int rows = infos.Zip(states, (info, state) => new { info, state }).Count(i => i.info == null || (i.info.Tries == 1 && Done(i.state)));
      int cols = variableNames.Count();

      var variableValues = new double[rows, cols];
      int n = 0;
      foreach (var pair in states.Zip(infos, Tuple.Create)) {
        var state = pair.Item1;
        var info = (DefaultPolicyActionInfo)pair.Item2;
        if (info == null || (info.Tries == 1 && Done(state))) {
          if (info != null) {
            variableValues[n, 0] = info.MaxReward;
          }
          int col = 1;
          foreach (var f in GenerateFeaturesPoly10(state)) {
            variableValues[n, col++] = f;
          }
          n++;
        }
      }

      ds = new Dataset(variableNames, variableValues);
    }


    private IBanditPolicyActionInfo GetStateInfo(string state) {
      var s = CanonicalState(state);
      IBanditPolicyActionInfo info;
      if (!stateInfo.TryGetValue(s, out info)) {
        info = banditPolicy.CreateActionInfo();
        stateInfo[s] = info;
      }
      return info;
    }

    private int rewardUpdatesSinceLastTraining = 0;
    private HashSet<string> statesWritten = new HashSet<string>();
    public void UpdateReward(IEnumerable<string> stateTrajectory, double reward) {
      rewardUpdatesSinceLastTraining++;
      if (rewardUpdatesSinceLastTraining == 5000) {
        rewardUpdatesSinceLastTraining = 0;
        //// write
        //foreach (var pair in stateInfo) {
        //  var state = pair.Key;
        //  var info = (DefaultPolicyActionInfo)pair.Value;
        //  if (!statesWritten.Contains(state) && info.Tries > 0) {
        //    Console.WriteLine("{0};{1};{2};{3};{4}", state, info.Tries, info.Value, info.MaxReward, string.Join(";", GenerateFeaturesPoly10(state)));
        //    statesWritten.Add(state);
        //  }
        //}
        //
        //Console.WriteLine();
        //UpdateValueApproximation();
      }

      int lvl = 0;
      foreach (var state in stateTrajectory) {
        double alpha = 0.99;
        OptimalPulls[lvl] = alpha * OptimalPulls[lvl] + (1 - alpha) * (problem.IsOptimalPhrase(state) ? 1.0 : 0.0);
        lvl++;

        GetStateInfo(state).UpdateReward(reward);
        //reward *= 0.95;
        // only the last state can be terminal
        if (problem.Grammar.IsTerminal(state)) {
          MarkAsDone(state);
        }
      }
    }

    private IEnumerable<double> GenerateFeaturesPoly10(string state) {
      // yield return problem.IsOptimalPhrase(state) ? 1 : 0;
      foreach (var f in problem.GetFeatures(state)) yield return f.Value;

      //if (!state.EndsWith("E")) state = state + "+E";
      //int len = state.Length;
      //Debug.Assert(state[len - 1] == 'E');
      //foreach (var sy0 in new char[] { '+', '*' }) {
      //  foreach (var sy1 in problem.Grammar.TerminalSymbols) {
      //    foreach (var sy2 in new char[] { '+', '*' }) {
      //      yield return state.Length > 3 && state[len - 4] == sy0 && state[len - 3] == sy1 && state[len - 2] == sy2 ? 1 : 0;
      //    }
      //  }
      //}

      //yield return state.Length;
      //foreach (var terminalSy in problem.Grammar.TerminalSymbols) {
      //  yield return state.Length > 2 && state[0] == terminalSy && state[1] == '+' ? 1 : 0;
      //  yield return state.Length > 2 && state[0] == terminalSy && state[1] == '*' ? 1 : 0;
      //}
      // yield return optimalSolutions.Contains(problem.CanonicalRepresentation(state)) ? 1 : 0;
      //foreach (var term in optimalTerms) yield return Regex.Matches(problem.CanonicalRepresentation(state), term).Count == 1 ? 1 : 0;
      //var len = state.Length;
      //yield return len;
      //foreach (var t in problem.Grammar.TerminalSymbols) {
      //  yield return state.Count(ch => ch == t);
      //}
      //// pairs
      //foreach (var u in problem.Grammar.TerminalSymbols) {
      //  foreach (var v in problem.Grammar.TerminalSymbols) {
      //    int n = 0;
      //    for (int i = 0; i < state.Length - 1; i++) {
      //      if (state[i] == u && state[i + 1] == v) n++;
      //    }
      //    yield return n;
      //  }
      //}
    }


    public void Reset() {
      stateInfo.Clear();
      done.Clear();
      OptimalPulls = new double[300]; // max sentence length is limited anyway
    }

    public int GetTries(string state) {
      var s = CanonicalState(state);
      if (stateInfo.ContainsKey(s)) return stateInfo[s].Tries;
      else return 0;
    }

    public double GetValue(string state) {
      var s = CanonicalState(state);
      if (stateInfo.ContainsKey(s)) return stateInfo[s].MaxReward;
      else return 0.0; // TODO: check alternatives
    }

    // the canonical states for the value function (banditInfos) and the done set must be distinguished
    // sequences of different length could have the same canonical representation and can have the same value (banditInfo)
    // however, if the canonical representation of a state is shorter then we must not mark the canonical state as done when all possible derivations from the initial state have been explored
    // eg. in the ant problem the canonical representation for ...lllA is ...rA
    // even though all possible derivations (of limited length) of lllA have been visited we must not mark the state rA as done
    private void MarkAsDone(string state) {
      var s = CanonicalState(state);
      // when the lengths of the canonical string and the original string are the same we also disable the actions
      // always disable terminals
      Debug.Assert(s.Length <= state.Length);
      if (s.Length == state.Length || problem.Grammar.IsTerminal(state)) {
        Debug.Assert(!done.Contains(s));
        done.Add(s);
      } else {
        // for non-terminals where the canonical string is shorter than the original string we can only disable the canonical representation for all states in the same level
        Debug.Assert(!done.Contains(s + state.Length));
        done.Add(s + state.Length); // encode the original length of the state, states in the same level of the tree are treated as equivalent
      }
    }

    // symmetric to MarkDone
    private bool Done(string state) {
      var s = CanonicalState(state);
      if (s.Length == state.Length || problem.Grammar.IsTerminal(state)) {
        return done.Contains(s);
      } else {
        // it is not necessary to visit states if the canonical representation has already been fully explored
        if (done.Contains(s)) return true;
        if (done.Contains(s + state.Length)) return true;
        for (int i = 1; i < state.Length; i++) {
          if (done.Contains(s + i)) return true;
        }
        return false;
      }
    }

    private string CanonicalState(string state) {
      if (useCanonicalPhrases) {
        return problem.CanonicalRepresentation(state);
      } else
        return state;
    }
  }
}