package org.simantics.scl.compiler.internal.codegen.ssa.statements;

import java.util.ArrayList;

import org.simantics.scl.compiler.common.exceptions.InternalCompilerError;
import org.simantics.scl.compiler.constants.SCLConstant;
import org.simantics.scl.compiler.internal.codegen.references.BoundVar;
import org.simantics.scl.compiler.internal.codegen.references.ValRef;
import org.simantics.scl.compiler.internal.codegen.ssa.SSAFunction;
import org.simantics.scl.compiler.internal.codegen.ssa.SSAStatement;
import org.simantics.scl.compiler.internal.codegen.ssa.binders.FunctionBinder;
import org.simantics.scl.compiler.internal.codegen.utils.CopyContext;
import org.simantics.scl.compiler.internal.codegen.utils.MethodBuilder;
import org.simantics.scl.compiler.internal.codegen.utils.PrintingContext;
import org.simantics.scl.compiler.internal.codegen.utils.SSALambdaLiftingContext;
import org.simantics.scl.compiler.internal.codegen.utils.SSASimplificationContext;
import org.simantics.scl.compiler.internal.codegen.utils.SSAValidationContext;
import org.simantics.scl.compiler.types.TVar;
import org.simantics.scl.compiler.types.Type;

import gnu.trove.map.hash.THashMap;
import gnu.trove.set.hash.THashSet;

public class LetFunctions extends SSAStatement implements FunctionBinder {
    long recursiveGroupLocation;
    SSAFunction firstFunction;

    public LetFunctions() {
    }
    
    public LetFunctions(SSAFunction function) {
        firstFunction = function;
        function.setParent(this);
    }

    @Override
    public void toString(PrintingContext context) {
        for(SSAFunction function = firstFunction; function != null; function = function.getNext()) {
            context.indentation();
            context.append(function.getTarget());
            context.append("(" + function.getTarget().occurrenceCount() + ")");
            context.append(" :: ");
            context.append(function.getTarget().getType());
            context.append(" = \n");
            context.indent();
            function.toString(context);
            context.dedent();
        }
    }
    
    public void addFunction(SSAFunction function) {
        function.setParent(this);        
        function.setNext(firstFunction);
        if(firstFunction != null)
            firstFunction.setPrev(function);
        firstFunction = function;
    }

    @Override
    public void generateCode(MethodBuilder mb) {
        throw new InternalCompilerError("Functions should be lambda lifted before code generation");        
    }

    @Override
    public void validate(SSAValidationContext context) {
        for(SSAFunction function = firstFunction; function != null; function = function.getNext()) {
            if(!(function.getTarget() instanceof BoundVar))
                throw new InternalCompilerError();
            function.validate(context);
        }
    }

    @Override
    public void destroy() {
        for(SSAFunction function = firstFunction; function != null; function = function.getNext())
            function.destroy();
    }

    @Override
    public SSAStatement copy(CopyContext context) {
        LetFunctions result = new LetFunctions();
        for(SSAFunction function = firstFunction; function != null; function = function.getNext()) {
            SSAFunction newFunction = function.copy(context);
            newFunction.setTarget(context.copy(function.getTarget()));
            result.addFunction(newFunction);
        }
        return result;        
    }

    @Override
    public void replace(TVar[] vars, Type[] replacements) {
        for(SSAFunction function = firstFunction; function != null; function = function.getNext()) {
            ((BoundVar)function.getTarget()).replace(vars, replacements);
            function.replace(vars, replacements);
        }
    }

    @Override
    public void addBoundVariablesTo(SSAValidationContext context) {
        for(SSAFunction function = firstFunction; function != null; function = function.getNext())
            context.validBoundVariables.add((BoundVar)function.getTarget());        
    }

    @Override
    public SSAFunction getFirstFunction() {
        return firstFunction;
    }

    @Override
    public void setFirstFunction(SSAFunction function) {
        this.firstFunction = function;     
        if(function == null)
            detach();
    }

    @Override
    public void collectFreeVariables(SSAFunction parentFunction,
            ArrayList<ValRef> vars) {
        throw new InternalCompilerError("Should not be called for non-lambda-lifted functions.");
        // FIXME inefficient, some kind of caching needed here
        /*THashSet<BoundVar> tempVars = new THashSet<BoundVar>();
        for(SSAFunction function = firstFunction; function != null; function = function.getNext())
            function.collectFreeVariables(tempVars);
        
        for(BoundVar var : tempVars)
            if(var.getFunctionParent() != parentFunction)
                vars.add(var);*/
    }

