#region License Information
/* HeuristicLab
 * Copyright (C) 2002-2015 Heuristic and Evolutionary Algorithms Laboratory (HEAL)
 *
 * This file is part of HeuristicLab.
 *
 * HeuristicLab is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * HeuristicLab is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with HeuristicLab. If not, see <http://www.gnu.org/licenses/>.
 */
#endregion

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using HeuristicLab.Common;
using HeuristicLab.Core;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.EvolutionTracking;
using HeuristicLab.Persistence.Default.CompositeSerializers.Storable;

namespace HeuristicLab.Problems.DataAnalysis.Symbolic {
  [Item("TraceCalculator", "Walks a genealogy graph and produces a trace of the specified subtree")]
  [StorableClass]
  public class TraceCalculator : Item {
    private Dictionary<IGenealogyGraphNode<ISymbolicExpressionTree>, TraceData> traceMap;
    private Dictionary<ISymbolicExpressionTree, List<ISymbolicExpressionTreeNode>> nodeListCache;
    private HashSet<Tuple<IGenealogyGraphNode<ISymbolicExpressionTree>, IGenealogyGraphNode<ISymbolicExpressionTree>, int>> traceCache;
    public IGenealogyGraph<ISymbolicExpressionTree> TraceGraph { get; private set; }

    public bool UpdateVertexWeights { get; set; }
    public bool UpdateSubtreeWeights { get; set; }
    public bool CacheTraceNodes { get; set; }

    public TraceCalculator() {
      ResetState();
    }

    protected TraceCalculator(TraceCalculator original, Cloner cloner)
      : base(original, cloner) {
    }

    public override IDeepCloneable Clone(Cloner cloner) {
      return new TraceCalculator(this, cloner);
    }

    public void ResetState() {
      TraceGraph = new GenealogyGraph<ISymbolicExpressionTree>();
      traceMap = new Dictionary<IGenealogyGraphNode<ISymbolicExpressionTree>, TraceData>();
      nodeListCache = new Dictionary<ISymbolicExpressionTree, List<ISymbolicExpressionTreeNode>>();
      traceCache = new HashSet<Tuple<IGenealogyGraphNode<ISymbolicExpressionTree>, IGenealogyGraphNode<ISymbolicExpressionTree>, int>>();
    }

    public static IGenealogyGraph<ISymbolicExpressionTree> TraceSubtree(IGenealogyGraphNode<ISymbolicExpressionTree> node, int subtreeIndex, bool updateVertexWeights = false, bool updateSubtreeWeights = false, bool cacheTraceNodes = true) {
      var tc = new TraceCalculator {
        UpdateVertexWeights = updateSubtreeWeights,
        UpdateSubtreeWeights = updateSubtreeWeights,
        CacheTraceNodes = cacheTraceNodes
      };
      tc.Trace(node, subtreeIndex);
      return tc.TraceGraph;
    }

    public IGenealogyGraph<ISymbolicExpressionTree> Trace(IGenealogyGraphNode<ISymbolicExpressionTree> node, int subtreeIndex, bool resetState = true) {
      if (resetState) ResetState();
      TraceRecursive(node, subtreeIndex);
      return TraceGraph;
    }

