using System;
using System.Collections.Generic;
using System.Linq;

namespace miew.Multimap
{
	/// <summary>
	/// Simple non-unique map wrapper
	/// </summary>
	/// <remarks>
	/// ApplyResultSelector (from Lookup[TKey, TElement] is not implemented, since the caller could just 
	/// as easily (or more-so) use .Select() with a Func[IGrouping[TKey, TElement], TResult], since
	/// IGrouping[TKey, TElement] already includes both the "TKey Key" and the IEnumerable[TElement].
	/// </remarks>
	public sealed class Multimap<TKey, TElement> : ILookup<TKey, TElement>
	{
		internal sealed class Grouping : List<TElement>, IGrouping<TKey, TElement>
		{
#if DEBUG
			readonly
#endif
			private TKey key;
			public TKey Key { get { return key; } }
			public Grouping(TKey key)
			{
				this.key = key;
			}
			public Grouping(TKey key, IEnumerable<TElement> ie)
				: base(ie)
			{
				this.key = key;
			}
		}

		private readonly Dictionary<TKey, Grouping> groups;
		/// <summary>
		/// Creates a new EditableLookup using the default key-comparer
		/// </summary>
		public Multimap() : this(default(IEqualityComparer<TKey>)) { }
		/// <summary>
		/// Creates a new EditableLookup using the specified key-comparer
		/// </summary>
		/// <param name="keyComparer"></param>
		public Multimap(IEqualityComparer<TKey> keyComparer)
		{
			groups = new Dictionary<TKey, Grouping>(keyComparer ?? EqualityComparer<TKey>.Default);
		}
		/// <summary>
		/// 
		/// </summary>
		/// <param name="keyComparer"></param>
		public Multimap(ILookup<TKey, TElement> il)
		{
			groups = new Dictionary<TKey, Grouping>();
			foreach (IGrouping<TKey, TElement> q in il)
				groups.Add(q.Key, new Grouping(q.Key, q));
		}

		public IEnumerable<KeyValuePair<TKey, TElement>> RecoverKeys(IEnumerable<TElement> values)
		{
			Dictionary<TElement, TKey> rev = groups
				.SelectMany(e => e.Value.Select(f => new { obj = f, k = e.Key }))
				.ToDictionary(e => e.obj, e => e.k);

			foreach (TElement e in values)
			{
				TKey recover_key;
				if (rev.TryGetValue(e, out recover_key))
					yield return new KeyValuePair<TKey, TElement>(recover_key, e);
			}
		}

		/// <summary>
		/// 
		/// </summary>
		public Multimap(IEnumerable<TElement> ie, Func<TElement, TKey> keySelector)
		{
			groups = new Dictionary<TKey, Grouping>();
			foreach (TElement e in ie)
				this.Add(keySelector(e), e);
		}

		public IEnumerable<TKey> Keys { get { return groups.Keys; } }

		/// <summary>
		/// Does the lookup contain any value(s) for the given key?
		/// </summary>
		public bool Contains(TKey key)
		{
			Grouping group;
			return groups.TryGetValue(key, out group) ? group.Count > 0 : false;
		}
		/// <summary>
		/// Does the lookup the specific key/value pair?
		/// </summary>
		public bool Contains(TKey key, TElement value)
		{
			Grouping group;
			return groups.TryGetValue(key, out group) ? group.Contains(value) : false;
		}
		/// <summary>
		/// Adds a key/value pair to the lookup
		/// </summary>
		/// <remarks>If the value is already present it will be duplicated</remarks>
		public void Add(TKey key, TElement value)
		{
			Grouping group;
			if (!groups.TryGetValue(key, out group))
			{
				group = new Grouping(key);
				groups.Add(key, group);
			}
			group.Add(value);
		}
		/// <summary>
		/// Adds a range of values against a single key
		/// </summary>
		/// <remarks>Any values already present will be duplicated</remarks>
		public void AddRange(TKey key, IEnumerable<TElement> values)
		{
			Grouping group;
			if (!groups.TryGetValue(key, out group))
			{
				group = new Grouping(key);
				groups.Add(key, group);
			}
			foreach (TElement value in values)
			{
				group.Add(value);
			}
			if (group.Count == 0)
			{
				groups.Remove(key); // nothing there after all!
			}
		}
		/// <summary>
		/// Add all key/value pairs from the supplied lookup to the current lookup
		/// </summary>
		/// <remarks>Any values already present will be duplicated</remarks>
		public void AddRange(ILookup<TKey, TElement> lookup)
		{
			foreach (IGrouping<TKey, TElement> group in lookup)
			{
				AddRange(group.Key, group);
			}
		}
		/// <summary>
		/// Remove all values from the lookup for the given key
		/// </summary>
		/// <returns>True if any items were removed, else false</returns>
		public bool Remove(TKey key)
		{
			return groups.Remove(key);
		}
		/// <summary>
		/// Remove the specific key/value pair from the lookup
		/// </summary>
		/// <returns>True if the item was found, else false</returns>
		public bool Remove(TKey key, TElement value)
		{
			Grouping group;
			if (groups.TryGetValue(key, out group))
			{
				bool removed = group.Remove(value);
				if (removed && group.Count == 0)
				{
					groups.Remove(key);
				}
				return removed;
			}
			return false;
		}
		/// <summary>
		/// Trims the inner data-structure to remove any surplus space
		/// </summary>
		public void TrimExcess()
		{
			foreach (var group in groups.Values)
			{
				group.TrimExcess();
			}
		}
		/// <summary>
		/// Returns the number of dictinct keys in the lookup
		/// </summary>
		public int Count
		{
			get { return groups.Count; }
		}

		private static readonly IEnumerable<TElement> arr_telement_empty = new TElement[0];
		/// <summary>
		/// Returns the set of values for the given key
		/// </summary>
		public IEnumerable<TElement> this[TKey key]
		{
			get
			{
				Grouping group;
				if (groups.TryGetValue(key, out group))
					return group;
				return arr_telement_empty;
			}
		}
		/// <summary>
		/// Returns the sequence of keys and their contained values
		/// </summary>
		public IEnumerator<IGrouping<TKey, TElement>> GetEnumerator()
		{
			foreach (var group in groups.Values)
			{
				yield return group;
			}
		}
		System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
		{
			return GetEnumerator();
		}
	};
}