using System;
using System.Collections.Generic;
using System.Drawing;
using System.Linq;
using HeuristicLab.Encodings.SymbolicExpressionTreeEncoding;
using HeuristicLab.EvolutionaryTracking.Views.Primitives;
using HeuristicLab.Visualization;

namespace HeuristicLab.EvolutionaryTracking.Views {
  public class ExtendedSymbolicExpressionTreeCanvas : ChartControl {
    private const int HSize = 20;
    private const int VSize = 15;
    private const int HDistance = 5;
    private const int VDistance = 5;
    private const int MinGroupHorizontalSpacing = 200;
    private const int MinGroupVerticalSpacing = 200;
    private const int HorizontalGroupDistance = 20;
    private const int VerticalGroupDistance = 20;

    public ReingoldTilfordLayoutEngine<ISymbolicExpressionTreeNode> TreeLayoutEngine { get; set; }
    public ReingoldTilfordLayoutEngine<ISymbolicExpressionTree> GroupsLayoutEngine { get; set; }

    private SymbolicExpressionTreeLayoutAdapter layoutAdapter;

    private Dictionary<ISymbolicExpressionTreeNode, IPrimitive> nodesToPrimitives;
    private Dictionary<ISymbolicExpressionTree, IPrimitive> treesToGroups;

    private Dictionary<Tuple<IPrimitive, IPrimitive>, Line> primitiveConnections;
    private Dictionary<Tuple<IPrimitive, IPrimitive>, Line> groupsConnections;

    private ISymbolicExpressionTreeNode currentSelected;
    private ISymbolicExpressionTreeNode SelectedSubtree {
      get { return currentSelected; }
      set {
        if (Chart == null)
          Chart = new Chart(0, 0, PreferredSize.Width, PreferredSize.Height);

        if (value == currentSelected) return;
        var previousSelected = currentSelected;
        currentSelected = value;

        DisableUpdate();
        if (previousSelected != null) HighlightSubtree(previousSelected, Color.Black); // un-highlight the previously selected subtree
        if (currentSelected != null) HighlightSubtree(currentSelected, Color.DeepSkyBlue); // highlight the new selected subtree
        EnableUpdate();
        EnforceUpdate();
      }
    }

    public ISymbolicExpressionTreeGenealogyGraph GenealogyGraph { get; set; }

    private void ClearDictionaries() {
      if (nodesToPrimitives == null) nodesToPrimitives = new Dictionary<ISymbolicExpressionTreeNode, IPrimitive>();
      nodesToPrimitives.Clear();

      if (treesToGroups == null) treesToGroups = new Dictionary<ISymbolicExpressionTree, IPrimitive>();
      treesToGroups.Clear();

      if (groupsConnections == null) groupsConnections = new Dictionary<Tuple<IPrimitive, IPrimitive>, Line>();
      groupsConnections.Clear();

      if (primitiveConnections == null) primitiveConnections = new Dictionary<Tuple<IPrimitive, IPrimitive>, Line>();
      primitiveConnections.Clear();
    }

    private ISymbolicExpressionTree root;
    public ISymbolicExpressionTree Root {
      get {
        return root;
      }
      set {
        if (value == null)
          throw new ArgumentNullException();

        if (root == value) return;

        ClearDictionaries();
        TreeLayoutEngine.Reset();
        GroupsLayoutEngine.Reset();

        root = value;
        currentSelected = null;

        if (Chart == null)
          Chart = new Chart(0, 0, PreferredSize.Width, PreferredSize.Height);

        DisableUpdate();
        Chart.Group.Clear();
        var group = LayoutTree(root);
        group.Ancestor = group;
        group.Content = root;
        Chart.Group.Add(group);
        GroupsLayoutEngine.Root = group;
        GroupsLayoutEngine.AddNode(group.Content, group);
        treesToGroups.Add(group.Content, group);

        // highlight the received fragment
        var graphNode = GenealogyGraph.GetGraphNodes(root).First();
        if (graphNode.InEdges != null) {
          var arc = graphNode.InEdges.Last(x => x.Source != x.Target);
          var fragment = (IFragment)arc.Data;
          if (fragment != null && fragment.Root != null) {
            HighlightSubtree(fragment.Root, Color.Orange);
          }
        }

        EnableUpdate();
        EnforceUpdate();
      }
    }

