package edu.stanford.nlp.ling.tokensregex.matcher;

import edu.stanford.nlp.util.Generics;
import edu.stanford.nlp.util.StringUtils;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Map that takes a iterable as key, and maps it to an value.
 *
 * This implementation is not particularly memory efficient, but will have relatively
 * fast lookup times for sequences where there are many possible keys (e.g. sequences over Strings).
 * Can be used for fairly efficient look up of sequence by prefix.
 *
 * @author Angel Chang
 *
 * @param <K> Key of the elements of the trie
 * @param <V> The value
 *
 */
public class TrieMap<K, V> extends AbstractMap<Iterable<K>, V> {

  /**
   * Child tries
   */
  protected Map<K, TrieMap<K, V>> children;

  /**
   * Value at a leaf node (leaf node is indicated by non-null value)
   */
  protected V value;
  // Should we have explicit marking if this element is a leaf node without requiring value?

  public TrieMap() {
  }

  public TrieMap(int initialCapacity) {
    // TODO: initial capacity implementation
  }

  // Trie specific functions
  public TrieMap<K,V> getChildTrie(K key) {
    return (children != null)? children.get(key):null;
  }

  public TrieMap<K,V> getChildTrie(Iterable<K> key) {
    TrieMap<K, V> curTrie = this;
    // go through each element
    for (K element : key){
      curTrie = (curTrie.children != null)? curTrie.children.get(element):null;
      if (curTrie == null) {
        return null;
      }
    }
    return curTrie;
  }

  public TrieMap<K,V> putChildTrie(Iterable<K> key, TrieMap<K,V> child) {
    TrieMap<K,V> parentTrie = null;
    TrieMap<K, V> curTrie = this;
    Iterator<K> keyIter = key.iterator();
    // go through each element
    while (keyIter.hasNext()) {
      K element = keyIter.next();
      boolean isLast = !keyIter.hasNext();
      if (curTrie.children == null) {
        curTrie.children = new ConcurrentHashMap<>();//Generics.newConcurrentHashMap();
      }
      parentTrie = curTrie;
      curTrie = curTrie.children.get(element);
      if (isLast) {
        parentTrie.children.put(element, child);
      } else if(curTrie == null){
        parentTrie.children.put(element, curTrie = new TrieMap<>());
      }
    }
    if (parentTrie == null) {
      throw new IllegalArgumentException("Cannot put a child trie with no keys");
    }
    return curTrie;
  }

  public Map<K, TrieMap<K, V>> getChildren() {
    return children;
  }

  public V getValue() {
    return value;
  }

  public boolean isLeaf() {
    return value != null;
  }


  public String toFormattedString() {
    List<String> strings = new LinkedList<>();
    updateTrieStrings(strings, "");
    return StringUtils.join(strings, "\n");
  }

  protected void updateTrieStrings(List<String> strings, String prefix) {
    if (children != null) {
      for (Entry<K, TrieMap<K, V>> kTrieMapEntry : children.entrySet()) {
        kTrieMapEntry.getValue().updateTrieStrings(strings, prefix + " - " + kTrieMapEntry.getKey());
      }
    }
    if (isLeaf()) {
      strings.add( prefix + " -> " + value);
    }
  }

  // Functions to support map interface to lookup using sequence

  @Override
  public int size() {
    int s = 0;
    if (children != null) {
      for (Entry<K, TrieMap<K, V>> kTrieMapEntry : children.entrySet()) {
        s += kTrieMapEntry.getValue().size();
      }
    }
    if (isLeaf()) s++;
    return s;
  }

  @Override
  public boolean isEmpty() {
    return (children == null && !isLeaf());
  }

  @Override
  public boolean containsKey(Object key) {
    return get(key) != null;
  }

  @Override
  public boolean containsValue(Object value) {
    return values().contains(value);
  }

  @Override
  public V get(Object key) {
    if (key instanceof Iterable) {
      return get( (Iterable<K>) key);
    } else if (key instanceof Object[]) {
      return get( Arrays.asList( (Object[]) key) );
    }
    return null;
  }

  public V get(Iterable<K> key) {
    TrieMap<K, V> curTrie = getChildTrie(key);
    return (curTrie != null) ? curTrie.value: null;
  }

