using System;
using System.Collections;
using System.Collections.Generic;
using System.Collections.Specialized;
using System.Diagnostics;
using System.Linq;
using HeuristicLab.Persistence.Default.CompositeSerializers.Storable;

namespace HeuristicLab.Data.PersistentDataStructures.Implementations {

  [StorableClass]
  public class ArrayMappedTrie<T> : IEnumerable<T> {

    [StorableClass]
    private class Node : IEnumerable {

      [Storable]
      private BitVector32 bitmap;
      [Storable]
      private readonly object[] values;

      private static UInt16 CalculateLocalIndex(UInt16 subIndex, BitVector32 bitmap) {
        UInt32 relevantBits = ((((uint)1)<<subIndex) - 1) & (uint) bitmap.Data;
        return (UInt16) PopCount(relevantBits);
      }

      public bool Contains(UInt16 localIndex) { return bitmap[1 << localIndex]; }

      [StorableConstructor]
      protected Node(bool isDeserializing) { }

      protected Node(BitVector32 bitmap, object[] values) {
        this.bitmap = bitmap;
        this.values = values;
      }

      public Node(UInt16 subIndex, object value) {
        bitmap[1 << subIndex] = true;
        values = new[] {value};
      }

      public object Get(UInt16 subIndex) {
        return Contains(subIndex) ? values[CalculateLocalIndex(subIndex, bitmap)] : null;
      }

      public Node Set(UInt16 subIndex, object value) {
        Debug.Assert(subIndex < 32);
        if (Contains(subIndex)) {
          var newValues = new object[values.Length];
          Array.Copy(values, newValues, newValues.Length);
          newValues[CalculateLocalIndex(subIndex, bitmap)] = value;
          return new Node(bitmap, newValues);
        } else {
          var newValues = new object[values.Length + 1];
          BitVector32 newBitmap = bitmap;
          newBitmap[1 << subIndex] = true;
          var localIndex = CalculateLocalIndex(subIndex, newBitmap);
          if (localIndex > 0)
            Array.ConstrainedCopy(values, 0, newValues, 0, localIndex);
          newValues[localIndex] = value;
          if (values.Length - localIndex > 0)
            Array.ConstrainedCopy(values, localIndex, newValues, localIndex + 1, values.Length - localIndex);
          return new Node(newBitmap, newValues);
        }
      }

      public Node Remove(UInt16 subIndex) {
        Debug.Assert(subIndex < 32);
        if (!Contains(subIndex)) return this;
        if (Length == 1) return null;
        var newBitmap = bitmap;
        var localIndex = CalculateLocalIndex(subIndex, bitmap);
        newBitmap[1 << subIndex] = false;
        var newValues = new object[values.Length - 1];
        Array.ConstrainedCopy(values, 0, newValues, 0, localIndex);
        Array.ConstrainedCopy(values, localIndex + 1, newValues, localIndex, newValues.Length - localIndex);
        return new Node(newBitmap, newValues);
      }

      public UInt16 Length { get { return (UInt16) values.Length; } }

      public IEnumerator GetEnumerator() {
        return values.GetEnumerator();
      }
    }

    [Storable]
    private Node root;

    [Storable]
    private List<Node> oldRoots;

    [Storable]
    private UInt32 stepsSinceSnapshot = UInt32.MaxValue;

    [Storable]
    private readonly bool assumeFilledWithDefaultValue;

    public ArrayMappedTrie(bool assumeFilledWithDefaultValue = true) {
      this.assumeFilledWithDefaultValue = assumeFilledWithDefaultValue;
      root = null;
      ClonedFrom = null;
      TrackClonedFrom = true;
      SnapshotInterval = 1;
    }

    public T this[UInt32 index] {
      get { return Get(root, index); }
      set { Set(index, value); }
    }

    public UInt32 Count { get { return DoCount(root); } }

    public UInt32 SnapshotInterval { get; set; }

    public bool TrackClonedFrom { get; set; }

    public ArrayMappedTrie<T> ClonedFrom { get; private set; }

    public void CreateSnapshot() {
      if (oldRoots == null) oldRoots = new List<Node>();
      oldRoots.Add(root);
      stepsSinceSnapshot = 0;
    }

    public void Remove(UInt32 index) {
      CreateSnapshotIfNeccessary();
      Remove(root, index, 6);
      stepsSinceSnapshot++;
    }

    public void Clear() {
      CreateSnapshotIfNeccessary();
      root = null;
      stepsSinceSnapshot++;
    }

    private void CreateSnapshotIfNeccessary() {
      if (SnapshotInterval > 0 && stepsSinceSnapshot > SnapshotInterval)
        CreateSnapshot();
    }

    private void Set(UInt32 index, T value) {
      if (assumeFilledWithDefaultValue) {
        T oldValue;
        var containsIndex = Get(root, index, out oldValue);
        if (containsIndex && oldValue.Equals(value) ||
            !containsIndex && value.Equals(default(T))) return;
      }
      CreateSnapshotIfNeccessary();
      root = assumeFilledWithDefaultValue && value.Equals(default(T)) ? Remove(root, index, 6) : Set(root, index, 6, value);
      stepsSinceSnapshot++;
    }

    private static Node Set(Node n, UInt32 index, UInt16 level, T value) {
      var subIndex = ExtractSubIndex(index, level);
      if (level == 0)
        return n == null ?
           new Node(subIndex, value) :
           n.Set(subIndex, value);
      return n == null ?
        new Node(subIndex, Set(null, index, (UInt16) (level - 1), value)) :
        n.Set(subIndex, Set((Node)n.Get(subIndex), index, (UInt16)(level - 1), value));
    }

