using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using HeuristicLab.Algorithms.DataAnalysis.SymRegGrammarEnumeration.GrammarEnumeration;
using HeuristicLab.Common;
using HeuristicLab.Core;
using HeuristicLab.Data;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.Optimization;
using HeuristicLab.Parameters;
using HeuristicLab.Persistence.Default.CompositeSerializers.Storable;
using HeuristicLab.Problems.DataAnalysis;
using HeuristicLab.Problems.DataAnalysis.Symbolic;
using HeuristicLab.Problems.DataAnalysis.Symbolic.Regression;

namespace HeuristicLab.Algorithms.DataAnalysis.SymRegGrammarEnumeration {
  [Item("Grammar Enumeration Symbolic Regression", "Iterates all possible model structures for a fixed grammar.")]
  [StorableClass]
  [Creatable(CreatableAttribute.Categories.DataAnalysisRegression, Priority = 250)]
  public class GrammarEnumerationAlgorithm : FixedDataAnalysisAlgorithm<IRegressionProblem> {
    #region properties and result names
    private readonly string BestTrainingQualityName = "Best R² (Training)";
    private readonly string BestTrainingSolutionName = "Best solution (Training)";
    private readonly string GeneratedSentencesName = "Generated Sentences";
    private readonly string DistinctSentencesName = "Distinct Sentences";
    private readonly string PhraseExpansionsName = "Phrase Expansions";
    private readonly string AverageTreeLengthName = "Avg. Sentence Length among Distinct";
    private readonly string GeneratedEqualSentencesName = "Generated equal sentences";


    private readonly string SearchDataStructureParameterName = "Search Data Structure";
    private readonly string MaxTreeSizeParameterName = "Max. Tree Nodes";
    private readonly string GuiUpdateIntervalParameterName = "GUI Update Interval";

    public override bool SupportsPause { get { return false; } }

    protected IValueParameter<IntValue> MaxTreeSizeParameter {
      get { return (IValueParameter<IntValue>)Parameters[MaxTreeSizeParameterName]; }
    }
    public int MaxTreeSize {
      get { return MaxTreeSizeParameter.Value.Value; }
      set { MaxTreeSizeParameter.Value.Value = value; }
    }

    protected IValueParameter<IntValue> GuiUpdateIntervalParameter {
      get { return (IValueParameter<IntValue>)Parameters[GuiUpdateIntervalParameterName]; }
    }
    public int GuiUpdateInterval {
      get { return GuiUpdateIntervalParameter.Value.Value; }
      set { GuiUpdateIntervalParameter.Value.Value = value; }
    }

    protected IValueParameter<EnumValue<StorageType>> SearchDataStructureParameter {
      get { return (IValueParameter<EnumValue<StorageType>>)Parameters[SearchDataStructureParameterName]; }
    }
    public StorageType SearchDataStructure {
      get { return SearchDataStructureParameter.Value.Value; }
      set { SearchDataStructureParameter.Value.Value = value; }
    }

    public SymbolString BestTrainingSentence;

    #endregion

    public Dictionary<int, SymbolString> DistinctSentences;   // Semantically distinct sentences in a run.
    public Dictionary<int, List<SymbolString>> AllSentences;  // All sentences ever generated in a run.
    public Dictionary<int, SymbolString> ArchivedPhrases;     // Nodes in the search tree
    SearchDataStore OpenPhrases;                              // Stack/Queue/etc. for fetching the next node in the search tree.  

    public int EqualGeneratedSentences; // It is not guaranteed that shorter solutions are found first. 
                                        // When longer solutions are overwritten with shorter ones, 
                                        // this counter is increased.
    public int Expansions;              // Number, how many times a nonterminal symbol is replaced with a production rule.
    public Grammar Grammar;

    public StringBuilder dotFileBuilder;

    #region ctors
    public override IDeepCloneable Clone(Cloner cloner) {
      return new GrammarEnumerationAlgorithm(this, cloner);
    }

    public GrammarEnumerationAlgorithm() {
      var provider = new HeuristicLab.Problems.Instances.DataAnalysis.NguyenInstanceProvider(seed: 1234);
      var regProblem = provider.LoadData(provider.GetDataDescriptors().Single(x => x.Name.Contains("F9 ")));

      Problem = new RegressionProblem() {
        ProblemData = regProblem
      };

      Parameters.Add(new ValueParameter<IntValue>(MaxTreeSizeParameterName, "The number of clusters.", new IntValue(6)));
      Parameters.Add(new ValueParameter<IntValue>(GuiUpdateIntervalParameterName, "Number of generated sentences, until GUI is refreshed.", new IntValue(4000)));

      Parameters.Add(new ValueParameter<EnumValue<StorageType>>(SearchDataStructureParameterName, new EnumValue<StorageType>(StorageType.Stack)));
    }

