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

namespace Implab.Automaton.RegularExpressions {
    public class TaggedDFA<TInput, TTag> : DFATable, ITaggedDFABuilder<TTag> {

        readonly Dictionary<int,TTag[]> m_tags = new Dictionary<int, TTag[]>();
        readonly IAlphabet<TInput> m_alphabet;

        public TaggedDFA(IAlphabet<TInput> alphabet) {
            Safe.ArgumentNotNull(alphabet, "aplhabet");

            m_alphabet = alphabet;
        }


        public IAlphabet<TInput> InputAlphabet {
            get {
                return m_alphabet;
            }
        }

        public void MarkFinalState(int s, TTag[] tags) {
            MarkFinalState(s);
            SetStateTag(s, tags);
        }

        public void SetStateTag(int s, TTag[] tags) {
            Safe.ArgumentNotNull(tags, "tags");
            m_tags[s] = tags;
        }

        public TTag[] GetStateTag(int s) {
            TTag[] tags;
            return m_tags.TryGetValue(s, out tags) ? tags : new TTag[0];
        }

        public TTag[][] CreateTagTable() {
            var table = new TTag[StateCount][];

            foreach (var pair in m_tags)
                table[pair.Key] = pair.Value;

            return table;
        }

        /// <summary>
        /// Optimize the specified alphabet.
        /// </summary>
        /// <param name="alphabet">Пустой алфавит, который будет зполнен в процессе оптимизации.</param>
        public TaggedDFA<TInput,TTag> Optimize(IAlphabetBuilder<TInput> alphabet) {
            Safe.ArgumentNotNull(alphabet, "alphabet");

            var dfa = new TaggedDFA<TInput, TTag>(alphabet);

            var states = new DummyAlphabet(StateCount);
            var alphaMap = new Dictionary<int,int>();
            var stateMap = new Dictionary<int,int>();

            Optimize(dfa, alphaMap, stateMap); 

            // mark tags in the new DFA
            foreach (var g in m_tags.Where(x => x.Key < StateCount).GroupBy(x => stateMap[x.Key], x => x.Value ))
                dfa.SetStateTag(g.Key, g.SelectMany(x => x).ToArray());

            // make the alphabet for the new DFA
            foreach (var pair in alphaMap)
                alphabet.DefineClass(m_alphabet.GetSymbols(pair.Key), pair.Value);
            
            return dfa;
        }

        protected override IEnumerable<HashSet<int>> GroupFinalStates() {
            var arrayComparer = new CustomEqualityComparer<TTag[]>(
                (x,y) => x.Length == y.Length && x.All(it => y.Contains(it)),
                x => x.Sum(it => x.GetHashCode())
            );
            return FinalStates.GroupBy(x => m_tags[x], arrayComparer).Select(g => new HashSet<int>(g));
        }

    }
}