package de.dhbwstuttgart.typeinference.unify.model;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

import de.dhbwstuttgart.typeinference.unify.MartelliMontanariUnify;
import de.dhbwstuttgart.typeinference.unify.interfaces.IFiniteClosure;
import de.dhbwstuttgart.typeinference.unify.interfaces.IUnify;

/**
 * The finite closure for the type unification
 * @author Florian Steurer
 */
public class FiniteClosure implements IFiniteClosure {
	
	/**
	 * A map that maps every type to the node in the inheritance graph that contains that type.
	 */
	private HashMap<UnifyType, Node<UnifyType>> inheritanceGraph;
	
	/**
	 * A map that maps every typename to the nodes of the inheritance graph that contain a type with that name.
	 */
	private HashMap<String, Set<Node<UnifyType>>> strInheritanceGraph;
	
	/**
	 * The initial pairs of that define the inheritance tree
	 */
	private Set<UnifyPair> pairs;
	
	/**
	 * Creates a new instance using the inheritance tree defined in the pairs.
	 */
	public FiniteClosure(Set<UnifyPair> pairs) {
		this.pairs = new HashSet<>(pairs);
		inheritanceGraph = new HashMap<UnifyType, Node<UnifyType>>();
		
		// Build the transitive closure of the inheritance tree
		for(UnifyPair pair : pairs) {
			if(pair.getPairOp() != PairOperator.SMALLER)
				continue;
			
			// Add nodes if not already in the graph
			if(!inheritanceGraph.containsKey(pair.getLhsType())) 
				inheritanceGraph.put(pair.getLhsType(), new Node<UnifyType>(pair.getLhsType()));
			if(!inheritanceGraph.containsKey(pair.getRhsType()))
				inheritanceGraph.put(pair.getRhsType(), new Node<UnifyType>(pair.getRhsType()));
			
			Node<UnifyType> childNode = inheritanceGraph.get(pair.getLhsType());
			Node<UnifyType> parentNode = inheritanceGraph.get(pair.getRhsType());
			
			// Add edge
			parentNode.addDescendant(childNode);
			
			// Add edges to build the transitive closure
			parentNode.getPredecessors().stream().forEach(x -> x.addDescendant(childNode));
			childNode.getDescendants().stream().forEach(x -> x.addPredecessor(parentNode));
		}
		
		// Build the alternative representation with strings as keys
		strInheritanceGraph = new HashMap<>();
		for(UnifyType key : inheritanceGraph.keySet()) {
			if(!strInheritanceGraph.containsKey(key.getName()))
				strInheritanceGraph.put(key.getName(), new HashSet<>());
			
			strInheritanceGraph.get(key.getName()).add(inheritanceGraph.get(key));
		}
	}

	/**
	 * Returns all types of the finite closure that are subtypes of the argument.
	 * @return The set of subtypes of the argument.
	 */
	@Override
	public Set<UnifyType> smaller(UnifyType type) {
		if(type instanceof FunNType)
			return computeSmallerFunN((FunNType) type);
		
		Set<UnifyType> ts = new HashSet<>();
		ts.add(type);
		return computeSmaller(ts);
	}
	
	/**
	 * Computes the smaller functions for every type except FunNTypes.
	 */
	private Set<UnifyType> computeSmaller(Set<UnifyType> types) {
		HashSet<UnifyType> result = new HashSet<>();
		
		IUnify unify = new MartelliMontanariUnify();
		
		for(UnifyType t : types) {
			
			// if T = T' then T <* T'
			result.add(t);
			
			// if C<...> <* C<...> then ... (third case in definition of <*)
			if(t.getTypeParams().size() > 0) {
				ArrayList<Set<UnifyType>> paramCandidates = new ArrayList<>();
				for (int i = 0; i < t.getTypeParams().size(); i++)
					paramCandidates.add(smArg(t.getTypeParams().get(i)));
				permuteParams(paramCandidates).forEach(x -> result.add(t.setTypeParams(x)));
			}
			
			if(!strInheritanceGraph.containsKey(t.getName()))
				continue;
			
			// if T <* T' then sigma(T) <* sigma(T')
			Set<Node<UnifyType>> candidates = strInheritanceGraph.get(t.getName());
			for(Node<UnifyType> candidate :  candidates) {
				UnifyType theta2 = candidate.getContent();
				Optional<Unifier> optSigma = unify.unify(theta2, t);
				if(!optSigma.isPresent())
					continue;
				
				Unifier sigma = optSigma.get();
				sigma.swapPlaceholderSubstitutions(t.getTypeParams());
				
				Set<UnifyType> theta1Set = candidate.getContentOfDescendants();
				
				for(UnifyType theta1 : theta1Set)
					result.add(theta1.apply(sigma));
			}
		}
		
		if(result.equals(types))
			return result;
		return computeSmaller(result);
	}
	