    public GrammarEnumerationAlgorithm(GrammarEnumerationAlgorithm original, Cloner cloner) : base(original, cloner) { }
    #endregion

    protected override void Run(CancellationToken cancellationToken) {
      #region init
      InitResults();

      dotFileBuilder = new StringBuilder();

      AllSentences = new Dictionary<int, List<SymbolString>>();
      ArchivedPhrases = new Dictionary<int, SymbolString>(); // Nodes in the search tree
      DistinctSentences = new Dictionary<int, SymbolString>();
      Expansions = 0;
      EqualGeneratedSentences = 0;

      Grammar = new Grammar(Problem.ProblemData.AllowedInputVariables.ToArray());

      OpenPhrases = new SearchDataStore(SearchDataStructure); // Select search strategy
      var phrase0 = new SymbolString(new[] { Grammar.StartSymbol });
      #endregion

      OpenPhrases.Store(Grammar.CalcHashCode(phrase0), phrase0);

      while (OpenPhrases.Count > 0) {
        if (cancellationToken.IsCancellationRequested) break;

        StoredSymbolString fetchedPhrase = OpenPhrases.GetNext();
        SymbolString currPhrase = fetchedPhrase.SymbolString;

        if (dotFileBuilder.Length == 0) {
          dotFileBuilder.AppendLine("digraph searchgraph {");
          dotFileBuilder.AppendLine(fetchedPhrase.Hash + " [label=\"" + Grammar.PostfixToInfixParser(fetchedPhrase.SymbolString) + "\", shape=doublecircle];");
        } else {
          dotFileBuilder.AppendLine(fetchedPhrase.Hash + " [label=\"" + Grammar.PostfixToInfixParser(fetchedPhrase.SymbolString) + "\"];");
        }

        ArchivedPhrases.Add(fetchedPhrase.Hash, currPhrase);

        // expand next nonterminal symbols
        int nonterminalSymbolIndex = currPhrase.FindIndex(s => s is NonterminalSymbol);
        NonterminalSymbol expandedSymbol = currPhrase[nonterminalSymbolIndex] as NonterminalSymbol;

        foreach (Production productionAlternative in expandedSymbol.Alternatives) {
          SymbolString newPhrase = new SymbolString(currPhrase);
          newPhrase.RemoveAt(nonterminalSymbolIndex);
          newPhrase.InsertRange(nonterminalSymbolIndex, productionAlternative);

          Expansions++;
          if (newPhrase.Count <= MaxTreeSize) {
            var phraseHash = Grammar.CalcHashCode(newPhrase);

            dotFileBuilder.AppendLine(fetchedPhrase.Hash + " -> " + phraseHash + " [label=\"" + expandedSymbol.StringRepresentation + "&rarr;" + productionAlternative + "\"];");

            if (newPhrase.IsSentence()) { // Sentence was generated.
              SaveToAllSentences(phraseHash, newPhrase);

              if (!DistinctSentences.ContainsKey(phraseHash) || DistinctSentences[phraseHash].Count > newPhrase.Count) {
                if (DistinctSentences.ContainsKey(phraseHash)) EqualGeneratedSentences++; // for analysis only

                DistinctSentences[phraseHash] = newPhrase;
                EvaluateSentence(newPhrase);

                dotFileBuilder.AppendLine(phraseHash + " [label=\"" + Grammar.PostfixToInfixParser(newPhrase) + "\", style=\"filled\"];");
              }
              UpdateView(AllSentences, DistinctSentences, Expansions);

            } else if (!OpenPhrases.Contains(phraseHash) && !ArchivedPhrases.ContainsKey(phraseHash)) {
              OpenPhrases.Store(phraseHash, newPhrase);
            }
          }
        }
      }

      dotFileBuilder.AppendLine(Grammar.CalcHashCode(BestTrainingSentence) + " [label=\"" + Grammar.PostfixToInfixParser(BestTrainingSentence) + "\", shape=Mcircle, style=\"filled,bold\"];");
      dotFileBuilder.AppendLine("}");

      System.IO.File.WriteAllText(Environment.GetFolderPath(System.Environment.SpecialFolder.DesktopDirectory)+ @"\searchgraph.dot", dotFileBuilder.ToString());

      UpdateView(AllSentences, DistinctSentences, Expansions, force: true);
      UpdateFinalResults();
    }

    // Store sentence to "MultiDictionary"
    private void SaveToAllSentences(int hash, SymbolString sentence) {
      if (AllSentences.ContainsKey(hash))
        AllSentences[hash].Add(sentence);
      else
        AllSentences[hash] = new List<SymbolString> { sentence };
    }

