/**
 * 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.backend.hadoop.executionengine.tez.plan;

import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.apache.pig.PigException;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.PhysicalOperator;
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.POLocalRearrange;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POPackage;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POSplit;
import org.apache.pig.backend.hadoop.executionengine.tez.plan.operator.BloomPackager;
import org.apache.pig.backend.hadoop.executionengine.tez.plan.operator.POLocalRearrangeTez;
import org.apache.pig.impl.plan.DepthFirstWalker;
import org.apache.pig.impl.plan.VisitorException;
import org.apache.pig.impl.plan.optimizer.OptimizerException;
import org.apache.pig.impl.util.Pair;

/**
 * A port of the POPackageAnnotator from MR to Tez.
 *
 */
public class TezPOPackageAnnotator extends TezOpPlanVisitor {

    /**
     * @param plan Tez plan to visit
     */
    public TezPOPackageAnnotator(TezOperPlan plan) {
        super(plan, new DepthFirstWalker<TezOperator, TezOperPlan>(plan));
    }

    @Override
    public void visitTezOp(TezOperator tezOp) throws VisitorException {
        if(!tezOp.plan.isEmpty()) {
            PackageDiscoverer pkgDiscoverer = new PackageDiscoverer(tezOp.plan);
            pkgDiscoverer.visit();
            POPackage pkg = pkgDiscoverer.getPkg();
            if(pkg != null) {
                handlePackage(tezOp, pkg);
            }
        }
    }

    private void handlePackage(TezOperator pkgTezOp, POPackage pkg) throws VisitorException {
        // the LocalRearrange(s) must be in the plan of a predecessor tez op
        int lrFound = 0;
        List<TezOperator> preds = this.mPlan.getPredecessors(pkgTezOp);
        for (Iterator<TezOperator> it = preds.iterator(); it.hasNext();) {
            TezOperator predTezOp = it.next();
            TezOperator predTezOpVertexGrp = null;
            if (predTezOp.isVertexGroup()) {
                predTezOpVertexGrp = predTezOp;
                // Just get one of the inputs to vertex group
                predTezOp = getPlan().getOperator(predTezOp.getVertexGroupMembers().get(0));
            }
            lrFound += patchPackage(predTezOp, predTezOpVertexGrp, pkgTezOp, pkg);
            if(lrFound == pkg.getNumInps()) {
                break;
            }
        }

        if(lrFound != pkg.getNumInps()) {
            int errCode = 2086;
            String msg = "Unexpected problem during optimization. "
                    + "Could not find all LocalRearrange operators. Expected: "
                    + pkg.getNumInps() + ", Found: " + lrFound;
            throw new OptimizerException(msg, errCode, PigException.BUG);
        }
    }

    private int patchPackage(TezOperator predTezOp,
            TezOperator predTezOpVertexGrp,
            TezOperator pkgTezOp,
            POPackage pkg) throws VisitorException {
        LoRearrangeDiscoverer lrDiscoverer = new LoRearrangeDiscoverer(
                predTezOp.plan, predTezOpVertexGrp, pkgTezOp, pkg);
        lrDiscoverer.visit();
        // let our caller know if we managed to patch
        // the package
        return lrDiscoverer.getLoRearrangeFound();
    }

    /**
     * Simple visitor of the "Reduce" physical plan
     * which will get a reference to the POPacakge
     * present in the plan
     */
    static class PackageDiscoverer extends PhyPlanVisitor {

        private POPackage pkg;

        public PackageDiscoverer(PhysicalPlan plan) {
            super(plan, new DepthFirstWalker<PhysicalOperator, PhysicalPlan>(plan));
        }

        @Override
        public void visitPackage(POPackage pkg) throws VisitorException {
            this.pkg = pkg;
        };

        /**
         * @return the pkg
         */
        public POPackage getPkg() {
            return pkg;
        }

    }

    /**
     * Physical Plan visitor which tries to get the
     * LocalRearrange(s) present in the plan (if any) and
     * annotate the POPackage given to it with the information
     * in the LocalRearrange (regarding columns in the "value"
     * present in the "key")
     */
    public static class LoRearrangeDiscoverer extends PhyPlanVisitor {

        private int loRearrangeFound = 0;
        private TezOperator pkgTezOp;
        private POPackage pkg;
        private TezOperator predTezOpVertexGrp;
        private boolean isPOSplit;

        public LoRearrangeDiscoverer(PhysicalPlan predPlan, TezOperator predTezOpVertexGrp, TezOperator pkgTezOp, POPackage pkg) {
            super(predPlan, new DepthFirstWalker<PhysicalOperator, PhysicalPlan>(predPlan));
            this.pkgTezOp = pkgTezOp;
            this.pkg = pkg;
            this.predTezOpVertexGrp = predTezOpVertexGrp;
        }

        @Override
        public void visitSplit(POSplit spl) throws VisitorException {
            isPOSplit = true;
            super.visitSplit(spl);
        }


        @Override
        public void visitLocalRearrange(POLocalRearrange lrearrange) throws VisitorException {
            POLocalRearrangeTez lr = (POLocalRearrangeTez) lrearrange;
            if (!(lr.isConnectedToPackage() && lr.containsOutputKey(pkgTezOp.getOperatorKey().toString()))) {
                return;
            }
            loRearrangeFound++;
            Map<Integer,Pair<Boolean, Map<Integer, Integer>>> keyInfo;

            if (pkg.getPkgr() instanceof LitePackager) {
                if(lrearrange.getIndex() != 0) {
                    // Throw some exception here
                    throw new RuntimeException("POLocalRearrange for POPackageLite cannot have index other than 0, but has index - "+lrearrange.getIndex());
                }
            }

            // annotate the package with information from the LORearrange
            // update the keyInfo information if already present in the POPackage
            keyInfo = pkg.getPkgr().getKeyInfo();
            if(keyInfo == null)
                keyInfo = new HashMap<Integer, Pair<Boolean, Map<Integer, Integer>>>();

            // For BloomPackager there is only one input, but the
            // POBuildBloomRearrangeTez index is that of the join's index and can be non-zero
            Integer index = (pkg.getPkgr() instanceof BloomPackager) ? 0 : Integer.valueOf(lrearrange.getIndex());
            if(keyInfo.get(index) != null) {
                if (isPOSplit) {
                    // Case of POSplit having more than one input in case of self join or union
                    loRearrangeFound--;
                } else {
                    // something is wrong - we should not be getting key info
                    // for the same index from two different Local Rearranges
                    int errCode = 2087;
                    String msg = "Unexpected problem during optimization." +
                            " Found index:" + lrearrange.getIndex() +
                            " in multiple LocalRearrange operators.";
                    throw new OptimizerException(msg, errCode, PigException.BUG);
                }

            }

            if (pkg.getPkgr() instanceof BloomPackager ) {
                keyInfo.put(index,
                        new Pair<Boolean, Map<Integer, Integer>>(
                                Boolean.FALSE, new HashMap<Integer, Integer>()));
                pkg.getPkgr().setKeyInfo(keyInfo);
            } else {
                keyInfo.put(index,
                        new Pair<Boolean, Map<Integer, Integer>>(
                                lrearrange.isProjectStar(), lrearrange.getProjectedColsMap()));
                pkg.getPkgr().setKeyInfo(keyInfo);
                pkg.getPkgr().setKeyTuple(lrearrange.isKeyTuple());
                pkg.getPkgr().setKeyCompound(lrearrange.isKeyCompound());
            }

        }

        /**
         * @return the loRearrangeFound
         */
        public int getLoRearrangeFound() {
            return loRearrangeFound;
        }

    }
}