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

import java.util.ArrayList;

import org.simantics.scl.compiler.common.exceptions.InternalCompilerError;
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.environment.Environment;
import org.simantics.scl.compiler.errors.Locations;
import org.simantics.scl.compiler.internal.codegen.references.IVal;
import org.simantics.scl.compiler.internal.codegen.writer.CodeWriter;
import org.simantics.scl.compiler.internal.elaboration.utils.ExpressionDecorator;
import org.simantics.scl.compiler.internal.interpreted.IExpression;
import org.simantics.scl.compiler.internal.interpreted.ILambda;
import org.simantics.scl.compiler.top.ExpressionInterpretationContext;
import org.simantics.scl.compiler.top.SCLCompilerConfiguration;
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 ESimpleLambda extends Expression {
    public Variable parameter;
    public Expression value;
    public Type effect = Types.NO_EFFECTS;
    
    public ESimpleLambda(Variable parameter, Expression value) {
        this.parameter = parameter;
        this.value = value;
    }
    
    public ESimpleLambda(Type effect, Variable parameter, Expression value) {
        this.parameter = parameter;
        this.value = value;
        this.effect = effect;
    }

    public ESimpleLambda(long loc, Variable parameter, Expression value) {
        super(loc);
        this.parameter = parameter;
        this.value = value;
    }
    
    public ESimpleLambda(long loc, Variable parameter, Type effect, Expression value) {
        super(loc);
        this.parameter = parameter;
        this.value = value;
        this.effect = effect;
    }

	public void collectRefs(TObjectIntHashMap<Object> allRefs, TIntHashSet refs) {
        value.collectRefs(allRefs, refs);
    }
	
	@Override
	public void collectVars(TObjectIntHashMap<Variable> allVars,
	        TIntHashSet vars) {
	    value.collectVars(allVars, vars);
	}

	public Expression decomposeMatching() {
        value = value.decomposeMatching();
        return this;
    }

	@Override
	protected void updateType() throws MatchException {
	    setType(Types.functionE(Types.canonical(parameter.type),
	            effect, value.getType()));
	}
	
	@Override
	public IVal toVal(Environment env, CodeWriter w) {
	    return lambdaToVal(env, w);
    }

    @Override
    public void collectFreeVariables(THashSet<Variable> vars) {
        value.collectFreeVariables(vars);
        vars.remove(parameter);
    }

    @Override
    public Expression simplify(SimplificationContext context) {
        value = value.simplify(context);
        return this;
    }

    @Override
    public Expression resolve(TranslationContext context) {
        value = value.resolve(context);
        return this;
    }

    @Override
    public Expression replace(ReplaceContext context) {
        Variable newParameter = parameter.copy();
        context.varMap.put(parameter, new EVariable(newParameter));
        ESimpleLambda result = new ESimpleLambda(getLocation(),
                newParameter, 
                effect.replace(context.tvarMap),
                value.replace(context));
        // not absolutely needed, but maybe good for performance
        context.varMap.remove(parameter); 
        return result;
    }
        
    public Type getLocalEffect() {
        if(SCLCompilerConfiguration.DEBUG)
            if(effect == null)
                throw new InternalCompilerError();
        return effect;
    }
    
    public void setEffect(Type effect) {
        if(effect == null)
            throw new InternalCompilerError();
        this.effect = effect;
    }
    
    @Override
    public IExpression toIExpression(ExpressionInterpretationContext context) {
        // Find parameters of the whole function definition
        ArrayList<Variable> parameters = new ArrayList<Variable>(2);
        parameters.add(parameter);
        Expression cur = value;
        while(true) {
            if(cur instanceof ESimpleLambda) {
                ESimpleLambda lambda = (ESimpleLambda)cur;
                parameters.add(lambda.parameter);
                cur = lambda.value;
            }
            else if(cur instanceof ELambdaType) {
                cur = ((ELambdaType)cur).value;
            }
            else
                break;
            
        }
        
        // Free variables;
        ExpressionInterpretationContext innerContext =  context.createNewContext();
        THashSet<Variable> freeVariables = cur.getFreeVariables();
        for(Variable parameter : parameters)
            freeVariables.remove(parameter);
        int i=0;
        int[] inheritedVariableIds = new int[freeVariables.size()];
        for(Variable var : freeVariables) {
            innerContext.push(var);
            inheritedVariableIds[i++] = context.getVariableId(var);
        }
        
        // Parameters
        for(Variable parameter : parameters)
            innerContext.push(parameter);
        
        // Construct lambda
        IExpression body = cur.toIExpression(innerContext); 
        return new ILambda(inheritedVariableIds,
                parameters.size(),
                innerContext.getMaxVariableId(),
                body);
    }
    
    public Expression checkBasicType(TypingContext context, Type requiredType) {
        MultiFunction mfun;
        try {
            mfun = Types.unifyFunction(requiredType, 1);
        } catch (UnificationException e) {
            context.getErrorLog().log(location, "Required type is <" + requiredType + "> which is incompatible with lambda.");
            setType(Types.metaVar(Kinds.STAR));
            return this;
        }
        
        effect = mfun.effect;
        context.pushEffectUpperBound(location, mfun.effect);
        parameter.setType(mfun.parameterTypes[0]);
        value = value.checkType(context, mfun.returnType);
        context.popEffectUpperBound();
        return this;
    }
    
    @Override
    public Expression inferType(TypingContext context) {
        effect = Types.metaVar(Kinds.EFFECT);
        context.pushEffectUpperBound(location, effect);
        if(parameter.type == null)
            parameter.setType(Types.metaVar(Kinds.STAR));
        value = value.checkType(context, Types.metaVar(Kinds.STAR));
        context.popEffectUpperBound();
        return this;
    }

    @Override
    public Expression decorate(ExpressionDecorator decorator) {
        if(decorator.decorateSubstructure(this))
            value = value.decorate(decorator);
        return decorator.decorate(this);
    }

    @Override
    public boolean isEffectful() {
    	return false;
    }
    
    @Override
    public void collectEffects(THashSet<Type> effects) {
    }
    
    @Override
    public void setLocationDeep(long loc) {
        if(location == Locations.NO_LOCATION) {
            location = loc;
            value.setLocationDeep(loc);
        }
    }
    
    @Override
    public void accept(ExpressionVisitor visitor) {
        visitor.visit(this);
    }
    
    public Expression getValue() {
        return value;
    }
    
    public Variable getParameter() {
        return parameter;
    }

    @Override
    public void forVariables(VariableProcedure procedure) {
        value.forVariables(procedure);
    }
    
    @Override
    public Expression accept(ExpressionTransformer transformer) {
        return transformer.transform(this);
    }

}