    /// <summary>
    /// This method starts from a given vertex in the genealogy graph and works its way 
    /// up the ancestry trying to track the structure of the subtree given by subtreeIndex.
    /// This method will skip genealogy graph nodes that did not have an influence on the 
    /// structure of the tracked subtree. 
    /// 
    /// Only genealogy nodes which did have an influence are added (as copies) to the trace 
    /// and are consequently called 'trace nodes'. 
    /// 
    /// The arcs connecting trace nodes hold information about the locations of the subtrees
    /// and fragments that have been swapped in the form of a tuple (si, fi, lastSi, lastFi),
    /// where:
    /// - si is the subtree index in the current trace node
    /// - fi is the fragment index in the current trace node
    /// - lastSi is the subtree index in the previous trace node
    /// - lastFi is the subtree index in the previous trace node
    /// </summary>
    /// <param name="node">The current node in the genealogy graph</param>
    /// <param name="subtreeIndex">The index of the traced subtree</param>
    /// <param name="last">The last added node in the trace graph</param>
    private void TraceRecursive(IGenealogyGraphNode<ISymbolicExpressionTree> node, int subtreeIndex, IGenealogyGraphNode<ISymbolicExpressionTree> last = null) {
      var g = node;
      int si = subtreeIndex; // subtree index
      int fi = 0; // fragment index
      while (((List<IArc>)((IVertex)g).InArcs).Count > 0) {
        Debug.Assert(si < g.Data.Length);
        var inArcs = (List<IArc>)((IVertex)g).InArcs;
        var fragment = (IFragment<ISymbolicExpressionTreeNode>)((IGenealogyGraphArc)inArcs.Last()).Data;
        if (fragment == null) {
          // TODO: think about what the correct behavior should be here (seems good so far)
          // the node is either an elite node or (in rare cases) no fragment was transferred
          g = (IGenealogyGraphNode<ISymbolicExpressionTree>)inArcs[0].Source;
          continue;
        }

        fi = fragment.Index1; // fragment index
        int fl = fragment.Root.GetLength(); // fragment length
        int sl = NodeAt(g.Data, si).GetLength(); // subtree length

        #region trace crossover
        if (inArcs.Count == 2) {
          var parent0 = (IGenealogyGraphNode<ISymbolicExpressionTree>)inArcs[0].Source;
          var parent1 = (IGenealogyGraphNode<ISymbolicExpressionTree>)inArcs[1].Source;
          if (fi == si) {
            g = parent1;
            si = fragment.Index2;
            continue;
          }
          if (fi < si) {
            if (fi + fl > si) {
              // fragment contains subtree
              g = parent1;
              si += fragment.Index2 - fi;
            } else {
              // fragment distinct from subtree
              g = parent0;
              si += NodeAt(g.Data, fi).GetLength() - fl;
            }
            continue;
          }
          if (fi > si) {
            if (fi < si + sl) {
              // subtree contains fragment => branching point in the fragment graph
              var n = AddTraceNode(g, si, fi); // current node becomes "last" as we restart tracing from the parent
              var t0 = new Tuple<IGenealogyGraphNode<ISymbolicExpressionTree>, IGenealogyGraphNode<ISymbolicExpressionTree>, int>(parent0, n, si);
              if (!(CacheTraceNodes && traceCache.Contains(t0))) {
                TraceRecursive(parent0, si, n);
                traceCache.Add(t0);
              }
              var t1 = new Tuple<IGenealogyGraphNode<ISymbolicExpressionTree>, IGenealogyGraphNode<ISymbolicExpressionTree>, int>(parent1, n, fragment.Index2);
              if (!(CacheTraceNodes && traceCache.Contains(t1))) {
                TraceRecursive(parent1, fragment.Index2, n);
                traceCache.Add(t1);
              }
              break;
            } else {
              // subtree and fragment are distinct. 
              g = parent0;
              continue;
            }
          }
        }
        #endregion
        #region trace mutation
        // mutation is handled in a simple way: we branch every time there is an overlap between the subtree and the fragment
        // (since mutation effects can be quite unpredictable: replace branch, change node, shake tree, etc)
        if (inArcs.Count == 1) {
          var parent0 = (IGenealogyGraphNode<ISymbolicExpressionTree>)inArcs[0].Source;
          Debug.Assert(fragment.Index1 == fragment.Index2);
          // check if the subtree and the fragment overlap => branch out
          if ((si == fi) || (si < fi && fi < si + sl) || (fi < si && si < fi + fl)) {
            var n = AddTraceNode(g, si, fi); // current node becomes "last" as we restart tracing from the parent
            int i = si < fi ? si : fi;
            var t = new Tuple<IGenealogyGraphNode<ISymbolicExpressionTree>, IGenealogyGraphNode<ISymbolicExpressionTree>, int>(parent0, n, i);
            if (!(CacheTraceNodes && traceCache.Contains(t))) {
              TraceRecursive(parent0, i, n);
              traceCache.Add(t);
            }
            break;
          } else {
            // if they don't overlap, go up
            g = parent0;
            if (fi < si)
              si += NodeAt(g.Data, fi).GetLength() - fl;
            continue;
          }
        }
        #endregion
        throw new InvalidOperationException("A node cannot have more than two parents");
      }
      // when we are out of the while the last vertex must be connected with the current one
      // if there is no last vertex, it means the tracing reached the top of the genealogy graph
      var current = AddTraceNode(g, si, fi);
      if (last != null)
        ConnectLast(current, last, si, fi);
    }

