using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using HeuristicLab.Grammars;

namespace CodeGenerator {
  public class BruteForceCodeGen {
    private string solverTemplate = @"
namespace ?PROBLEMNAME? {
  public sealed class ?IDENT?Solver {

    private static Dictionary<int, int[]> transition = new Dictionary<int, int[]>() {
?TRANSITIONTABLE?
    };
    private static Dictionary<int, int> subtreeCount = new Dictionary<int, int>() {
       { -1, 0 }, // terminals
?SUBTREECOUNTTABLE?
    };

    private static string[] symb = new string[] { ?SYMBOLNAMES? };
    private static bool[] isTerminal = new bool[] { ?TERMINALTABLE? };
    private static bool[] hasAlternatives = new bool[] {?HASALTERNATIVESTABLE? };

    
    public sealed class SolverState : ISolverState {
      private class Node {
        public int state;
        public int count;
        public List<Node> nodes;
        public bool done;
        private int nextAlt;
        public Node(int state) {
          this.state = state;
          nodes = new List<Node>(transition[state].Length);
          if (!hasAlternatives[state]) {
            foreach (var t in transition[state]) { nodes.Add(new Node(t)); }
          }
        }

        public int GetNextAlternative() {
          System.Diagnostics.Debug.Assert(hasAlternatives[state]);
          if(nodes.Count == 0 && !isTerminal[state]) { 
            foreach(var targetState in transition[state]) nodes.Add(new Node(targetState));
            nextAlt = 0; 
            // begin with a terminal if possible
            if(!isTerminal[nodes[nextAlt].state]) 
              do { nextAlt = (nextAlt + 1) % nodes.Count; } while(nextAlt != 0 /* full circle*/ && !isTerminal[nodes[nextAlt].state]);
          }

          throw new NotImplementedException(); // TODO: handle terminal classes correctly
          return nextAlt;
        }

        public void Update() {
          count++;
          if(hasAlternatives[state] && !isTerminal[state]) {
            nodes[GetNextAlternative()].Update();
            // check if all nodes done and set this node to done
            if(nodes.All(n=>n.done)) { 
              this.done = true;
            } else {
              // must not walk nodes that are done
              do { nextAlt = (nextAlt + 1) % nodes.Count; } while (nodes[nextAlt].done);
            }
          } else {
            if(isTerminal[state]) {
              throw new NotImplementedException(); // TODO: handle terminal classes correctly
              done = (nextAlt + 1) >= GetAlternativesCount(symb[state]);
            } else {
              // update the smallest count. this basically means keeping all sub-trees the same and only altering the updated one
              Node minNode = nodes.First();
              foreach(var node in nodes.Skip(1)) {
                if(node.count < minNode.count) {
                  minNode = node;
                }
              }
              minNode.Update();
              if(nodes.All(n=>n.done)) this.done = true;
            }
          }
        }
      }


      ?ALTERNATIVESCOUNTMETHOD?


      public int curDepth;
      public int steps;
      public int depth;
      private readonly Stack<Node> nodes;
      private readonly IGpdlProblem problem;
      private readonly Node searchTree;

      public SolverState(IGpdlProblem problem) {
        this.problem = problem;
        nodes = new Stack<Node>();
        searchTree = new Node(0);
        nodes.Push(searchTree);
      }

      public void Reset() {
        System.Diagnostics.Debug.Assert(nodes.Count == 1);
        steps = 0;
        depth = 0;
        curDepth = 0;
      }

      public int PeekNextAlternative() {
        var curNode = nodes.Peek();
        System.Diagnostics.Debug.Assert(hasAlternatives[curNode.state]);
        return curNode.GetNextAlternative();
      }

      public void Follow(int idx) {
        steps++;
        curDepth++;
        depth = Math.Max(depth, curDepth);
        var curNode = nodes.Peek();
        if(hasAlternatives[curNode.state]) {
          nodes.Push(curNode.nodes[curNode.GetNextAlternative()]);
        } else {
          nodes.Push(curNode.nodes[idx]);
        }
      }