    private static Node Remove(Node n, UInt32 index, UInt16 level) {
      if (n == null) return null;
      var subIndex = ExtractSubIndex(index, level);
      if (level == 0) return n.Remove(subIndex);
      var newN = Remove((Node) n.Get(subIndex), index, (UInt16) (level - 1));
      return newN == null ? n.Remove(subIndex) : n.Set(subIndex, newN);
    }

    private bool Get(Node n, UInt32 index, out T value) {
      value = default(T);
      var leafNode = Get(root, index, 6);
      if (leafNode == null) return false;
      var subIndex = ExtractSubIndex(index, 0);
      if (!leafNode.Contains(subIndex)) return false;
      value = (T)leafNode.Get(ExtractSubIndex(index, 0));
      return true;
    }

    private T Get(Node n, UInt32 index) {
      T value;
      if (Get(n, index, out value) || assumeFilledWithDefaultValue)
        return value;
      throw new IndexOutOfRangeException();
    }

    private static Node Get(Node n, UInt32 index, UInt16 level) {
      while (true) {
        var subIndex = ExtractSubIndex(index, level);
        if (level == 0) return n;
        if (n == null) return null;
        n = (Node) n.Get(subIndex);
        level = (UInt16) (level - 1);
      }
    }

    private UInt32 DoCount(Node n) { return DoCount(root, 6); }

    private static UInt32 DoCount(Node n, UInt16 level) {
      if (n == null) return 0;
      return level == 0 ? n.Length :
        n.Cast<Node>().Aggregate<Node, uint>(0,
        (current, child) =>
          current + DoCount(child, (UInt16) (level - 1)));
    }

    public UInt32 CountNodes() { return CountNodes(root, 6);  }

    private static UInt32 CountNodes(Node n, UInt16 level) {
      if (n == null) return 0;
      if (level == 0) return 1;
      return 1 + n.Cast<Node>().Aggregate<Node, uint>(0,
        (current, child) =>
          current + CountNodes(child, (UInt16) (level - 1)));
    }

    public UInt32 Overhead() { return Overhead(root, 6); }

    private static UInt32 Overhead(Node n, UInt16 level) {
      if (n == null) return 0;
      if (level == 0) return 1 + 1; /* bitmap + child, leaf arrays are not overhead */
      UInt32 total = 0;
      foreach (Node child in n) {
        total += Overhead(child, (UInt16)(level - 1));
      }
      return 1 + 1 + (UInt32)n.Length + total; /* node + bitmap + array + children */
    }

    public long SizeWithHistory() {
      var nodes = new Dictionary<Node, UInt16>();
      AddNodeSizeRecursively(root, nodes);
      var size = nodes.Sum(kvp => (long) kvp.Value);
      if (oldRoots != null) {
        foreach (var node in oldRoots) {
          AddNodeSizeRecursively(node, nodes);
        }
      }
      var sizeWithHistory = nodes.Sum(kvp => (long)kvp.Value);
      Console.WriteLine("size = {0}, size with history = {1}", size, sizeWithHistory);
      return sizeWithHistory;
    }

    private static void AddNodeSizeRecursively(Node n, IDictionary<Node, UInt16> d) {
      if (n == null) return;
      if (d.ContainsKey(n)) return;
      d[n] = (UInt16)(n.Length + 1 + 1);
      foreach (var child in n) {
        var childNode = child as Node;
        if (child == null) break;
        AddNodeSizeRecursively(childNode, d);
      }
    }

    public ArrayMappedTrie<T> Clone() {
      return new ArrayMappedTrie<T> {
        TrackClonedFrom = this.TrackClonedFrom,
        ClonedFrom = TrackClonedFrom ? this : null,
        oldRoots = new List<Node> {root},
        root = root
      };
    } 

    public static UInt16 ExtractSubIndex(UInt32 index, int level) {
      int offset = level*5;
      const uint window = 0x1f;
      UInt32 shiftedWindow = window << offset;
      UInt32 maskedValue = index & shiftedWindow;
      UInt32 subindex = maskedValue >> offset;
      Debug.Assert(subindex < 32);
      // Console.WriteLine("index {0} level {1} => {2} ({3})", index, level, subindex, subindex * Math.Pow(32, level));
      return (UInt16)subindex;
    }

    private static IEnumerable<Node> GetChildrenRecursively(Node n, int level) {
      if (level == 0) yield return n;
      foreach (Node child in n) {
        foreach (var grandchild in GetChildrenRecursively(child, level-1)) {
          yield return grandchild;
        }
      }
    }

    IEnumerator IEnumerable.GetEnumerator() { return GetEnumerator(); }

    public IEnumerator<T> GetEnumerator() {
      foreach (var leaf in GetChildrenRecursively(root, 6)) {
        foreach (var value in leaf)
          yield return (T)value;
      }
    }

    public const UInt32 SK5  = 0x55555555;
    public const UInt32 SK3  = 0x33333333;
    public const UInt32 SKF0 = 0x0f0f0f0f;
    public const UInt32 SKFF = 0x00ff00ff;

    public static UInt32 PopCount(UInt32 x) {
      x -= (x>>1)&SK5;
      x = (x & SK3) + ((x >> 2) & SK3);
      x = (x & SKF0) + ((x >> 4) & SKF0);
      x += x >> 8;
      return (x + (x >> 16)) & 0x3F;
    }

  }
}