    /// <summary>
    /// Get the trace node from the trace graph which corresponds to node g from the genealogy graph.
    /// If the trace graph does not contain such a node, one is created by performing a shallow copy of g, then inserted into the trace graph.
    /// </summary>
    /// <param name="g">The genealogy graph node</param>
    /// <param name="si">The subtree index</param>
    /// <param name="fi">The fragment index</param>
    /// <returns></returns>
    private IGenealogyGraphNode<ISymbolicExpressionTree> AddTraceNode(IGenealogyGraphNode<ISymbolicExpressionTree> g, int si, int fi) {
      var n = TraceGraph.GetByContent(g.Data);
      if (n == null) {
        n = g.Copy();
        TraceGraph.AddVertex(n);
        Debug.Assert(!traceMap.ContainsKey(n));
        traceMap[n] = new TraceData(si, fi, -1, -1); // only the first two fields are needed
      }
      return n;
    }

    // caching node lists brings ~2.5-2.7x speed improvement (since graph nodes are visited multiple times)
    // this caching will be even more effective with larger tree sizes
    private ISymbolicExpressionTreeNode NodeAt(ISymbolicExpressionTree tree, int index) {
      List<ISymbolicExpressionTreeNode> list;
      nodeListCache.TryGetValue(tree, out list);
      if (list == null) {
        list = tree.IterateNodesPrefix().ToList();
        nodeListCache[tree] = list;
      }
      return list[index];
    }

    /// <summary>
    /// Connect the current node of the trace graph with the node that was previously added (@last). The current node of the trace graph is determined by the content
    /// of the genealogy graph node @g.  
    /// </summary>
    /// <param name="current">The current node in the genealogy graph</param>
    /// <param name="last">The last added node in the trace graph</param>   
    /// <param name="si">The index of the traced subtree</param>
    /// <param name="fi">The index of the fragment</param>
    private void ConnectLast(IGenealogyGraphNode<ISymbolicExpressionTree> current, IGenealogyGraphNode<ISymbolicExpressionTree> last, int si, int fi) {
      var lastTraceData = traceMap[last];
      int lastSi = lastTraceData.SubtreeIndex; // last subtree index (index of the traced subtree in the previous trace node)
      int lastFi = lastTraceData.FragmentIndex; // last fragment index (index of the fragment in the previous trace node)
      var td = new TraceData(si, fi, lastSi, lastFi); // trace data

      // TODO: more testing
      var inArcs = (List<IArc>)((IVertex)last).InArcs; // using the InArcs seems to be slightly more efficient than using the OutArcs
      var arc = inArcs.FirstOrDefault(a => a.Source == current && ((IArc<IDeepCloneable>)a).Data.Equals(td));
      if (arc == null) {
        arc = new GenealogyGraphArc(current, last) { Data = td };
        TraceGraph.AddArc(arc);
      }
      if (UpdateVertexWeights) {
        arc.Weight++;
        current.Weight++;
      }
      if (UpdateSubtreeWeights) {
        var subtree = NodeAt(current.Data, td.SubtreeIndex);
        foreach (var s in subtree.IterateNodesPrefix())
          s.NodeWeight++;
      }
    }
  }

  public class TraceData : Tuple<int, int, int, int>, IDeepCloneable {
    public TraceData(int currentSubtreeIndex, int currentFragmentIndex, int lastSubtreeIndex, int lastFragmentIndex)
      : base(currentSubtreeIndex, currentFragmentIndex, lastSubtreeIndex, lastFragmentIndex) {
    }

    public int SubtreeIndex { get { return Item1; } }
    public int FragmentIndex { get { return Item2; } }
    public int LastSubtreeIndex { get { return Item3; } }
    public int LastFragmentIndex { get { return Item4; } }
    public object Clone() {
      return new TraceData(SubtreeIndex, FragmentIndex, LastSubtreeIndex, LastFragmentIndex);
    }

    public IDeepCloneable Clone(Cloner cloner) {
      return cloner.Clone(this);
    }
  }

  internal static class Util {
    // shallow node copy (does not clone the data or the arcs)
    #region some helper methods for shortening the tracing code
    public static IGenealogyGraphNode<ISymbolicExpressionTree> Copy(this IGenealogyGraphNode<ISymbolicExpressionTree> node) {
      return new GenealogyGraphNode<ISymbolicExpressionTree>(node.Data) { Rank = node.Rank, Quality = node.Quality };
    }
    #endregion
  }
}