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

namespace CodeGenerator {
  public class RandomSearchCodeGen {

    private string solverTemplate = @"
namespace ?PROBLEMNAME? {
  public sealed class ?IDENT?Solver {
    private static double baseTerminalProbability = 0.05; // 5% of all samples are only a terminal node
    private static double terminalProbabilityInc = 0.05; // for each level the probability to sample a terminal grows by 5%

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

    
    public sealed class SolverState : ISolverState {
      private class Tree {
        public int altIdx;
        // public string symbol; // for debugging
        public List<Tree> subtrees;
        public Tree(int state, int altIdx) {
          subtrees = new List<Tree>(subtreeCount[state]);
          this.altIdx = altIdx;
        }
      }

      ?TERMINALNODEDEFINITIONS?

      public int curDepth;
      public int steps;
      public int depth;
      private readonly Stack<Tree> nodes;
      private readonly ?IDENT?Problem problem;

      public SolverState(?IDENT?Problem problem, int seed) {
        this.problem = problem;
        this.nodes = new Stack<Tree>();
        
        // create a random tree
        var tree = SampleTree(new Random(seed), 0, -1);  // state 0 is the state for the start symbol
        nodes.Push(tree);
      }

      public void Reset() {
        // stack must contain only the root of the tree
        System.Diagnostics.Debug.Assert(nodes.Count == 1);
      }

      private Tree SampleTree(Random random, int state, int altIdx) {
        // Console.Write(state + "" "");        
        curDepth += 1;
        steps += 1;
        depth = Math.Max(depth, curDepth);
        var t = new Tree(state, altIdx);

        // t.symbol = symb.Length > state ? symb[state] : ""TERM"";
        // if the symbol has alternatives then we must choose one randomly (only one sub-tree in this case)
        if(subtreeCount[state] == 1) {
          var targetStates = transition[state];
          var i = SampleAlternative(random, state);
          if(targetStates.Length == 0) {
            //terminal
            t.subtrees.Add(SampleTree(random, -1, i));
          } else {
            t.subtrees.Add(SampleTree(random, targetStates[i], i));
          }
        } else {
          // if the symbol contains only one sequence we must use create sub-trees for each symbol in the sequence
          for(int i = 0; i < subtreeCount[state]; i++) {
            t.subtrees.Add(SampleTree(random, transition[state][i], i));
          }
        }
        curDepth -=1;
        return t;
      }

      public int PeekNextAlternative() {
        // this must only be called nodes that contain alternatives and therefore must only have single-symbols alternatives
        System.Diagnostics.Debug.Assert(nodes.Peek().subtrees.Count == 1);
        return nodes.Peek().subtrees[0].altIdx;
      }

      public void Follow(int idx) {
        nodes.Push(nodes.Peek().subtrees[idx]);
      }

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

      private int SampleAlternative(Random random, int state) {
        switch(state) {

?SAMPLEALTERNATIVECODE?

          default: throw new InvalidOperationException();
        }
      }

      private double TerminalProbForDepth(int depth) {
        return baseTerminalProbability + depth * terminalProbabilityInc;
      }
    }

