/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.pig.pen;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import org.apache.pig.backend.hadoop.executionengine.physicalLayer.PhysicalOperator;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.EqualToExpr;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.GTOrEqualToExpr;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.GreaterThanExpr;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.LTOrEqualToExpr;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.LessThanExpr;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.NotEqualToExpr;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POAnd;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POBinCond;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POCast;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POIsNull;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POMapLookUp;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.PONegative;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.PONot;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POOr;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POProject;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.PORegexp;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POUserComparisonFunc;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POUserFunc;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.plans.PhyPlanVisitor;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.plans.PhysicalPlan;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.LitePackager;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POCounter;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.PODemux;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.PODistinct;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POFilter;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POForEach;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLimit;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLoad;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLocalRearrange;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POOptimizedForEach;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POPackage;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.PORank;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POSort;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POSplit;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POStore;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POStream;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POUnion;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.Tuple;
import org.apache.pig.impl.PigContext;
import org.apache.pig.impl.plan.DepthFirstWalker;
import org.apache.pig.impl.plan.PlanWalker;
import org.apache.pig.impl.plan.VisitorException;
import org.apache.pig.impl.util.IdentityHashSet;
import org.apache.pig.newplan.logical.relational.LogicalSchema;
import org.apache.pig.pen.util.LineageTracer;

/**
 * The class used to (re)attach illustrators to physical operators
 *
 *
 */
public class IllustratorAttacher extends PhyPlanVisitor {

    PigContext pigContext;

    LineageTracer lineage;

    HashMap<PhysicalOperator, Collection<IdentityHashSet<Tuple>>> poToEqclassesMap;

    private HashMap<PhysicalOperator, DataBag> poToDataMap;
    private int maxRecords;
    private boolean revisit = false;
    private ArrayList<Boolean[]> subExpResults = null;
    private final Map<POLoad, LogicalSchema> poloadToSchemaMap;

    public IllustratorAttacher(PhysicalPlan plan, LineageTracer lineage, int maxRecords,
        Map<POLoad, LogicalSchema> poLoadToSchemaMap, PigContext hadoopPigContext) throws VisitorException {
        super(plan, new DepthFirstWalker<PhysicalOperator, PhysicalPlan>(plan));
        pigContext = hadoopPigContext;
        this.lineage = lineage;
        poToEqclassesMap = new HashMap<PhysicalOperator, Collection<IdentityHashSet<Tuple>>>();
        poToDataMap = new HashMap<PhysicalOperator, DataBag>();
        this.maxRecords = maxRecords;
        this.poloadToSchemaMap = poLoadToSchemaMap;
    }

    /**
     * revisit an enhanced physical plan from MR compilation
     * @param plan a physical plan to be traversed
     */
    public void revisit(PhysicalPlan plan) throws VisitorException {
        pushWalker(new DepthFirstWalker<PhysicalOperator, PhysicalPlan>(plan));
        revisit = true;
        PhysicalPlan oriPlan = mPlan;
        mPlan = plan;
        visit();
        mPlan = oriPlan;
        popWalker();
    }

    private void setIllustrator(PhysicalOperator po, int nEqClasses) {
        if (revisit && po.getIllustrator() != null)
            return;
        LinkedList<IdentityHashSet<Tuple>> eqClasses = new LinkedList<IdentityHashSet<Tuple>>();
        poToEqclassesMap.put(po, eqClasses);
        for (int i = 0; i < nEqClasses; ++i)
        {
            IdentityHashSet<Tuple> eqClass = new IdentityHashSet<Tuple>();
            eqClasses.add(eqClass);
        }
        Illustrator illustrator = new Illustrator(lineage, eqClasses, this, pigContext);
        po.setIllustrator(illustrator);
        poToDataMap.put(po, illustrator.getData());
    }

    private void setIllustrator(PhysicalOperator po, LinkedList<IdentityHashSet<Tuple>> eqClasses) {
        if (revisit && po.getIllustrator() != null)
            return;
        Illustrator illustrator = new Illustrator(lineage, eqClasses, this, pigContext);
        po.setIllustrator(illustrator);
        if (eqClasses != null)
            poToEqclassesMap.put(po, eqClasses);
        poToDataMap.put(po, illustrator.getData());
    }