  public V get(K[] key) {
    return get(Arrays.asList(key));
  }

  @Override
  public V put(Iterable<K> key, V value) {
    if (value == null) throw new IllegalArgumentException("Value cannot be null");
    TrieMap<K, V> curTrie = this;
    // go through each element
    for(K element:key){
      if (curTrie.children == null) {
        curTrie.children = new ConcurrentHashMap<>();//Generics.newConcurrentHashMap();
      }
      TrieMap<K, V> parent = curTrie;
      curTrie = curTrie.children.get(element);
      if(curTrie == null){
        parent.children.put(element, curTrie = new TrieMap<>());
      }
    }
    V oldValue = curTrie.value;
    curTrie.value = value;
    return oldValue;
  }

  public V put(K[] key, V value) {
    return put(Arrays.asList(key), value);
  }

  @Override
  public V remove(Object key) {
    if (key instanceof Iterable) {
      return remove( (Iterable) key );
    }
    return null;
  }

  public V remove(Iterable key) {
    TrieMap<K, V> parent = null;
    TrieMap<K, V> curTrie = this;
    Object lastKey = null;
    // go through each element
    for(Object element : key){
      if (curTrie.children == null) return null;
      lastKey = element;
      parent = curTrie;
      curTrie = curTrie.children.get(element);
      if (curTrie == null){
        return null;
      }
    }
    V v = curTrie.value;
    if (parent != null) {
      parent.children.remove(lastKey);
    } else {
      value = null;
    }
    return v;
  }

  public V remove(K[] key) {
    return remove(Arrays.asList(key));
  }

  @Override
  public void putAll(Map<? extends Iterable<K>, ? extends V> m) {
    for (Entry<? extends Iterable<K>, ? extends V> entry : m.entrySet()) {
      put(entry.getKey(), entry.getValue());
    }
  }

  @Override
  public void clear() {
    value = null;
    children = null;
  }

  @Override
  public Set<Iterable<K>> keySet() {
    Set<Iterable<K>> keys = new LinkedHashSet<>();
    updateKeys(keys, new ArrayList<>());
    return keys;
  }

  protected void updateKeys(Set<Iterable<K>> keys, List<K> prefix) {
    if (children != null) {
      for (Entry<K, TrieMap<K, V>> kTrieMapEntry : children.entrySet()) {
        List<K> p = new ArrayList<>(prefix.size() + 1);
        p.addAll(prefix);
        p.add(kTrieMapEntry.getKey());
        kTrieMapEntry.getValue().updateKeys(keys, p);
      }
    }
    if (value != null) {
      keys.add(prefix);
    }
  }

  @Override
  public Collection<V> values() {
    List<V> values = new ArrayList<>();
    updateValues(values);
    return values;
  }

  protected void updateValues(List<V> values) {
    if (children != null) {
      for (Entry<K, TrieMap<K, V>> kTrieMapEntry : children.entrySet()) {
        kTrieMapEntry.getValue().updateValues(values);
      }
    }
    if (value != null) {
      values.add(value);
    }
  }

  @Override
  public Set<Entry<Iterable<K>, V>> entrySet() {
    Set<Entry<Iterable<K>, V>> entries = new LinkedHashSet<>();
    updateEntries(entries, new ArrayList<>());
    return entries;
  }

  protected void updateEntries(Set<Entry<Iterable<K>,V>> entries, final List<K> prefix) {
    if (children != null) {
      for (Entry<K, TrieMap<K, V>> kTrieMapEntry : children.entrySet()) {
        List<K> p = new ArrayList<>(prefix.size() + 1);
        p.addAll(prefix);
        p.add(kTrieMapEntry.getKey());
        kTrieMapEntry.getValue().updateEntries(entries, p);
      }
    }
    if (value != null) {
      entries.add(new Map.Entry<Iterable<K>,V>() {
        @Override
        public Iterable<K> getKey() {
          return prefix;
        }

        @Override
        public V getValue() {
          return value;
        }

        @Override
        public V setValue(V value) {
          throw new UnsupportedOperationException();
        }
      });
    }
  }

}