    public ExtendedSymbolicExpressionTreeCanvas() {
      TreeLayoutEngine = new ReingoldTilfordLayoutEngine<ISymbolicExpressionTreeNode> { MinHorizontalSpacing = HDistance + HSize, MinVerticalSpacing = VDistance + VSize };
      layoutAdapter = new SymbolicExpressionTreeLayoutAdapter();
      GroupsLayoutEngine = new ReingoldTilfordLayoutEngine<ISymbolicExpressionTree> { MinHorizontalSpacing = MinGroupHorizontalSpacing, MinVerticalSpacing = MinGroupVerticalSpacing };
      if (Chart == null)
        Chart = new Chart(0, 0, PreferredSize.Width, PreferredSize.Height);
      nodesToPrimitives = new Dictionary<ISymbolicExpressionTreeNode, IPrimitive>();
      treesToGroups = new Dictionary<ISymbolicExpressionTree, IPrimitive>();
      groupsConnections = new Dictionary<Tuple<IPrimitive, IPrimitive>, Line>();
      primitiveConnections = new Dictionary<Tuple<IPrimitive, IPrimitive>, Line>();
    }

    /// <summary>
    /// Draws a symbolic expression tree using HeuristicLab.Visualization primitives and the ReingoldTilford layout algorithm.
    /// </summary>
    /// <param name="tree">The symbolic expression tree</param>
    /// <returns>An IChart containing a group of primitives representing the tree.</returns>
    private PrimitiveGroup LayoutTree(ISymbolicExpressionTree tree) {
      var group = LayoutTree(tree.Root.GetSubtree(0).GetSubtree(0));
      group.Content = tree;
      return group;
    }

    // this function is used by the SymbolicExpressionTreeLayoutAdapter
    private ILayoutNode<ISymbolicExpressionTreeNode> ConvertFunction(ISymbolicExpressionTreeNode node) {
      ILayoutNode<ISymbolicExpressionTreeNode> layoutNode;
      if (node.SubtreeCount == 0) {
        layoutNode = new LabeledRectangle(Chart, 0, 0, HSize, VSize) { Text = node.ToString(), Font = new Font(FontFamily.GenericSansSerif, 12) };
      } else {
        layoutNode = new LabeledEllipse(Chart, 0, 0, HSize, VSize) { Text = node.ToString(), Font = new Font(FontFamily.GenericSansSerif, 12) };
      }
      // layoutNode.Children are not initialized here, but when the adapter expands the tree
      layoutNode.Ancestor = layoutNode;
      layoutNode.Content = node;
      return layoutNode;
    }