    public static void Main(string[] args) {
      if(args.Length >= 1) ParseArguments(args);

      var problem = new ?IDENT?Problem();
      var solver = new ?IDENT?Solver(problem);
      solver.Start();
    }
    private static void ParseArguments(string[] args) {
      var baseTerminalProbabilityRegex = new Regex(@""--terminalProbBase=(?<prob>.+)"");
      var terminalProbabilityIncRegex = new Regex(@""--terminalProbInc=(?<prob>.+)"");
      var helpRegex = new Regex(@""--help|/\?"");
      
      foreach(var arg in args) {
        var baseTerminalProbabilityMatch = baseTerminalProbabilityRegex.Match(arg);
        var terminalProbabilityIncMatch = terminalProbabilityIncRegex.Match(arg);
        var helpMatch = helpRegex.Match(arg);
        if(helpMatch.Success) { PrintUsage(); Environment.Exit(0); }
        else if(baseTerminalProbabilityMatch.Success) {
          baseTerminalProbability = double.Parse(baseTerminalProbabilityMatch.Groups[""prob""].Captures[0].Value, System.Globalization.CultureInfo.InvariantCulture);
          if(baseTerminalProbability < 0.0 || baseTerminalProbability > 1.0) throw new ArgumentException(""base terminal probability must lie in range [0.0 ... 1.0]"");
        } else if(terminalProbabilityIncMatch.Success) {
           terminalProbabilityInc = double.Parse(terminalProbabilityIncMatch.Groups[""prob""].Captures[0].Value, System.Globalization.CultureInfo.InvariantCulture);
           if(terminalProbabilityInc < 0.0 || terminalProbabilityInc > 1.0) throw new ArgumentException(""terminal probability increment must lie in range [0.0 ... 1.0]"");
        } else {
           Console.WriteLine(""Unknown switch {0}"", arg); PrintUsage(); Environment.Exit(0); 
        }
      }
    }
    private static void PrintUsage() {
      Console.WriteLine(""Find a solution using random tree search."");
      Console.WriteLine();
      Console.WriteLine(""Parameters:"");
      Console.WriteLine(""\t--terminalProbBase=<prob>\tSets the probability of sampling a terminal alternative in a rule [Default: 0.05]"");
      Console.WriteLine(""\t--terminalProbInc=<prob>\tSets the increment for the probability of sampling a terminal alternative for each level in the syntax tree [Default: 0.05]"");
    }


    private readonly ?IDENT?Problem problem;
    public ?IDENT?Solver(?IDENT?Problem problem) {
      this.problem = problem;
    }