      public void Unwind() {
        nodes.Pop();
        curDepth--;
      }

      public void Update() {
        searchTree.Update();
      }
    }

    public static void Main(string[] args) {
      var problem = new ?IDENT?Problem();
      var solver = new ?IDENT?Solver(problem);
      solver.Start();
    }
    
    private readonly ?IDENT?Problem problem;
    public ?IDENT?Solver(?IDENT?Problem problem) {
      this.problem = problem;
    }

    private void Start() {
      var bestF = ?MAXIMIZATION? ? double.NegativeInfinity : double.PositiveInfinity;
      int n = 0;
      long sumDepth = 0;
      long sumSize = 0;
      var sumF = 0.0;
      var sw = new System.Diagnostics.Stopwatch();
      sw.Start();
      var _state = new SolverState(problem);
      while (true) {

        var f = problem.Evaluate(_state);
        _state.Update();

        n++;
        sumSize += _state.steps;
        sumDepth += _state.depth;
        sumF += f;
        if (IsBetter(f, bestF)) {
          // evaluate again with tracing to console
          // problem.Evaluate(new SolverState(_state.seed, true));
          bestF = f;
          Console.WriteLine(""{0}\t{1}\t(size={2}, depth={3})"", n, bestF, _state.steps, _state.depth);
        }
        if (n % 1000 == 0) {
          sw.Stop();
          Console.WriteLine(""{0}\tbest: {1:0.000}\t(avg: {2:0.000})\t(avg size: {3:0.0})\t(avg. depth: {4:0.0})\t({5:0.00} sols/ms)"", n, bestF, sumF/1000.0, sumSize/1000.0, sumDepth/1000.0, 1000.0 / sw.ElapsedMilliseconds);
          sw.Reset();
          sumSize = 0;
          sumDepth = 0;
          sumF = 0.0;
          sw.Start();
        }
      }
    }

    private bool IsBetter(double a, double b) { 
      return ?MAXIMIZATION? ? a > b : a < b;
    }
  }
}";


    public void Generate(IGrammar grammar, bool maximization, SourceBuilder problemSourceCode) {
      var solverSourceCode = new SourceBuilder();
      solverSourceCode.Append(solverTemplate)
        .Replace("?MAXIMIZATION?", maximization.ToString().ToLowerInvariant())
        .Replace("?SYMBOLNAMES?", grammar.Symbols.Select(s => s.Name).Aggregate(string.Empty, (str, symb) => str + "\"" + symb + "\", "))
        .Replace("?TERMINALTABLE?", grammar.Symbols.Select(grammar.IsTerminal).Aggregate(string.Empty, (str, b) => str + b.ToString().ToLowerInvariant() + ", "))
        .Replace("?HASALTERNATIVESTABLE?", grammar.Symbols.Select(s => grammar.IsNonTerminal(s) && (grammar.NumberOfAlternatives(s) > 1)).Aggregate(string.Empty, (str, b) => str + b.ToString().ToLowerInvariant() + ", "))
        .Replace("?TRANSITIONTABLE?", GenerateTransitionTable(grammar))
        .Replace("?SUBTREECOUNTTABLE?", GenerateSubtreeCountTable(grammar))
        .Replace("?ALTERNATIVESCOUNTMETHOD?", GenerateAlternativesCountMethod(grammar))
      ;

      problemSourceCode.Append(solverSourceCode.ToString());
    }

