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?RandomSearchSolver {
    private double baseTerminalProbability = 0.05; // 5% of all samples are only a terminal node
    private double terminalProbabilityInc = 0.05; // for each level the probability to sample a terminal grows by 5%
    private static int maxDepth = 20;

    private readonly ?IDENT?Problem problem;
    private readonly Random random;

    public Tree SampleTree(int maxDepth, out int steps, out int depth) {
      steps = 0;
      depth = 0;
      int curDepth = 0;
      return SampleTree(0, maxDepth, ref steps, ref curDepth, ref depth);
    }

    public Tree SampleTree(int state, int maxDepth, ref int steps, ref int curDepth, ref int depth) {
      curDepth += 1;
      Debug.Assert(maxDepth > 0);
      steps += 1;
      depth = Math.Max(depth, curDepth);
      Tree t = null;

      // terminals
      if(Grammar.subtreeCount[state] == 0) {
        t = CreateTerminalNode(state, random, problem);
      } else {
        // if the symbol has alternatives then we must choose one randomly (only one sub-tree in this case)
        if(Grammar.subtreeCount[state] == 1) {
          var targetStates = Grammar.transition[state];
          var altIdx = SampleAlternative(random, maxDepth - 1, state, curDepth);
          var alternative = SampleTree(targetStates[altIdx], maxDepth - 1, ref steps, ref curDepth, ref depth);
          t = new Tree(altIdx, new Tree[] { alternative });
        } else {
          // if the symbol contains only one sequence we must use create sub-trees for each symbol in the sequence
          Tree[] subtrees = new Tree[Grammar.subtreeCount[state]];
          for(int i = 0; i < Grammar.subtreeCount[state]; i++) {
            subtrees[i] = SampleTree(Grammar.transition[state][i], maxDepth - 1, ref steps, ref curDepth, ref depth);
          }
          t = new Tree(-1, subtrees); // alternative index is ignored
        }
      }
      curDepth -=1;
      return t;
    }

    private static Tree CreateTerminalNode(int state, Random random, ?IDENT?Problem problem) {
      switch(state) {
        ?CREATETERMINALNODECODE?
        default: { throw new ArgumentException(""Unknown state index "" + state); }
      }
    }

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

?SAMPLEALTERNATIVECODE?

        default: throw new InvalidOperationException();
      }
    }

    private double TerminalProbForDepth(int depth) {
      if(depth>=maxDepth) return 1.0;
      return baseTerminalProbability + depth * terminalProbabilityInc;
    }

    private void ParseArguments(string[] args) {
      var baseTerminalProbabilityRegex = new Regex(@""--terminalProbBase=(?<prob>.+)"");
      var terminalProbabilityIncRegex = new Regex(@""--terminalProbInc=(?<prob>.+)"");
      var maxDepthRegex = new Regex(@""--maxDepth=(?<d>.+)"");

      var helpRegex = new Regex(@""--help|/\?"");
      
      foreach(var arg in args) {
        var baseTerminalProbabilityMatch = baseTerminalProbabilityRegex.Match(arg);
        var terminalProbabilityIncMatch = terminalProbabilityIncRegex.Match(arg);
        var maxDepthMatch = maxDepthRegex.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 if(maxDepthMatch.Success) {
           maxDepth = int.Parse(maxDepthMatch.Groups[""d""].Captures[0].Value, System.Globalization.CultureInfo.InvariantCulture);
           if(maxDepth < 1 || maxDepth > 100) throw new ArgumentException(""max depth must lie in range [1 ... 100]"");
        } else {
           Console.WriteLine(""Unknown switch {0}"", arg); PrintUsage(); Environment.Exit(0); 
        }
      }
    }
    private void PrintUsage() {
      Console.WriteLine(""Find a solution using random tree search."");
      Console.WriteLine();
      Console.WriteLine(""Parameters:"");
      Console.WriteLine(""\t--maxDepth=<depth>\tSets the maximal depth of sampled trees [Default: 20]"");
      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]"");
    }


    public ?IDENT?RandomSearchSolver(?IDENT?Problem problem, string[] args) {
      if(args.Length >= 1) ParseArguments(args);

      this.problem = problem;
      this.random = new Random();
    }

    public 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();
      while (true) {

        int steps, depth;
        var _t = SampleTree(maxDepth, out steps, out depth);
        Debug.Assert(depth <= maxDepth);
        // _t.PrintTree(0); Console.WriteLine();
        var f = problem.Evaluate(_t);
 
        n++;    
        sumSize += steps;
        sumDepth += depth;
        sumF += f;
        if (problem.IsBetter(f, bestF)) {
          bestF = f;
          _t.PrintTree(0); Console.WriteLine();
          Console.WriteLine(""{0}\t{1}\t(size={2}, depth={3})"", n, bestF, steps, 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);
          sumSize = 0;
          sumDepth = 0;
          sumF = 0.0;
          sw.Restart();
        }
      }
    }
  }
}";

    public void Generate(IGrammar grammar, IEnumerable<TerminalNode> terminals, bool maximization, SourceBuilder problemSourceCode) {
      var solverSourceCode = new SourceBuilder();
      solverSourceCode.Append(solverTemplate)
        .Replace("?MAXIMIZATION?", maximization.ToString().ToLowerInvariant())
        .Replace("?SAMPLEALTERNATIVECODE?", GenerateSampleAlternativeSource(grammar))
        .Replace("?CREATETERMINALNODECODE?", GenerateCreateTerminalCode(grammar, terminals))
      ;

      problemSourceCode.Append(solverSourceCode.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(maxDepth <= 1 || random.NextDouble() < TerminalProbForDepth(depth)) {").BeginBlock();
            GenerateSampleTerminalAlternativesStatement(terminalAltIndexes, sb);
            sb.Append("} else {");
            GenerateSampleNonterminalAlternativesStatement(nonTerminalAltIndexes, sb);
            sb.Append("}").EndBlock();
          } else {
            GenerateReturnStatement(grammar.NumberOfAlternatives(s), sb);
          }
        }
      }
      return sb.ToString();
    }
    private string GenerateCreateTerminalCode(IGrammar grammar, IEnumerable<TerminalNode> terminals) {
      Debug.Assert(grammar.Symbols.First().Equals(grammar.StartSymbol));
      var sb = new SourceBuilder();
      var allSymbols = grammar.Symbols.ToList();
      foreach (var s in grammar.Symbols) {
        if (grammar.IsTerminal(s)) {
          sb.AppendFormat("case {0}: {{", allSymbols.IndexOf(s)).BeginBlock();
          sb.AppendFormat("var t = new {0}Tree();", s.Name).AppendLine();
          var terminal = terminals.Single(t => t.Ident == s.Name);
          foreach (var constr in terminal.Constraints) {
            if (constr.Type == ConstraintNodeType.Set) {
              sb.Append("{").BeginBlock();
              sb.AppendFormat("var elements = problem.GetAllowed{0}_{1}().ToArray();", terminal.Ident, constr.Ident).AppendLine();
              sb.AppendFormat("t.{0} = elements[random.Next(elements.Length)]; ", constr.Ident).EndBlock();
              sb.AppendLine("}");
            } else {
              sb.Append("{").BeginBlock();
              sb.AppendFormat(" var max = problem.GetMax{0}_{1}();", terminal.Ident, constr.Ident).AppendLine();
              sb.AppendFormat(" var min = problem.GetMin{0}_{1}();", terminal.Ident, constr.Ident).AppendLine();
              sb.AppendFormat("t.{0} = random.NextDouble() * (max - min) + min;", constr.Ident).EndBlock();
              sb.AppendLine("}");
            }
          }
          sb.AppendLine("return t;").EndBlock();
          sb.Append("}");
        }
      }
      return sb.ToString();
    }
    private void GenerateSampleTerminalAlternativesStatement(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 GenerateSampleNonterminalAlternativesStatement(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.AppendLine("{");
        sb.AppendFormat("var allIdx = new int[] {{ {0} }}; ", idxStr).AppendLine();
        sb.AppendFormat(
          "var allowedIdx = (from idx in allIdx let targetState = Grammar.transition[state][idx] where Grammar.minDepth[targetState] <= maxDepth select idx).ToArray();")
          .AppendLine();
        sb.AppendLine(
          "if(allowedIdx.Length==0) { allowedIdx = Enumerable.Range(0, Grammar.transition[state].Length).Except(allIdx).ToArray(); } ")
          .AppendLine();
        sb.AppendLine("return allowedIdx[random.Next(allowedIdx.Length)];");
        sb.AppendLine("}");
      }
    }

    private void GenerateReturnStatement(int nAlts, SourceBuilder sb) {
      if (nAlts > 1)
      {
        sb.AppendLine("{");
        sb.AppendFormat(
          "var allowedIdx = (from idx in Enumerable.Range(0, {0}) let targetState = Grammar.transition[state][idx] where Grammar.minDepth[targetState] <= maxDepth select idx).ToArray();", nAlts)
          .AppendLine();
        sb.AppendLine("return allowedIdx[random.Next(allowedIdx.Length)];");
        sb.AppendLine("}");
      } else if (nAlts == 1) {
        sb.AppendLine("return 0; ");
      } else {
        sb.AppendLine("throw new InvalidProgramException();");
      }
    }
  }
}