    private PrimitiveGroup LayoutTree(ISymbolicExpressionTreeNode treeRoot) {
      if (treeRoot == null) throw new ArgumentNullException();

      var layoutNodes = layoutAdapter.Convert(treeRoot, ConvertFunction).ToList(); // returns a list of layoutNodes in breadth order (relative to the tree traversal)
      TreeLayoutEngine.Reset();
      var group = new PrimitiveGroup(Chart);
      // add primitives (rectangles and ellipses corresponding to terminal/function nodes)
      foreach (var layoutNode in layoutNodes) {
        var brush = new SolidBrush(Color.Transparent);
        var pen = new Pen(Color.Black);
        var fontBrush = new SolidBrush(pen.Color);
        var treeNode = layoutNode.Content;
        layoutNode.Number = layoutNode.Parent == null ? 0 : layoutNode.Parent.Content.IndexOfSubtree(layoutNode.Content);
        var primitive = (IPrimitive)layoutNode;
        primitive.Brush = brush;
        primitive.Pen = pen;
        if (treeNode.SubtreeCount == 0) {
          var rectangle = (LabeledRectangle)primitive;
          rectangle.FontBrush = fontBrush;
        } else {
          var ellipse = (LabeledEllipse)primitive;
          ellipse.FontBrush = fontBrush;
        }
        TreeLayoutEngine.AddNode(treeNode, layoutNode);
        group.Add(primitive);
        nodesToPrimitives.Add(treeNode, primitive);
      }
      TreeLayoutEngine.Root = (ILayoutNode<ISymbolicExpressionTreeNode>)nodesToPrimitives[treeRoot];
      TreeLayoutEngine.CalculateLayout(); // calculate X,Y coordinates for each layout node

      foreach (var layoutNode in layoutNodes) {
        var primitive = (RectangularPrimitiveBase)layoutNode;
        var lowerLeft = new PointD(layoutNode.X, layoutNode.Y);
        primitive.SetPosition(lowerLeft, lowerLeft + new Offset(primitive.Size));
      }
      // connect nodes with lines
      foreach (var layoutNode in layoutNodes) {
        var primitive = (RectangularPrimitiveBase)layoutNode;
        var pen = new Pen(Color.Black);
        var node = layoutNode.Content;
        if (node.SubtreeCount == 0) continue;
        foreach (var subtree in node.Subtrees) {
          var subtreePrimitive = (RectangularPrimitiveBase)nodesToPrimitives[subtree];
          var start = new PointD(primitive.UpperRight.X - primitive.Size.Width / 2f, primitive.UpperRight.Y);
          var end = new PointD(subtreePrimitive.LowerLeft.X + subtreePrimitive.Size.Width / 2f, subtreePrimitive.LowerLeft.Y);
          var line = new Line(Chart, start, end, pen);
          group.Add(line);
          primitiveConnections.Add(new Tuple<IPrimitive, IPrimitive>(primitive, subtreePrimitive), line);
        }
      }
      // once the coordonates for each primitive are updated according to the calculated layout, we update the bounds of the group (LowerLeft and UpperRight)
      group.UpdateBounds();
      return group;
    }
    /// <summary>
    /// Handler for the user click event inside the ExtendedSymbolicExpressionTreeCanvas. When the user clicks, the following things should happen:
    /// 1) The group of primitives containing the clicked screen point is identified
    /// 2) The corresponding symbolic expression tree is retrieved from the groupsToTreesDictionary
    /// 3) The corresponding genealogy graph node is retrieved from the GenealogyGraph
    /// 4) The parents and the fragment are retrieved from the graph
    /// 5) Both parents are laid out using the symbolic expression tree layout engine and the resulting groups of primitives are added to the canvas
    /// 6) The fragment is highlighted in the root parent and in the other parent 
    /// </summary>
    /// <param name="sender"></param>
    /// <param name="e"></param>
    protected override void pictureBox_MouseUp(object sender, System.Windows.Forms.MouseEventArgs e) {
      // STEP 1: Get the primitives and the group
      var primitives = Chart.GetAllPrimitives(e.Location).Where(p => p is PrimitiveGroup).ToList();
      if (primitives.Count == 0) return;
      if (primitives.Count > 1) throw new Exception("Groups/trees should not overlap."); // in the extended canvas, the PrimitiveGroups (each representing a symbolic expression tree) should not overlap
      var selectedGroup = (PrimitiveGroup)primitives[0];
      var selectedPrimitive = selectedGroup.GetAllPrimitives(Chart.TransformPixelToWorld(e.Location)).First(p => p is ILayoutNode<ISymbolicExpressionTreeNode>); // similarly with the above condition, there should only be one primitive in that location
      // STEP 2: Get the selected tree and the selected subtree
      SelectedSubtree = ((ILayoutNode<ISymbolicExpressionTreeNode>)selectedPrimitive).Content;
      var selectedTree = ((ILayoutNode<ISymbolicExpressionTree>)selectedGroup).Content;
      // STEP 3: Get the genealogy graph node, the parents and the fragment
      var graphNode = GenealogyGraph.GetGraphNodes(selectedTree).First(); // get the GenealogyGraphNode corresponding to the tree
      var parentVertices = graphNode.InEdges.Select(edge => (SymbolicExpressionTreeGenealogyGraphNode)edge.Source).ToList();

      DisableUpdate();

      foreach (var edge in graphNode.InEdges) {
        var parentVertex = (SymbolicExpressionTreeGenealogyGraphNode)edge.Source;
        var parent = parentVertex.SymbolicExpressionTree;
        if (treesToGroups.ContainsKey(parent)) continue;
        var childGroup = LayoutTree(parent); // from the layout perspective, roles are reversed and genealogical parents become layout children
        treesToGroups.Add(parent, childGroup);
        Chart.Group.Add(childGroup);
        if (selectedGroup.Children == null) selectedGroup.Children = new List<ILayoutNode<ISymbolicExpressionTree>>();
        selectedGroup.Children.Add(childGroup);
        childGroup.Parent = selectedGroup;
        childGroup.Level = childGroup.Parent.Level + 1;
        childGroup.Number = childGroup.Parent.Children.Count - 1;
        childGroup.Ancestor = childGroup;
        childGroup.Content = parent;
        GroupsLayoutEngine.AddNode(childGroup.Content, childGroup);
      }
      var parents = parentVertices.Select(x => x.SymbolicExpressionTree).ToList();
      // highlight the fragment
      if (graphNode.InEdges.Count == 2) {
        // highlight the fragment
        var fragment = (IFragment)graphNode.InEdges.Last().Data;
        if (!(fragment == null || fragment.Root == null)) {

          var subtree0 = parents[0].IterateNodesPrefix().ElementAt(fragment.Index);
          HighlightSubtree(subtree0, Color.Red);

          var subtree1 = parents[1].IterateNodesPrefix().ElementAt(fragment.OldIndex);
          HighlightSubtree(subtree1, Color.LimeGreen);
        }
      }
      // layout the groups
      var groups = Chart.Group.Primitives.Where(p => p is PrimitiveGroup).Cast<PrimitiveGroup>().ToList();
      var maxSize = new SizeF(groups.Max(g => g.Size.Width), groups.Max(g => g.Size.Height));
      GroupsLayoutEngine.MinHorizontalSpacing = maxSize.Width + HorizontalGroupDistance;
      GroupsLayoutEngine.MinVerticalSpacing = maxSize.Height + VerticalGroupDistance;

      GroupsLayoutEngine.CalculateLayout();
      foreach (var group in Chart.Group.Primitives.Where(p => p is PrimitiveGroup).Cast<PrimitiveGroup>()) {
        group.Move(new PointD(group.X, group.Y) - group.LowerLeft);
      }
      // draw or update (start,end) coordinates of group connections
      foreach (var group in Chart.Group.Primitives.Where(p => p is PrimitiveGroup).Cast<PrimitiveGroup>()) {
        if (group.Children == null) continue;
        var start = new PointD(group.UpperRight.X - group.Size.Width / 2, group.UpperRight.Y);
        foreach (var childGroup in group.Children.Cast<PrimitiveGroup>()) {
          var tuple = new Tuple<IPrimitive, IPrimitive>(group, childGroup);
          var childRoot = (RectangularPrimitiveBase)nodesToPrimitives[childGroup.Content.Root.GetSubtree(0).GetSubtree(0)];
          var end = new PointD(childRoot.LowerLeft.X + childRoot.Size.Width / 2, childRoot.LowerLeft.Y);
          if (groupsConnections.ContainsKey(tuple)) {
            var line = groupsConnections[tuple];
            line.SetPosition(start, end);
          } else {
            var line = new Line(Chart, start, end) { Pen = new Pen(Color.Black) };
            Chart.Group.Add(line);
            groupsConnections.Add(tuple, line);
          }
        }
      }

      EnableUpdate();
      EnforceUpdate();

      base.pictureBox_MouseUp(sender, e);
    }
    // highlight a given subtree with the specified color
    public void HighlightSubtree(ISymbolicExpressionTreeNode subtree, Color color) {
      foreach (var node in subtree.IterateNodesBreadth()) {
        var primitive = nodesToPrimitives[node];
        primitive.Pen.Brush = new SolidBrush(color);
      }
    }
    // the (0,0) point of the HeuristicLab.Visualization.Chart is at the bottom left corner of the screen, while the (0,0) point of the layout engine 
    // and of the actual user control is at the top left corner. so we have to vertically flip everything to maintain correct orientation.
    private void FlipVertical() {
      foreach (var primitive in Chart.Group.Primitives) {
        var group = primitive as PrimitiveGroup;
        if (group != null) {
          foreach (var groupPrimitive in group.Primitives) {
            var rp = groupPrimitive as RectangularPrimitiveBase;
            if (rp != null) {
              rp.SetPosition(rp.LowerLeft.X, Height - rp.UpperRight.Y, rp.UpperRight.X, Height - rp.LowerLeft.Y);
            } else {
              var lp = groupPrimitive as LinearPrimitiveBase;
              if (lp != null) {
                lp.SetPosition(lp.Start.X, Height - lp.Start.Y, lp.End.X, Height - lp.End.Y);
              }
            }
          }
        } else {
          var lp = primitive as LinearPrimitiveBase;
          if (lp != null) {
            lp.SetPosition(lp.Start.X, Height - lp.Start.Y, lp.End.X, Height - lp.End.Y);
          }
        }
      }
    }
    // some convenience methods for the lazy (enabling, disabling updates for the Chart)
    private void EnableUpdate() {
      Chart.UpdateEnabled = true;
    }
    private void DisableUpdate() {
      Chart.UpdateEnabled = false;
    }
    private void EnforceUpdate() {
      DisableUpdate();
      FlipVertical();
      EnableUpdate();
      Chart.EnforceUpdate();
    }
  }
}