    void setIllustrator(PhysicalOperator po) {
        if (revisit && po.getIllustrator() != null)
            return;
        LinkedList<IdentityHashSet<Tuple>> eqClasses = new LinkedList<IdentityHashSet<Tuple>>();
        IdentityHashSet<Tuple> eqClass = new IdentityHashSet<Tuple>();
        eqClasses.add(eqClass);
        Illustrator illustrator = new Illustrator(lineage, eqClasses, this, pigContext);
        po.setIllustrator(illustrator);
        poToEqclassesMap.put(po, eqClasses);
        poToDataMap.put(po, illustrator.getData());
    }

    public Map<PhysicalOperator, DataBag> getDataMap() {
        return poToDataMap;
    }

    @Override
    public void visitLoad(POLoad ld) throws VisitorException{
        // LOAD from temporary files need no illustrator
        if (revisit)
            return;

        LinkedList<IdentityHashSet<Tuple>> eqClasses = new LinkedList<IdentityHashSet<Tuple>>();
        poToEqclassesMap.put(ld, eqClasses);

        IdentityHashSet<Tuple> eqClass = new IdentityHashSet<Tuple>();
        eqClasses.add(eqClass);
        Illustrator illustrator;
        illustrator = new Illustrator(lineage, eqClasses, maxRecords, this, poloadToSchemaMap.get(ld), pigContext);
        ld.setIllustrator(illustrator);
        poToDataMap.put(ld, illustrator.getData());
    }

    @Override
    public void visitStore(POStore st) throws VisitorException{
        setIllustrator(st, 1);
    }

    @Override
    public void visitFilter(POFilter fl) throws VisitorException{
        setIllustrator(fl, 0);
        subExpResults = fl.getIllustrator().getSubExpResults();
        innerPlanAttach(fl, fl.getPlan());
        subExpResults = null;
    }

    @Override
    public void visitLocalRearrange(POLocalRearrange lr) throws VisitorException{
      super.visitLocalRearrange(lr);
      setIllustrator(lr);
    }

    @Override
    public void visitPackage(POPackage pkg) throws VisitorException{
        setIllustrator(pkg, pkg.numberOfEquivalenceClasses());
    }

    @Override
    public void visitPOForEach(POForEach nfe) throws VisitorException {
        if (revisit && nfe.getIllustrator() != null)
            return;
        List<PhysicalPlan> innerPlans = nfe.getInputPlans();
        for (PhysicalPlan innerPlan : innerPlans)
          innerPlanAttach(nfe, innerPlan);
        List<PhysicalOperator> preds = mPlan.getPredecessors(nfe);
        if (preds != null && preds.size() == 1
                && preds.get(0) instanceof POPackage
                && ((POPackage) preds.get(0)).getPkgr().isDistinct()) {
            // equivalence class of POPackage for DISTINCT needs to be used
            //instead of the succeeding POForEach's equivalence class
            setIllustrator(nfe, preds.get(0).getIllustrator().getEquivalenceClasses());
            nfe.getIllustrator().setEqClassesShared();
        } else
            setIllustrator(nfe, 1);
    }

    @Override
    public void visitUnion(POUnion un) throws VisitorException{
        if (revisit && un.getIllustrator() != null)
          return;
        setIllustrator(un, null);
    }

    @Override
    public void visitSplit(POSplit spl) throws VisitorException{
        if (revisit && spl.getIllustrator() != null)
            return;
        for (PhysicalPlan poPlan : spl.getPlans())
          innerPlanAttach(spl, poPlan);
        setIllustrator(spl);
    }

    @Override
    public void visitDemux(PODemux demux) throws VisitorException{
      if (revisit && demux.getIllustrator() != null)
          return;
      List<PhysicalPlan> innerPlans = demux.getPlans();
      for (PhysicalPlan innerPlan : innerPlans)
        innerPlanAttach(demux, innerPlan);
      setIllustrator(demux);
    }

    @Override
    public void visitDistinct(PODistinct distinct) throws VisitorException {
        setIllustrator(distinct, 1);
    }

    @Override
    public void visitSort(POSort sort) throws VisitorException {
        setIllustrator(sort, 1);
    }

    @Override
    public void visitRank(PORank rank) throws VisitorException {
        setIllustrator(rank, 3);
    }

    @Override
    public void visitCounter(POCounter counter) throws VisitorException {
        setIllustrator(counter, 1);
    }

    @Override
    public void visitProject(POProject proj) throws VisitorException{
    }

