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

import java.util.ArrayList;
import java.util.Set;

import org.simantics.scl.compiler.elaboration.contexts.ReplaceContext;
import org.simantics.scl.compiler.elaboration.contexts.TypingContext;
import org.simantics.scl.compiler.elaboration.expressions.ESimpleLet;
import org.simantics.scl.compiler.elaboration.expressions.EVariable;
import org.simantics.scl.compiler.elaboration.expressions.Expression;
import org.simantics.scl.compiler.elaboration.expressions.QueryTransformer;
import org.simantics.scl.compiler.elaboration.expressions.Variable;
import org.simantics.scl.compiler.elaboration.expressions.VariableProcedure;
import org.simantics.scl.compiler.elaboration.java.EqRelation;
import org.simantics.scl.compiler.elaboration.query.compilation.ConstraintCollectionContext;
import org.simantics.scl.compiler.elaboration.query.compilation.DerivateException;
import org.simantics.scl.compiler.elaboration.query.compilation.EnforcingContext;
import org.simantics.scl.compiler.elaboration.query.compilation.ExpressionConstraint;
import org.simantics.scl.compiler.elaboration.query.compilation.RelationConstraint;
import org.simantics.scl.compiler.elaboration.relations.CompositeRelation;
import org.simantics.scl.compiler.elaboration.relations.LocalRelation;
import org.simantics.scl.compiler.elaboration.relations.SCLRelation;
import org.simantics.scl.compiler.errors.Locations;
import org.simantics.scl.compiler.types.TVar;
import org.simantics.scl.compiler.types.Type;
import org.simantics.scl.compiler.types.Types;

import gnu.trove.map.hash.THashMap;
import gnu.trove.map.hash.TIntObjectHashMap;
import gnu.trove.map.hash.TObjectIntHashMap;
import gnu.trove.set.hash.TIntHashSet;

public class QAtom extends Query {
    public SCLRelation relation;
    public Type[] typeParameters;
    public Expression[] parameters;

    public QAtom(SCLRelation relation, Expression ... parameters) {
        this.relation = relation;
        this.parameters = parameters;
    }

    public QAtom(SCLRelation relation, Type[] typeParameters, Expression ... parameters) {
        this.relation = relation;
        this.typeParameters = typeParameters;
        this.parameters = parameters;
    }

    @Override
    public void checkType(TypingContext context) {
        // Type parameters
        TVar[] typeVariables = relation.getTypeVariables();
        typeParameters = new Type[typeVariables.length];
        for(int i=0;i<typeVariables.length;++i)
            typeParameters[i] = Types.metaVar(typeVariables[i].getKind());

        // Check parameter types
        Type[] parameterTypes = relation.getParameterTypes();
        if(parameterTypes.length != parameters.length)
            context.getErrorLog().log(location, "Relation is applied with wrong number of parameters.");
        else
            for(int i=0;i<parameters.length;++i)
                parameters[i] = parameters[i]
                        .checkType(context, parameterTypes[i].replace(typeVariables, typeParameters));
    }

    public Expression generateEnforce(EnforcingContext context) {
        Variable[] variables = new Variable[parameters.length];
        for(int i=0;i<variables.length;++i)
            if(parameters[i] instanceof EVariable)
                variables[i] = ((EVariable)parameters[i]).getVariable();
            else
                variables[i] = new Variable("p" + i, parameters[i].getType());
        Expression result = relation.generateEnforce(location, context, typeParameters, variables);
        for(int i=variables.length-1;i>=0;--i)
            if(!(parameters[i] instanceof EVariable))
                result = new ESimpleLet(
                        variables[i],
                        parameters[i],
                        result
                        );
        return result;
    }

    private static class VariableMaskProcedure implements VariableProcedure {
        ConstraintCollectionContext context;
        long requiredVariablesMask = 0L;

        public VariableMaskProcedure(ConstraintCollectionContext context) {
            this.context = context;
        }

        @Override
        public void execute(long location, Variable variable) {
            int id = context.getVariableMap().get(variable);
            if(id >= 0)
                requiredVariablesMask |= 1L << id;
        }
    }