    #region Evaluation of generated models.

    // Evaluate sentence within an algorithm run.
    private void EvaluateSentence(SymbolString symbolString) {
      SymbolicExpressionTree tree = Grammar.ParseSymbolicExpressionTree(symbolString);
      SymbolicRegressionModel model = new SymbolicRegressionModel(
        Problem.ProblemData.TargetVariable,
        tree,
        new SymbolicDataAnalysisExpressionTreeLinearInterpreter());

      var probData = Problem.ProblemData;
      var target = probData.TargetVariableTrainingValues;
      var estVals = model.GetEstimatedValues(probData.Dataset, probData.TrainingIndices);
      OnlineCalculatorError error;
      var r2 = OnlinePearsonsRSquaredCalculator.Calculate(target, estVals, out error);
      if (error != OnlineCalculatorError.None) r2 = 0.0;

      var bestR2 = ((DoubleValue)Results[BestTrainingQualityName].Value).Value;
      if (r2 > bestR2) {
        ((DoubleValue)Results[BestTrainingQualityName].Value).Value = r2;
        BestTrainingSentence = symbolString;
      }
    }

    #endregion

    #region Visualization in HL
    // Initialize entries in result set.
    private void InitResults() {
      BestTrainingSentence = null;

      Results.Add(new Result(BestTrainingQualityName, new DoubleValue(-1.0)));

      Results.Add(new Result(GeneratedSentencesName, new IntValue(0)));
      Results.Add(new Result(DistinctSentencesName, new IntValue(0)));
      Results.Add(new Result(PhraseExpansionsName, new IntValue(0)));
      Results.Add(new Result(GeneratedEqualSentencesName, new IntValue(0)));
      Results.Add(new Result(AverageTreeLengthName, new DoubleValue(1.0)));
    }

    // Update the view for intermediate results in an algorithm run.
    private int updates;
    private void UpdateView(Dictionary<int, List<SymbolString>> allGenerated,
        Dictionary<int, SymbolString> distinctGenerated, int expansions, bool force = false) {
      updates++;

      if (force || updates % GuiUpdateInterval == 0) {
        var allGeneratedEnum = allGenerated.Values.SelectMany(x => x).ToArray();
        ((IntValue)Results[GeneratedSentencesName].Value).Value = allGeneratedEnum.Length;
        ((IntValue)Results[DistinctSentencesName].Value).Value = distinctGenerated.Count;
        ((IntValue)Results[PhraseExpansionsName].Value).Value = expansions;
        ((IntValue)Results[GeneratedEqualSentencesName].Value).Value = EqualGeneratedSentences;
        ((DoubleValue)Results[AverageTreeLengthName].Value).Value = allGeneratedEnum.Select(sentence => sentence.Count).Average();
      }
    }

    // Generate all Results after an algorithm run.
    private void UpdateFinalResults() {
      SymbolicExpressionTree tree = Grammar.ParseSymbolicExpressionTree(BestTrainingSentence);
      SymbolicRegressionModel model = new SymbolicRegressionModel(
        Problem.ProblemData.TargetVariable,
        tree,
        new SymbolicDataAnalysisExpressionTreeLinearInterpreter());

      IRegressionSolution bestTrainingSolution = new RegressionSolution(model, Problem.ProblemData);
      Results.AddOrUpdateResult(BestTrainingSolutionName, bestTrainingSolution);

      // Print generated sentences.
      string[,] sentencesMatrix = new string[AllSentences.Values.SelectMany(x => x).Count(), 3];

      int i = 0;
      foreach (var sentenceSet in AllSentences) {
        foreach (var sentence in sentenceSet.Value) {
          sentencesMatrix[i, 0] = sentence.ToString();
          sentencesMatrix[i, 1] = Grammar.PostfixToInfixParser(sentence).ToString();
          sentencesMatrix[i, 2] = sentenceSet.Key.ToString();
          i++;
        }
      }
      Results.Add(new Result("All generated sentences", new StringMatrix(sentencesMatrix)));

      string[,] distinctSentencesMatrix = new string[DistinctSentences.Count, 3];
      i = 0;
      foreach (KeyValuePair<int, SymbolString> distinctSentence in DistinctSentences) {
        distinctSentencesMatrix[i, 0] = distinctSentence.Key.ToString();
        distinctSentencesMatrix[i, 1] = Grammar.PostfixToInfixParser(distinctSentence.Value).ToString();
        distinctSentencesMatrix[i, 2] = distinctSentence.Key.ToString();
        i++;
      }
      Results.Add(new Result("Distinct generated sentences", new StringMatrix(distinctSentencesMatrix)));
    }
    #endregion

  }
}