package org.simantics.scl.compiler.elaboration.expressions;

import org.simantics.scl.compiler.elaboration.contexts.ReplaceContext;
import org.simantics.scl.compiler.elaboration.contexts.SimplificationContext;
import org.simantics.scl.compiler.elaboration.contexts.TranslationContext;
import org.simantics.scl.compiler.elaboration.contexts.TypingContext;
import org.simantics.scl.compiler.errors.Locations;
import org.simantics.scl.compiler.internal.elaboration.utils.ExpressionDecorator;
import org.simantics.scl.compiler.types.Type;
import org.simantics.scl.compiler.types.Types;
import org.simantics.scl.compiler.types.exceptions.MatchException;
import org.simantics.scl.compiler.types.exceptions.UnificationException;
import org.simantics.scl.compiler.types.kinds.Kinds;
import org.simantics.scl.compiler.types.util.MultiFunction;

import gnu.trove.map.hash.TObjectIntHashMap;
import gnu.trove.set.hash.THashSet;
import gnu.trove.set.hash.TIntHashSet;

public class ELambda extends SimplifiableExpression {
    public Case[] cases;
    Type effect = Types.NO_EFFECTS;
    
    public ELambda(Case[] cases) {
        this.cases = cases;
    }
    
    public ELambda(Case case_) {
        this(new Case[] {case_});
    }

    public ELambda(long loc, Case ... cases) {
        super(loc);
        this.cases = cases;
    }
    
    public ELambda(long loc, Expression pat, Expression exp) {
        this(loc, new Case(new Expression[] {pat}, exp));
    }

	public void collectRefs(TObjectIntHashMap<Object> allRefs, TIntHashSet refs) {
        for(Case case_ : cases)
            case_.collectRefs(allRefs, refs);
    }
	
	@Override
	public void collectVars(TObjectIntHashMap<Variable> allVars,
	        TIntHashSet vars) {
	    for(Case case_ : cases)
            case_.collectVars(allVars, vars);
	}

	public Expression decomposeMatching() {
	    Expression[] patterns = cases[0].patterns;
        int arity = patterns.length;
        
        // Simple cases
        if(cases.length == 1 && 
                !(cases[0].value instanceof GuardedExpressionGroup)) {
            boolean noMatchingNeeded = true;
            for(int i=0;i<arity;++i)
                if(!(patterns[i] instanceof EVariable)) {
                    noMatchingNeeded = false;
                    break;
                }
            if(noMatchingNeeded) {
                Expression decomposed = cases[0].value.decomposeMatching();
                Type effect = this.effect;
                for(int i=arity-1;i>=0;--i) {
                    Variable var = ((EVariable)patterns[i]).getVariable(); 
                    decomposed = new ESimpleLambda(getLocation(), var, effect, decomposed);
                    effect = Types.NO_EFFECTS;
                }
                return decomposed;
            }
        }
        
        // Complex case
        Variable[] vars = new Variable[arity];
        Expression[] scrutinee = new Expression[arity];
        for(int i=0;i<arity;++i) {
            vars[i] = new Variable("temp" + i);
            Type type = patterns[i].getType();
            vars[i].setType(type);
            scrutinee[i] = new EVariable(getLocation(), vars[i]);
            scrutinee[i].setType(type);
        }
        Expression decomposed = new EMatch(getLocation(), scrutinee, cases);
        Type curEffect = this.effect;
        for(int i=arity-1;i>=0;--i) {            
            decomposed = new ESimpleLambda(getLocation(), vars[i], curEffect, decomposed);            
            curEffect = Types.NO_EFFECTS;                 
        }
        return decomposed;
    }
	
	@Override
	protected void updateType() throws MatchException {
	    setType(Types.functionE(Types.getTypes(cases[0].patterns), effect, cases[0].value.getType()));
	}
	
	@Override
	public Expression simplify(SimplificationContext context) {
	    return decomposeMatching().simplify(context);
	}

    @Override
    public void collectFreeVariables(THashSet<Variable> vars) {
        for(Case case_ : cases)
            case_.collectFreeVariables(vars);
    }

    @Override
    public Expression resolve(TranslationContext context) {
        for(Case case_ : cases)
            case_.resolve(context);
        return this;
    }
    
    @Override
    public void setLocationDeep(long loc) {
        if(location == Locations.NO_LOCATION) {
            location = loc;
            for(Case case_ : cases)
                case_.setLocationDeep(loc);
        }
    }
    
    @Override
    public Expression replace(ReplaceContext context) {
        Case[] newCases = new Case[cases.length];
        for(int i=0;i<cases.length;++i)
            newCases[i] = cases[i].replace(context);
        return new ELambda(newCases);
    }

    public Case[] getCases() {
        return cases;
    }
    
    public Expression checkBasicType(TypingContext context, Type requiredType) {
        int arity = cases[0].patterns.length;
        MultiFunction mfun;
        try {
            mfun = Types.unifyFunction(requiredType, arity);
        } catch (UnificationException e) {
            int requiredArity = Types.getArity(requiredType);
            context.getErrorLog().log(cases[0].getLhs(), "Arity is " + requiredArity + " but "
                    + arity + " patterns have been given.");
            setType(Types.metaVar(Kinds.STAR));
            return this;
        }
        
        effect = mfun.effect;
        context.pushEffectUpperBound(location, mfun.effect);
        for(Case case_ : cases)
            case_.checkType(context, mfun.parameterTypes,  mfun.returnType);
        context.popEffectUpperBound();
        return this;
    }
    
    @Override
    public Expression inferType(TypingContext context) {
        int arity = cases[0].patterns.length;
        effect = Types.metaVar(Kinds.EFFECT);
        context.pushEffectUpperBound(location, effect);
        Type[] parameterTypes = new Type[arity];        
        for(int i=0;i<parameterTypes.length;++i)
            parameterTypes[i] = Types.metaVar(Kinds.STAR);
        Type requiredType = Types.metaVar(Kinds.STAR);
        for(Case case_ : cases)
            case_.checkType(context, parameterTypes, requiredType);
        context.popEffectUpperBound();
        return this;
    }

    @Override
    public Expression decorate(ExpressionDecorator decorator) {
        if(decorator.decorateSubstructure(this))
            for(Case case_ : cases)
                case_.decorate(decorator);
        return decorator.decorate(this);
    }
    
    @Override
    public boolean isEffectful() {
    	return false;
    }

    @Override
    public void collectEffects(THashSet<Type> effects) {
        for(Case case_ : cases) {
            for(Expression pattern : case_.patterns)
                pattern.collectEffects(effects);
            case_.value.collectEffects(effects);
        }
    }
    
    @Override
    public void accept(ExpressionVisitor visitor) {
        visitor.visit(this);
    }

    @Override
    public void forVariables(VariableProcedure procedure) {
        for(Case case_ : cases)
            case_.forVariables(procedure);
    }
    
    @Override
    public Expression accept(ExpressionTransformer transformer) {
        return transformer.transform(this);
    }

}