    @Override
    public void collectConstraints(ConstraintCollectionContext context) {
        try {
            // Analyze parameters and find required and optional variables
            VariableMaskProcedure procedure = new VariableMaskProcedure(context);
            int[] optionalVariableByParameter = new int[parameters.length];
            Variable[] varParameters = new Variable[parameters.length];
            for(int i=0;i<parameters.length;++i) {
                Expression parameter = parameters[i];
                if(parameter instanceof EVariable) {
                    Variable variable = ((EVariable)parameter).getVariable();
                    optionalVariableByParameter[i] = context.getVariableMap().get(variable);
                    varParameters[i] = variable;
                }
                else {
                    Variable temp = new Variable("temp", parameter.getType());
                    varParameters[i] = temp;
                    if(parameter.isPattern(0)) {
                        int tempId = context.addVariable(temp);
                        context.addConstraint(new ExpressionConstraint(context, temp, parameter, true));
                        optionalVariableByParameter[i] = tempId;
                    }
                    else {
                        optionalVariableByParameter[i] = -1;
                        parameter.forVariableUses(procedure);
                    }
                }
            }

            // Combine required and optional variables
            TIntHashSet allVariablesSet = new TIntHashSet();
            for(int v : optionalVariableByParameter)
                if(v >= 0)
                    allVariablesSet.add(v);

            context.addConstraint(new RelationConstraint(allVariablesSet.toArray(), varParameters, this,
                    optionalVariableByParameter, procedure.requiredVariablesMask));
        } catch(Exception e) {
            context.getQueryCompilationContext().getTypingContext().getErrorLog().log(location, e);
        }
    }

    @Override
    public void collectVars(TObjectIntHashMap<Variable> allVars,
            TIntHashSet vars) {
        for(Expression parameter : parameters)
            parameter.collectVars(allVars, vars);
    }

    @Override
    public Query replace(ReplaceContext context) {
        Type[] newTypeParameters;
        if(typeParameters == null)
            newTypeParameters = null;
        else {
            newTypeParameters = new Type[typeParameters.length];
            for(int i=0;i<typeParameters.length;++i)
                newTypeParameters[i] = typeParameters[i].replace(context.tvarMap);
        }
        return new QAtom(relation,
                newTypeParameters,
                Expression.replace(context, parameters));
    }

    @SuppressWarnings("unchecked")
    @Override
    public Diff[] derivate(THashMap<LocalRelation, Diffable> diffables) throws DerivateException {
        Diffable diffable = diffables.get(relation);
        if(diffable == null) {
            if(relation instanceof CompositeRelation && 
                    containsReferenceTo((CompositeRelation)relation,
                            (THashMap<SCLRelation, Diffable>)(THashMap)diffables))
                throw new DerivateException(location);
            return NO_DIFF;
        }
        else {
            Query[] eqs = new Query[parameters.length];
            for(int i=0;i<parameters.length;++i) {
                QAtom eq = new QAtom(EqRelation.INSTANCE, new Expression[] {
                        new EVariable(diffable.parameters[i]),
                        parameters[i]
                });
                eq.setLocationDeep(location);
                eq.typeParameters = new Type[] {parameters[i].getType()};
                eqs[i] = eq;
            }
            return new Diff[] { new Diff(diffable.id, new QConjunction(eqs)) };
        }
    }

    private static boolean containsReferenceTo(
            CompositeRelation relation,
            THashMap<SCLRelation, Diffable> diffables) {
        for(SCLRelation r : relation.getSubrelations())
            if(diffables.containsKey(r))
                return true;
            else if(r instanceof CompositeRelation &&
                    containsReferenceTo((CompositeRelation)r, diffables))
                return true;
        return false;
    }

    @Override
    public Query removeRelations(Set<SCLRelation> relations) {
        if(relations.contains(relation))
            return EMPTY_QUERY;
        else
            return this;
    }

    @Override
    public void setLocationDeep(long loc) {
        if(location == Locations.NO_LOCATION) {
            location = loc;
            for(Expression parameter : parameters)
                parameter.setLocationDeep(loc);
        }
    }

    @Override
    public void accept(QueryVisitor visitor) {
        visitor.visit(this);
    }

    @Override
    public void splitToPhases(TIntObjectHashMap<ArrayList<Query>> result) {
        int phase = relation.getPhase();
        ArrayList<Query> list = result.get(phase);
        if(list == null) {
            list = new ArrayList<Query>();
            result.put(phase, list);
        }
        list.add(this);
    }

    @Override
    public Query accept(QueryTransformer transformer) {
        return transformer.transform(this);
    }

}