    private void Start() {
      var seedRandom = new Random();
      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();
      while (true) {

        // must make sure that calling the start-symbol multiple times in the fitness function always leads to the same path through the grammar
        // so we use a PRNG for generating seeds for a separate PRNG that is reset each time the start symbol is called
        
        var _state = new SolverState(problem, seedRandom.Next());

        var f = problem.Evaluate(_state);
 
        n++;
        sumSize += _state.steps;
        sumDepth += _state.depth;
        sumF += f;
        if (IsBetter(f, bestF)) {
          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, IEnumerable<TerminalNode> terminals, bool maximization, SourceBuilder problemSourceCode) {
      var solverSourceCode = new SourceBuilder();
      solverSourceCode.Append(solverTemplate)
        // .Replace("?TERMINALFIELDS?", GenerateTerminalFields(grammar))
        .Replace("?TERMINALNODEDEFINITIONS?", GenerateTerminalNodeDefinitions(terminals))
        // .Replace("?CONSTRUCTORCODE?", GenerateConstructorCode(grammar))
        .Replace("?MAXIMIZATION?", maximization.ToString().ToLowerInvariant())
        .Replace("?SYMBOLNAMES?", grammar.Symbols.Select(s => s.Name).Aggregate(string.Empty, (str, symb) => str + "\"" + symb + "\", "))
        .Replace("?TRANSITIONTABLE?", GenerateTransitionTable(grammar))
        .Replace("?SUBTREECOUNTTABLE?", GenerateSubtreeCountTable(grammar))
        .Replace("?SAMPLEALTERNATIVECODE?", GenerateSampleAlternativeSource(grammar))
//        .Replace("?SAMPLETERMINALCODE?", GenerateSampleTerminalSource(grammar))
      ;

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

    private string GenerateTerminalNodeDefinitions(IEnumerable<TerminalNode> terminals) {
      var sb = new SourceBuilder();

      foreach (var terminal in terminals) {
        GenerateTerminalNodeDefinitions(terminal, sb);
      }
      return sb.ToString();
    }

    private void GenerateTerminalNodeDefinitions(TerminalNode terminal, SourceBuilder sb) {
      sb.AppendFormat("private class {0}Tree {{", terminal.Ident).BeginBlock();
      foreach (var att in terminal.FieldDefinitions) {
        sb.AppendFormat("{0} {1};", att.Type, att.Identifier).AppendLine();
      }
      sb.AppendFormat(" static {0}Tree Create(Random random, ?IDENT?Problem problem) {{", terminal.Ident).BeginBlock();
      sb.AppendFormat("var t = new {0}Tree();", terminal.Ident);
      foreach (var constr in terminal.Constraints) {
        if (constr.Type == ConstraintNodeType.Set) {
          sb.AppendLine("{").BeginBlock();
          sb.AppendFormat(" var elements = problem.GetAllowed{0}_{1}().ToArray();", terminal.Ident, constr.Ident);
          sb.AppendFormat("t.{0} = elements[random.Next(elements.Length)]; ", constr.Ident).AppendLine();
          sb.AppendLine("}").EndBlock();
        } else {
          sb.AppendLine("{").BeginBlock();
          sb.AppendFormat(" var max = problem.GetMax{0}_{1}();", terminal.Ident, constr.Ident);
          sb.AppendFormat(" var min = problem.GetMin{0}_{1}();", terminal.Ident, constr.Ident);
          sb.AppendFormat("t.{0} = random.NextDouble() * (max - min) + min ", constr.Ident).AppendLine();
          sb.AppendLine("}").EndBlock();
        }
      }
      sb.AppendLine("return t;");
      sb.AppendLine("}").EndBlock();
      sb.AppendLine("}").EndBlock();
    }

    //private string GenerateSampleTerminalSource(IGrammar grammar) {
    //  StringBuilder sb = new StringBuilder();
    //  foreach (var t in grammar.TerminalSymbols) {
    //    sb.AppendFormat("public void {0}(ISolverState _state, {1}) {{", t.Name, t.GetAttributeString()).AppendLine();
    //    foreach (var att in t.Attributes) {
    //      // need constraints
    //      sb.AppendFormat("{0}", att.Name);
    //    }
    //    sb.AppendLine("}");
    //  }
    //  return sb.ToString();
    //}


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

    private string GenerateSampleAlternativeSource(IGrammar grammar) {
      Debug.Assert(grammar.Symbols.First().Equals(grammar.StartSymbol));
      var sb = new SourceBuilder();
      int stateCount = 0;
      foreach (var s in grammar.Symbols) {
        sb.AppendFormat("case {0}: ", stateCount++);
        if (grammar.IsTerminal(s)) {
          // ignore
        } else {
          var terminalAltIndexes = grammar.GetAlternatives(s)
            .Select((alt, idx) => new { alt, idx })
            .Where((p) => p.alt.All(symb => grammar.IsTerminal(symb)))
            .Select(p => p.idx);
          var nonTerminalAltIndexes = grammar.GetAlternatives(s)
            .Select((alt, idx) => new { alt, idx })
            .Where((p) => p.alt.Any(symb => grammar.IsNonTerminal(symb)))
            .Select(p => p.idx);
          var hasTerminalAlts = terminalAltIndexes.Any();
          var hasNonTerminalAlts = nonTerminalAltIndexes.Any();
          if (hasTerminalAlts && hasNonTerminalAlts) {
            sb.Append("if(random.NextDouble() < TerminalProbForDepth(depth)) {").BeginBlock();
            GenerateReturnStatement(terminalAltIndexes, sb);
            sb.Append("} else {");
            GenerateReturnStatement(nonTerminalAltIndexes, sb);
            sb.Append("}").EndBlock();
          } else {
            GenerateReturnStatement(grammar.NumberOfAlternatives(s), sb);
          }
        }
      }
      return sb.ToString();
    }
    private void GenerateReturnStatement(IEnumerable<int> idxs, SourceBuilder sb) {
      if (idxs.Count() == 1) {
        sb.AppendFormat("return {0};", idxs.Single()).AppendLine();
      } else {
        var idxStr = idxs.Aggregate(string.Empty, (str, idx) => str + idx + ", ");
        sb.AppendFormat("return new int[] {{ {0} }}[random.Next({1})]; ", idxStr, idxs.Count()).AppendLine();
      }
    }

    private void GenerateReturnStatement(int nAlts, SourceBuilder sb) {
      if (nAlts > 1) {
        sb.AppendFormat("return random.Next({0});", nAlts).AppendLine();
      } else if (nAlts == 1) {
        sb.AppendLine("return 0; ");
      } else {
        sb.AppendLine("throw new InvalidProgramException();");
      }
    }
  }
}