    @Override
    public void lambdaLift(SSALambdaLiftingContext context) {
        boolean hasValues = false;
        boolean isRecursive = false;
        
        // Lambda lift functions and collect free variables
        THashSet<BoundVar> targets = new THashSet<BoundVar>();
        ArrayList<ValRef> freeVars = new ArrayList<ValRef>();        
        for(SSAFunction function = firstFunction; 
                function != null; 
                function = function.getNext()) {
            hasValues |= function.getArity() == 0;
            function.lambdaLift(context);
            targets.add((BoundVar)function.getTarget());
            function.collectFreeVariables(freeVars);
        }
                
        // Classify by BoundVars
        THashSet<BoundVar> boundVars = new THashSet<BoundVar>(); 
        ArrayList<BoundVar> boundVarsList = new ArrayList<BoundVar>(4);
        ArrayList<ValRef> newFreeVars = new ArrayList<ValRef>(freeVars.size()); 
        for(ValRef ref : freeVars) {
            BoundVar var = (BoundVar)ref.getBinding();
            if(targets.contains(var)) {
                isRecursive = true;
                continue;
            }
            if(boundVars.add(var))
                boundVarsList.add(var);
            newFreeVars.add(ref);
            /*BoundVar inVar = map.get(outVar);
            if(inVar == null) {
                inVar = new BoundVar(outVar.getType());
                map.put(outVar, inVar);
                outParameters.add(outVar);
                inParameters.add(inVar);
            }
            ref.replaceBy(inVar);*/
        }
        BoundVar[] outVars = boundVarsList.toArray(new BoundVar[boundVarsList.size()]);
        freeVars = newFreeVars;
        
        // Modify functions
        THashMap<SSAFunction, THashMap<BoundVar, BoundVar>> varMap = new THashMap<SSAFunction, THashMap<BoundVar, BoundVar>>();
        THashMap<SSAFunction, BoundVar[]> inVarsMap = new THashMap<SSAFunction, BoundVar[]>();
        THashMap<SSAFunction, BoundVar> oldTargets = new THashMap<SSAFunction, BoundVar>();
        for(SSAFunction function = firstFunction; 
                function != null; 
                function = function.getNext()) {
            THashMap<BoundVar, BoundVar> map = new THashMap<BoundVar, BoundVar>(2*outVars.length);
            BoundVar[] inVars = new BoundVar[outVars.length];            
            for(int i=0;i<inVars.length;++i) {
                inVars[i] = new BoundVar(outVars[i].getType());
                map.put(outVars[i], inVars[i]);
            }
            inVarsMap.put(function, inVars);
            varMap.put(function, map);
            
            function.addParametersInFront(inVars);            
            SCLConstant functionConstant = new SCLConstant(context.createName(), function.getType());
            context.addConstant(functionConstant);   
            oldTargets.put(function, (BoundVar)function.getTarget());
            function.setTarget(functionConstant);
            functionConstant.setDefinition(function);   
            functionConstant.setPrivate(true);
            // TODO handle type parameters
            
            // Define target by an application
            /*new LetApply(oldTarget, functionConstant.createOccurrence(), 
                    ValRef.createOccurrences(outVars)).insertBefore(this);*/
        }
        
        for(SSAFunction function = firstFunction; 
                function != null; 
                function = function.getNext()) {
            BoundVar oldTarget = oldTargets.get(function);
            for(ValRef ref : oldTarget.getOccurences()) {
                SSAFunction parent = ref.getParentFunction();
                BoundVar[] vars = inVarsMap.get(parent);
                if(vars == null)
                    vars = outVars;
                if(vars.length > 0)
                    ref.replaceByApply(function.getTarget(), vars);
                else
                    ref.replaceBy(function.getTarget());
            }
        }
            
        // Fix references
        for(ValRef ref : freeVars) {
            BoundVar inVar = (BoundVar)ref.getBinding();
            if(targets.contains(inVar))
                continue;
            BoundVar outVar = varMap.get(ref.getParentFunction()).get(inVar);
            ref.replaceBy(outVar);
        }
        
        detach();
        //context.validate();
        
        if(hasValues && isRecursive)
            context.getErrorLog().log(recursiveGroupLocation, "Variables defined recursively must all be functions.");
    }
    
    @Override
    public void simplify(SSASimplificationContext context) {
        for(SSAFunction function = firstFunction; 
                function != null; 
                function = function.getNext())
            function.simplify(context);
    }
    
    public void setRecursiveGroupLocation(long recursiveGroupLocation) {
        this.recursiveGroupLocation = recursiveGroupLocation;
    }
}