using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using HeuristicLab.Common;
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 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(Random random, string curState, IEnumerable<string> afterStates, out int selectedStateIdx) {
      // 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++;
      }

      selectedStateIdx = actionIndexMap[banditPolicy.SelectAction(random, activeAfterStates.Take(idx))];

      return true;
    }



    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;
    }

    public void UpdateReward(IEnumerable<string> stateTrajectory, double reward) {
      foreach (var state in stateTrajectory) {
        GetStateInfo(state).UpdateReward(reward);

        // only the last state can be terminal
        if (problem.Grammar.IsTerminal(state)) {
          MarkAsDone(state);
        }
      }
    }


    public void Reset() {
      stateInfo.Clear();
      done.Clear();
    }

    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].Value;
      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;
    }
  }
}