	/**
	 * Computes the smaller-Function for FunNTypes.
	 */
	private Set<UnifyType> computeSmallerFunN(FunNType type) {
		Set<UnifyType> result = new HashSet<>();
		
		// if T = T' then T <=* T'
	    result.add(type); 
	    
	    // Because real function types are implicitly variant
	    // it is enough to permute the params with the values of greater / smaller.
		ArrayList<Set<UnifyType>> paramCandidates = new ArrayList<>();
		paramCandidates.add(smaller(type.getTypeParams().get(0)));
		for (int i = 1; i < type.getTypeParams().size(); i++)
			paramCandidates.add(greater(type.getTypeParams().get(i)));
	
		permuteParams(paramCandidates).forEach(x -> result.add(type.setTypeParams(x)));
		return result;
	}

	/**
	 * Returns all types of the finite closure that are supertypes of the argument.
	 * @return The set of supertypes of the argument.
	 */
	@Override
	public Set<UnifyType> greater(UnifyType type) {
		if(type instanceof FunNType)
			return computeGreaterFunN((FunNType) type); 
		
		Set<UnifyType> ts = new HashSet<>();
		ts.add(type);
		return computeGreater(ts);
	}
 
	/**
	 * Computes the greater function for all types except function types.
	 */
	protected Set<UnifyType> computeGreater(Set<UnifyType> types) {
		HashSet<UnifyType> result = new HashSet<>();
		
		IUnify unify = new MartelliMontanariUnify();
		
		for(UnifyType t : types) {
			
			// if T = T' then T <=* T'
			result.add(t);
			
			// if C<...> <* C<...> then ... (third case in definition of <*)			
			if(t.getTypeParams().size() > 0) {
				ArrayList<Set<UnifyType>> paramCandidates = new ArrayList<>();
				for (int i = 0; i < t.getTypeParams().size(); i++)
					paramCandidates.add(grArg(t.getTypeParams().get(i)));
				permuteParams(paramCandidates).forEach(x -> result.add(t.setTypeParams(x)));
			}
			
			if(!strInheritanceGraph.containsKey(t.getName()))
				continue;
			
			// if T <* T' then sigma(T) <* sigma(T')
			Set<Node<UnifyType>> candidates = strInheritanceGraph.get(t.getName());
			for(Node<UnifyType> candidate :  candidates) {
				UnifyType theta1 = candidate.getContent();
				Optional<Unifier> optSigma = unify.unify(theta1, t);
				if(!optSigma.isPresent())
					continue;
				
				Unifier sigma = optSigma.get();
				sigma.swapPlaceholderSubstitutionsReverse(theta1.getTypeParams());
				
				Set<UnifyType> theta2Set = candidate.getContentOfPredecessors();
				
				for(UnifyType theta2 : theta2Set)
					result.add(theta2.apply(sigma));
			}
			
		}
		
		if(result.equals(types))
			return result;
		return computeGreater(result);
	}

	/**
	 * Computes the greater function for FunN-Types
	 */
	protected Set<UnifyType> computeGreaterFunN(FunNType type) {
		Set<UnifyType> result = new HashSet<>();
		
		// if T = T' then T <=* T'
	    result.add(type); 
	    
	    // Because real function types are implicitly variant
	    // it is enough to permute the params with the values of greater / smaller.
		ArrayList<Set<UnifyType>> paramCandidates = new ArrayList<>();
		paramCandidates.add(greater(type.getTypeParams().get(0)));
		for (int i = 1; i < type.getTypeParams().size(); i++)
				paramCandidates.add(smaller(type.getTypeParams().get(i)));
		permuteParams(paramCandidates).forEach(x -> result.add(type.setTypeParams(x)));
		return result;
	}

	
	@Override	
	public Set<UnifyType> grArg(UnifyType type) {
		return type.grArg(this);
	}

	@Override
	public Set<UnifyType> grArg(ReferenceType type) {
		Set<UnifyType> result = new HashSet<UnifyType>();
		result.add(type);	
        smaller(type).forEach(x -> result.add(new SuperType(x)));
        greater(type).forEach(x -> result.add(new ExtendsType(x)));
		return result;
	}
	
	@Override
	public Set<UnifyType> grArg(FunNType type) {
		Set<UnifyType> result = new HashSet<UnifyType>();
		result.add(type);	
        smaller(type).forEach(x -> result.add(new SuperType(x)));
        greater(type).forEach(x -> result.add(new ExtendsType(x)));
		return result;
	}
	
