package org.simantics.scl.compiler.parser.regexp.automata;

import gnu.trove.list.array.TIntArrayList;
import gnu.trove.map.hash.TIntObjectHashMap;
import gnu.trove.map.hash.TObjectIntHashMap;
import gnu.trove.procedure.TIntIntProcedure;
import gnu.trove.procedure.TIntObjectProcedure;
import gnu.trove.procedure.TIntProcedure;
import gnu.trove.set.hash.TIntHashSet;

import java.util.ArrayList;

public class NFA implements Automata {
    
    private static class State {
        TIntObjectHashMap<TIntArrayList> transitions = new TIntObjectHashMap<TIntArrayList>();
        TIntArrayList epsilonTransitions = new TIntArrayList();
        boolean accepts;
        
        private void addTransition(int symbol, int target) {
            TIntArrayList l = transitions.get(symbol);
            if(l == null) {
                l = new TIntArrayList();
                l.add(target);
                transitions.put(symbol, l);
            }
            else if(!l.contains(target))
                l.add(target);
        }
        
        private void addEpsilonTransition(int target) {
            if(!epsilonTransitions.contains(target))
                epsilonTransitions.add(target);
        }
    }
    
    private ArrayList<State> states = new ArrayList<State>();
    private int initialState;
    
    public int size() {
        return states.size();
    }
    
    public void setInitialState(int initialState) {
        this.initialState = initialState;
    }
    
    public int getInitialState() {
        return initialState;
    }
    
    public int newState() {
        int id = states.size();
        states.add(new State());
        return id;
    }
    
    public void setAccepts(int id, boolean accepts) {
        states.get(id).accepts = accepts;
    }
    
    public boolean getAccepts(int id) {
        return states.get(id).accepts;
    }
    
    public void addTransition(int source, int symbol, int target) {
        states.get(source).addTransition(symbol, target);
    }
    
    public void forEachTransition(int source, final TIntIntProcedure proc) {
        states.get(source).transitions.forEachEntry(new TIntObjectProcedure<TIntArrayList>() {
            @Override
            public boolean execute(int a, TIntArrayList b) {
                for(int i=0;i<b.size();++i)
                    if(!proc.execute(a, b.get(i)))
                        return false;
                return true;
            }
        });
    }
    
    public void addEpsilonTransition(int source, int target) {
        states.get(source).addEpsilonTransition(target);
    }
    
    public void forEachEpsilonTransition(int source, final TIntProcedure proc) {
        states.get(source).epsilonTransitions.forEach(proc);
    }
    
    private TIntArrayList closure(final TIntHashSet stateSet) {
        final TIntArrayList result = new TIntArrayList(stateSet);
        TIntProcedure proc = new TIntProcedure() {                
            @Override
            public boolean execute(int value) {
                if(stateSet.add(value))
                    result.add(value);
                return true;
            }
        };
        for(int i=0;i<result.size();++i)
            forEachEpsilonTransition(result.get(i), proc);                
        result.sort();
        return result;
    }
    
    public DFA determinize() {
        final DFA aut = new DFA();
        
        final ArrayList<TIntArrayList> stack = new ArrayList<TIntArrayList>(); 
        final TObjectIntHashMap<TIntArrayList> map = new TObjectIntHashMap<TIntArrayList>();
        {
            TIntHashSet initialSet = new TIntHashSet(4);
            initialSet.add(getInitialState());
            TIntArrayList stateList = closure(initialSet);
            
            int stateId = aut.newState();
            map.put(stateList, stateId);
            aut.setInitialState(stateId);
            
            stack.add(stateList);
        }
        
        while(!stack.isEmpty()) {
            TIntArrayList curList = stack.remove(stack.size()-1);
            int[] stateArray = curList.toArray();
            final int id = map.get(curList);
            
            // Transitions
            final TIntObjectHashMap<TIntHashSet> transitions =
                    new TIntObjectHashMap<TIntHashSet>();
            TIntIntProcedure proc = new TIntIntProcedure() {                
                @Override
                public boolean execute(int symbol, int b) {
                    TIntHashSet set = transitions.get(symbol);
                    if(set == null) {
                        set = new TIntHashSet();
                        transitions.put(symbol, set);
                    }
                    set.add(b);
                    return true;
                }
            };
            for(int s : stateArray)
                forEachTransition(s, proc);
            
            // Create transition targets
            transitions.forEachEntry(new TIntObjectProcedure<TIntHashSet>() {
                @Override
                public boolean execute(int symbol, TIntHashSet b) {
                    TIntArrayList stateList = closure(b);
                    
                    if(map.containsKey(stateList))
                        aut.addTransition(id, symbol, map.get(stateList));
                    else {
                        int stateId = aut.newState();
                        map.put(stateList, stateId);                        
                        stack.add(stateList);
                        
                        aut.addTransition(id, symbol, stateId);
                    }                    
                    return true;
                }
            });
            
            // Accepts
            for(int s : stateArray)
                if(getAccepts(s)) {
                    aut.setAccepts(id, true);
                    break;
                }
        }
        
        return aut;
    }
    
}