    #region same as random search
    private string GenerateTransitionTable(IGrammar grammar) {
      Debug.Assert(grammar.Symbols.First().Equals(grammar.StartSymbol));
      var sb = new SourceBuilder();

      // state idx = idx of the corresponding symbol in the grammar
      var allSymbols = grammar.Symbols.ToList();
      var attributes = new List<string>();
      foreach (var s in grammar.Symbols) {
        var targetStates = new List<int>();
        if (grammar.IsTerminal(s)) {
          foreach (var att in s.Attributes) {
            targetStates.Add(allSymbols.Count + attributes.Count);
            attributes.Add(s.Name + "_" + att);
          }
        } else {
          if (grammar.NumberOfAlternatives(s) > 1) {
            foreach (var alt in grammar.GetAlternatives(s)) {
              // only single-symbol alternatives are supported
              Debug.Assert(alt.Count() == 1);
              targetStates.Add(allSymbols.IndexOf(alt.Single()));
            }
          } else {
            // rule is a sequence of symbols
            var seq = grammar.GetAlternatives(s).Single();
            targetStates.AddRange(seq.Select(symb => allSymbols.IndexOf(symb)));
          }
        }

        var targetStateString = targetStates.Aggregate(string.Empty, (str, state) => str + state + ", ");

        var idxOfSourceState = allSymbols.IndexOf(s);
        sb.AppendFormat("// {0}", s).AppendLine();
        sb.AppendFormat("{{ {0} , new int[] {{ {1} }} }},", idxOfSourceState, targetStateString).AppendLine();
      }
      for (int attIdx = 0; attIdx < attributes.Count; attIdx++) {
        sb.AppendFormat("// {0}", attributes[attIdx]).AppendLine();
        sb.AppendFormat("{{ {0} , new int[] {{ }} }},", attIdx + allSymbols.Count).AppendLine();
      }
      return sb.ToString();
    }
    private string GenerateSubtreeCountTable(IGrammar grammar) {
      Debug.Assert(grammar.Symbols.First().Equals(grammar.StartSymbol));
      var sb = new SourceBuilder();

      // state idx = idx of the corresponding symbol in the grammar
      var allSymbols = grammar.Symbols.ToList();
      var attributes = new List<string>();
      foreach (var s in grammar.Symbols) {
        int subtreeCount;
        if (grammar.IsTerminal(s)) {
          subtreeCount = s.Attributes.Count();
          attributes.AddRange(s.Attributes.Select(att => s.Name + "_" + att.Name));
        } else {
          if (grammar.NumberOfAlternatives(s) > 1) {
            Debug.Assert(grammar.GetAlternatives(s).All(alt => alt.Count() == 1));
            subtreeCount = 1;
          } else {
            subtreeCount = grammar.GetAlternative(s, 0).Count();
          }
        }

        sb.AppendFormat("// {0}", s).AppendLine();
        sb.AppendFormat("{{ {0} , {1} }},", allSymbols.IndexOf(s), subtreeCount).AppendLine();
      }

      for (int attIdx = 0; attIdx < attributes.Count; attIdx++) {
        sb.AppendFormat("// {0}", attributes[attIdx]).AppendLine();
        sb.AppendFormat("{{ {0} , 1 }},", attIdx + allSymbols.Count).AppendLine();
      }
      return sb.ToString();
    }
    #endregion
    private string GenerateAlternativesCountMethod(IGrammar grammar) {
      Debug.Assert(grammar.Symbols.First().Equals(grammar.StartSymbol));
      var sb = new SourceBuilder();
      sb.Append("public int GetAlternativesCount(string symbol, string attribute = null) {").BeginBlock();
      sb.AppendLine("switch(symbol) {");
      foreach (var s in grammar.Symbols) {
        sb.AppendFormat("case \"{0}\":", s.Name).AppendLine();
        int altCount;
        if (grammar.IsTerminal(s)) {
          if (s.Attributes.Any()) {
            sb.Append("switch(attribute) {").BeginBlock();
            foreach (var att in s.Attributes) {
              sb.AppendFormat("case \"{1}\": return problem.GetCardinality(\"{0}\", \"{1}\");", s.Name, att.Name).AppendLine();
            }
            sb.Append("} // switch").EndBlock();
          } else {
            sb.AppendLine("return 0;");
          }
        } else {
          sb.AppendFormat("return {0};", grammar.NumberOfAlternatives(s)).AppendLine();
        }


      }
      sb.AppendLine("} // switch");
      sb.Append("}").EndBlock();
      return sb.ToString();
    }
  }
}