	@Override
	public Set<UnifyType> grArg(ExtendsType type) {
		Set<UnifyType> result = new HashSet<UnifyType>();
		result.add(type);
		UnifyType t = type.getExtendedType();
		greater(t).forEach(x -> result.add(new ExtendsType(x)));
		return result;
	}

	@Override
	public Set<UnifyType> grArg(SuperType type) {
		Set<UnifyType> result = new HashSet<UnifyType>();
		result.add(type);
		UnifyType t = type.getSuperedType();
		smaller(t).forEach(x -> result.add(new SuperType(x)));
		return result;
	}
	
	@Override
	public Set<UnifyType> grArg(PlaceholderType type) {
		HashSet<UnifyType> result = new HashSet<>();
		result.add(type);
		return result;
	}

	@Override	
	public Set<UnifyType> smArg(UnifyType type) {
		return type.smArg(this);
	}
	
	@Override
	public Set<UnifyType> smArg(ReferenceType type) {
		Set<UnifyType> result = new HashSet<UnifyType>();
		result.add(type);	
		return result;
	}
	
	@Override
	public Set<UnifyType> smArg(FunNType type) {
		Set<UnifyType> result = new HashSet<UnifyType>();
		result.add(type);	
		return result;
	}
	
	@Override
	public Set<UnifyType> smArg(ExtendsType type) {
		Set<UnifyType> result = new HashSet<UnifyType>();
		result.add(type);
		UnifyType t = type.getExtendedType();
		result.add(t);	
        smaller(t).forEach(x -> { 
        	result.add(new ExtendsType(x));
        	result.add(x); 
        });
		return result;
	}
	

	@Override
	public Set<UnifyType> smArg(SuperType type) {
		Set<UnifyType> result = new HashSet<UnifyType>();
		result.add(type);
		UnifyType t = type.getSuperedType();
		result.add(t);	
        greater(t).forEach(x -> { 
        	result.add(new SuperType(x));
        	result.add(x); 
        });
		return result;
	}
	
	@Override
	public Set<UnifyType> smArg(PlaceholderType type) {
		HashSet<UnifyType> result = new HashSet<>();
		result.add(type);
		return result;
	}

	@Override
	public Set<UnifyType> getAllTypesByName(String typeName) {
		if(!strInheritanceGraph.containsKey(typeName))
			return new HashSet<>();
		return strInheritanceGraph.get(typeName).stream().map(x -> x.getContent()).collect(Collectors.toCollection(HashSet::new));
	}
	
	@Override
	public Optional<UnifyType> getLeftHandedType(String typeName) {
		if(!strInheritanceGraph.containsKey(typeName))
			return Optional.empty();
		
		for(UnifyPair pair : pairs)
			if(pair.getLhsType().getName().equals(typeName))
					return Optional.of(pair.getLhsType());
		
		return Optional.empty();
	}
	
	@Override
	public Set<UnifyType> getAncestors(UnifyType t) {
		if(!inheritanceGraph.containsKey(t))
			return new HashSet<>();
		Set<UnifyType> result = inheritanceGraph.get(t).getContentOfPredecessors();
		result.add(t);
		return result;
	}
	
	@Override
	public Set<UnifyType> getChildren(UnifyType t) {
		if(!inheritanceGraph.containsKey(t))
			return new HashSet<>();
		Set<UnifyType> result = inheritanceGraph.get(t).getContentOfDescendants();
		result.add(t);
		return result;
	}
	
	/**
	 * Takes a set of candidates for each position and computes all possible permutations.
	 * @param candidates The length of the list determines the number of type params. Each set
	 * contains the candidates for the corresponding position.
	 */
	protected Set<TypeParams> permuteParams(ArrayList<Set<UnifyType>> candidates) {
		Set<TypeParams> result = new HashSet<>();
		permuteParams(candidates, 0, result, new UnifyType[candidates.size()]);
		return result;
	}
	
	/**
	 * Takes a set of candidates for each position and computes all possible permutations.
	 * @param candidates The length of the list determines the number of type params. Each set
	 * contains the candidates for the corresponding position.
	 * @param idx Idx for the current permutatiton.
	 * @param result Set of all permutations found so far
	 * @param current The permutation of type params that is currently explored
	 */
	protected void permuteParams(ArrayList<Set<UnifyType>> candidates, int idx, Set<TypeParams> result, UnifyType[] current) {
		if(candidates.size() == idx) {
			result.add(new TypeParams(Arrays.copyOf(current, current.length)));
			return;
		}
		
		Set<UnifyType> localCandidates = candidates.get(idx);
		
		for(UnifyType t : localCandidates) {
			current[idx] = t;
			permuteParams(candidates, idx+1, result, current);
		}
	}
	
	@Override
	public String toString(){
		return this.inheritanceGraph.toString();
	}
}