    @Override
    public void visitGreaterThan(GreaterThanExpr grt) throws VisitorException{
        setIllustrator(grt, 0);
        if (!revisit && subExpResults != null)
            subExpResults.add(grt.getIllustrator().getSubExpResult());
    }

    @Override
    public void visitLessThan(LessThanExpr lt) throws VisitorException{
        setIllustrator(lt, 0);
        if (!revisit && subExpResults != null)
            subExpResults.add(lt.getIllustrator().getSubExpResult());
    }

    @Override
    public void visitGTOrEqual(GTOrEqualToExpr gte) throws VisitorException{
        setIllustrator(gte, 0);
        if (!revisit && subExpResults != null)
            subExpResults.add(gte.getIllustrator().getSubExpResult());
    }

    @Override
    public void visitLTOrEqual(LTOrEqualToExpr lte) throws VisitorException{
        setIllustrator(lte, 0);
        if (!revisit && subExpResults != null)
            subExpResults.add(lte.getIllustrator().getSubExpResult());
    }

    @Override
    public void visitEqualTo(EqualToExpr eq) throws VisitorException{
        setIllustrator(eq, 0);
        if (!revisit && subExpResults != null)
            subExpResults.add(eq.getIllustrator().getSubExpResult());
    }

    @Override
    public void visitNotEqualTo(NotEqualToExpr eq) throws VisitorException{
        setIllustrator(eq, 0);
        if (!revisit && subExpResults != null)
            subExpResults.add(eq.getIllustrator().getSubExpResult());
    }

    @Override
    public void visitRegexp(PORegexp re) throws VisitorException{
        setIllustrator(re, 0);
        if (!revisit && subExpResults != null)
            subExpResults.add(re.getIllustrator().getSubExpResult());
    }

    @Override
    public void visitIsNull(POIsNull isNull) throws VisitorException {
        setIllustrator(isNull, 0);
        if (!revisit && subExpResults != null)
            subExpResults.add(isNull.getIllustrator().getSubExpResult());
    }

    @Override
    public void visitAnd(POAnd and) throws VisitorException {
        setIllustrator(and, 0);
    }

    @Override
    public void visitOr(POOr or) throws VisitorException {
        setIllustrator(or, 0);
    }

    @Override
    public void visitNot(PONot not) throws VisitorException {
        setIllustrator(not, 0);
        if (!revisit && subExpResults != null)
            subExpResults.add(not.getIllustrator().getSubExpResult());
    }

    @Override
    public void visitBinCond(POBinCond binCond) {

    }

    @Override
    public void visitNegative(PONegative negative) {
      setIllustrator(negative, 1);
    }

    @Override
    public void visitUserFunc(POUserFunc userFunc) throws VisitorException {
    }

    @Override
    public void visitComparisonFunc(POUserComparisonFunc compFunc) throws VisitorException {
        // one each for >, ==, and <
        setIllustrator(compFunc, 3);
    }

    @Override
    public void visitMapLookUp(POMapLookUp mapLookUp) {
      setIllustrator(mapLookUp, 1);
    }

    @Override
    public void visitCast(POCast cast) {
    }

    @Override
    public void visitLimit(POLimit lim) throws VisitorException {
        setIllustrator(lim, 1);
    }

    @Override
    public void visitStream(POStream stream) throws VisitorException {
        setIllustrator(stream, 1);
    }

    /**
     * @param optimizedForEach
     */
    @Override
    public void visitPOOptimizedForEach(POOptimizedForEach optimizedForEach) throws VisitorException {
        visitPOForEach(optimizedForEach);
    }

    private void innerPlanAttach(PhysicalOperator po, PhysicalPlan plan) throws VisitorException {
        PlanWalker<PhysicalOperator, PhysicalPlan> childWalker =
              mCurrentWalker.spawnChildWalker(plan);
        pushWalker(childWalker);
        childWalker.walk(this);
        popWalker();
        LinkedList<IdentityHashSet<Tuple>> eqClasses = new LinkedList<IdentityHashSet<Tuple>>();
        if (subExpResults != null && !revisit) {
            int size = 1 << subExpResults.size();
            for (int i = 0; i < size; ++i) {
                eqClasses.add(new IdentityHashSet<Tuple>());
            }
            po.getIllustrator().setEquivalenceClasses(eqClasses, po);
        }
    }
}