/*
 * 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.newplan.logical.visitor;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.pig.PigException;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.plan.VisitorException;
import org.apache.pig.impl.util.MultiMap;
import org.apache.pig.impl.util.Pair;
import org.apache.pig.newplan.DependencyOrderWalker;
import org.apache.pig.newplan.DepthFirstWalker;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
import org.apache.pig.newplan.PlanWalker;
import org.apache.pig.newplan.logical.expression.LogicalExpressionPlan;
import org.apache.pig.newplan.logical.expression.LogicalExpressionVisitor;
import org.apache.pig.newplan.logical.expression.ProjectExpression;
import org.apache.pig.newplan.logical.relational.LOCogroup;
import org.apache.pig.newplan.logical.relational.LOCube;
import org.apache.pig.newplan.logical.relational.LOForEach;
import org.apache.pig.newplan.logical.relational.LOGenerate;
import org.apache.pig.newplan.logical.relational.LOInnerLoad;
import org.apache.pig.newplan.logical.relational.LOJoin;
import org.apache.pig.newplan.logical.relational.LORank;
import org.apache.pig.newplan.logical.relational.LOSort;
import org.apache.pig.newplan.logical.relational.LogicalPlan;
import org.apache.pig.newplan.logical.relational.LogicalRelationalNodesVisitor;
import org.apache.pig.newplan.logical.relational.LogicalRelationalOperator;
import org.apache.pig.newplan.logical.relational.LogicalSchema;
import org.apache.pig.newplan.logical.relational.LogicalSchema.LogicalFieldSchema;

import com.google.common.primitives.Booleans;

/**
 * A visitor to walk operators that contain a nested plan and translate project( * )
 * operators to a list of projection operators, i.e., 
 * project( * ) -> project(0), project(1), ... project(n-2), project(n-1)
 * If input schema is null, project(*) is not expanded.
 * It also expands project range ( eg $1 .. $5). It won't expand project-range-to-end
 * (eg $3 ..) if the input schema is null.
 * 
 */
public class ProjectStarExpander extends LogicalRelationalNodesVisitor{

    public ProjectStarExpander(OperatorPlan plan)
    throws FrontendException {
        super(plan, new DependencyOrderWalker(plan));
    }

    @Override
    public void visit(LOSort sort) throws FrontendException{
        List<LogicalExpressionPlan> expPlans = sort.getSortColPlans();
        List<Boolean> ascOrder = sort.getAscendingCols();

        // new expressionplans and sort order list after star expansion
        List<LogicalExpressionPlan> newExpPlans =
            new ArrayList<LogicalExpressionPlan>();
        List<Boolean> newAscOrder =  new ArrayList<Boolean>();

        if(expPlans.size() != ascOrder.size()){
            throw new AssertionError("Size of expPlans and ascorder should be same");
        }            
        
        for(int i=0; i < expPlans.size(); i++){
            //expand the plan
            LogicalExpressionPlan ithExpPlan = expPlans.get(i);
            List<LogicalExpressionPlan> expandedPlans = expandPlan(ithExpPlan, 0);
            newExpPlans.addAll(expandedPlans);

            //add corresponding isAsc flags
            Boolean isAsc = ascOrder.get(i);
            for(int j=0; j < expandedPlans.size(); j++){
                newAscOrder.add(isAsc);
            }
        }

        //check if there is a project-star-to-end followed by another sort plan
        // in the expanded plans (can happen if there is no input schema)
        for(int i=0; i < newExpPlans.size(); i++){
            ProjectExpression proj = getProjectStar(newExpPlans.get(i));
            if(proj != null && 
                    proj.isRangeProject() && proj.getEndCol() == -1 &&
                    i != newExpPlans.size() -1
            ){
                //because of order by sampler logic limitation, this is not
                //supported right now
                String msg = "Project-range to end (eg. x..)" +
                " is supported in order-by only as last sort column";
                throw new FrontendException(
                        msg,
                        1128,
                        PigException.INPUT
                );
            }
        }

        sort.setSortColPlans(newExpPlans);
        sort.setAscendingCols(newAscOrder);
    }

    @Override
    public void visit(LORank rank) throws FrontendException {

        List<LogicalExpressionPlan> expPlans = rank.getRankColPlans();
        List<Boolean> ascOrder = rank.getAscendingCol();

        List<LogicalExpressionPlan> newExpPlans = new ArrayList<LogicalExpressionPlan>();
        List<Boolean> newAscOrder = new ArrayList<Boolean>();

        if (expPlans.size() != ascOrder.size()) {
            throw new AssertionError(
                    "Size of expPlans and ascorder should be same");
        }

        for (int i = 0; i < expPlans.size(); i++) {
            // expand the plan
            LogicalExpressionPlan ithExpPlan = expPlans.get(i);
            List<LogicalExpressionPlan> expandedPlans = expandPlan(ithExpPlan,
                    0);
            newExpPlans.addAll(expandedPlans);

            // add corresponding isAsc flags
            Boolean isAsc = ascOrder.get(i);
            for (int j = 0; j < expandedPlans.size(); j++) {
                newAscOrder.add(isAsc);
            }
        }

        // check if there is a project-star-to-end followed by another sort plan
        // in the expanded plans (can happen if there is no input schema)
        for (int i = 0; i < newExpPlans.size(); i++) {
            ProjectExpression proj = getProjectStar(newExpPlans.get(i));
            if (proj != null && proj.isRangeProject() && proj.getEndCol() == -1
                    && i != newExpPlans.size() - 1) {
                String msg = "Project-range to end (eg. x..)"
                        + " is supported in rank-by only as last rank column";
                throw new FrontendException(msg, 1128, PigException.INPUT);
            }
        }

        rank.setRankColPlan(newExpPlans);
        rank.setAscendingCol(newAscOrder);
    }

    /**
     * Expand plan into multiple plans if the plan contains a project star,
     * if there is no project star the returned list contains the plan argument.
     * @param plan
     * @return
     * @throws FrontendException
     */
    private List<LogicalExpressionPlan> expandPlan(LogicalExpressionPlan plan, int inputNum)
    throws FrontendException {
        List<LogicalExpressionPlan> expandedPlans;
        ProjectExpression projStar = getProjectStar(plan);
        if(projStar != null){
            // expand the plan into multiple plans
            return expandPlan(plan, projStar, inputNum);
        }else{
            //no project star to expand
            expandedPlans = new ArrayList<LogicalExpressionPlan>();
            expandedPlans.add(plan);
        }
        return expandedPlans;
    }

    @Override
    public void visit(LOCogroup cg) throws FrontendException{
        
        MultiMap<Integer, LogicalExpressionPlan> inpExprPlans = 
            cg.getExpressionPlans();
 
        //modify the plans if they have project-star
        expandPlans(inpExprPlans);
        
        
        //do some validations -
        List<Operator> inputs = cg.getInputs((LogicalPlan) cg.getPlan());

        // check if after translation none of group by plans in a cogroup
        // have a project(*) - if they still do it's because the input
        // for the project(*) did not have a schema - in this case, we should
        // error out since we could have different number/types of 
        // cogroup keys
        if(inputs.size() > 1) { // only for cogroups
            for(int i=0; i<inputs.size(); i++)
                for(LogicalExpressionPlan lp: inpExprPlans.get(i)) {
                    if(getProjectStar(lp) != null) {
                        String msg = "Cogroup/Group by '*' or 'x..' " +
                        "(range of columns to the end) " +
                        "is only allowed if the input has a schema";
                        throw new VisitorException( cg,
                                msg,
                                1123,
                                PigException.INPUT
                        );
                    }
                }
        }
        // check if after translation all group by plans have same arity
        int arity = inpExprPlans.get(0).size();
        for(int i=1; i<inputs.size(); i++){
            if(arity != inpExprPlans.get(i).size()) {
                String msg = "The arity of cogroup/group by columns " +
                "do not match";
                throw new VisitorException(cg,
                        msg,
                        1122,
                        PigException.INPUT
                );
            }
        }

    }

    @Override
    public void visit(LOCube cu) throws FrontendException {

	MultiMap<Integer, LogicalExpressionPlan> inpExprPlans = cu.getExpressionPlans();

	// modify the plans if they have project-star
	expandPlans(inpExprPlans);

    }

    @Override
    public void visit(LOJoin join) throws FrontendException{
        expandPlans(join.getExpressionPlans());
    }

    @Override
    public void visit(LOForEach foreach) throws FrontendException{
        //in case of LOForeach , expand when inner plan has a single project-star
        // and its input LOInnerLoad also is a project-star
        // then Reset the input number in project expressions
        
        LogicalPlan innerPlan = foreach.getInnerPlan();
        
        //visit the inner plan first
        PlanWalker newWalker = currentWalker.spawnChildWalker(innerPlan);
        pushWalker(newWalker);
        currentWalker.walk(this);
        popWalker();
        
        //get the LOGenerate
        List<Operator> feOutputs = innerPlan.getSinks();
        LOGenerate gen = null;
        for( Operator op  : feOutputs){
            if(op instanceof LOGenerate){
                if(gen != null){
                    String msg = "Expected single LOGenerate output in innerplan of foreach";
                    throw new VisitorException(foreach,
                            msg,
                            2266,
                            PigException.BUG
                    );
                }
                gen = (LOGenerate) op;
            }
        }
        
        //work on the generate plan, flatten and user schema
        List<LogicalExpressionPlan> expPlans = gen.getOutputPlans();
        List<LogicalExpressionPlan> newExpPlans = new ArrayList<LogicalExpressionPlan>();
        
        List<Operator> loGenPreds = innerPlan.getPredecessors(gen);
        
        if(loGenPreds == null){
            // there are no LOInnerLoads , must be working on just constants
            // no project-star expansion to be done
            return;
        }
        
        List<LogicalSchema> userSchema = gen.getUserDefinedSchema();
        List<LogicalSchema> newUserSchema = null;
        if(userSchema != null){
            newUserSchema = new ArrayList<LogicalSchema>();
        }
        
        boolean[] flattens = gen.getFlattenFlags();
        List<Boolean> newFlattens = new ArrayList<Boolean>(flattens.length);

        //get mapping of LOGenerate predecessor current position to object
        Map<Integer, LogicalRelationalOperator> oldPos2Rel =
            new HashMap<Integer, LogicalRelationalOperator>();
        
        for(int i=0; i<loGenPreds.size(); i++){
            oldPos2Rel.put(i, (LogicalRelationalOperator) loGenPreds.get(i));
        }
        
        //get schema of predecessor, project-star expansion needs a schema
        LogicalRelationalOperator pred =
            (LogicalRelationalOperator) foreach.getPlan().getPredecessors(foreach).get(0);
        LogicalSchema inpSch = pred.getSchema();
 
        //store mapping between the projection in inner plans of
        // of LOGenerate to the input relation object
        Map<ProjectExpression, LogicalRelationalOperator> proj2InpRel =
            new HashMap<ProjectExpression, LogicalRelationalOperator>();
        
        
        for(int i=0; i<expPlans.size(); i++){
            LogicalExpressionPlan expPlan = expPlans.get(i);
            ProjectExpression projStar = getProjectLonelyStar(expPlan, oldPos2Rel);

            boolean foundExpandableProject = false;
            if(projStar != null){              
                //there is a project-star to be expanded

                LogicalSchema userStarSch = null;
                if(userSchema != null && userSchema.get(i) != null){
                    userStarSch = userSchema.get(i);
                }


                //the range values are set in the project in LOInnerLoad
                ProjectExpression loInnerProj = ((LOInnerLoad)oldPos2Rel.get(projStar.getInputNum())).getProjection();

                int firstProjCol = 0;
                int lastProjCol = 0;
                
                if(loInnerProj.isRangeProject()){
                    loInnerProj.setColumnNumberFromAlias();
                    firstProjCol = loInnerProj.getStartCol();
                    lastProjCol = loInnerProj.getEndCol();
                }

                
                boolean isProjectToEnd = loInnerProj.isProjectStar() || 
                    (loInnerProj.isRangeProject() && lastProjCol == -1); 
                
                //can't expand if there is no input schema, and this is
                // as project star or project-range-to-end
                if( !(inpSch == null && isProjectToEnd) ){
                    
                    foundExpandableProject = true;

                    if(isProjectToEnd)
                        lastProjCol = inpSch.size() - 1;

                    //replacing the existing project star with new ones
                    expPlan.remove(projStar);

                    //remove the LOInnerLoad with star
                    LOInnerLoad oldLOInnerLoad = (LOInnerLoad)oldPos2Rel.get(projStar.getInputNum());
                    innerPlan.disconnect(oldLOInnerLoad, gen);
                    innerPlan.remove(oldLOInnerLoad);


                    //generate new exp plan, inner load for each field in schema
                    for(int j = firstProjCol; j <= lastProjCol; j++){

                        //add new LOInnerLoad
                        LOInnerLoad newInLoad = new LOInnerLoad(innerPlan, foreach, j);
                        innerPlan.add(newInLoad);
                        innerPlan.connect(newInLoad, gen);


                        // new expression plan and proj
                        LogicalExpressionPlan newExpPlan = new LogicalExpressionPlan();
                        newExpPlans.add(newExpPlan);

                        ProjectExpression newProj =
                            new ProjectExpression(newExpPlan, -2, -1, gen);

                        proj2InpRel.put(newProj, newInLoad);

                        newFlattens.add(flattens[i]);
                        if(newUserSchema != null ){
                            //index into user specified schema
                            int schIdx = j - firstProjCol;
                            if(userStarSch != null 
                                    && userStarSch.getFields().size() > schIdx
                                    && userStarSch.getField(schIdx) != null){

                                //if the project-star field has user specified schema, use the
                                // j'th field for this column
                                LogicalSchema sch = new LogicalSchema();
                                sch.addField(new LogicalFieldSchema(userStarSch.getField(schIdx)));
                                newUserSchema.add(sch);
                            }
                            else{
                                newUserSchema.add(null);
                            }
                        }
                    }
                }
            }

            if(!foundExpandableProject){ //no project-star that could be expanded

                //get all projects in here 
                FindProjects findProjs = new FindProjects(expPlan);
                findProjs.visit();
                List<ProjectExpression> projs = findProjs.getProjs();

                //create a mapping of project expression to their inputs
                for(ProjectExpression proj : projs){
                    proj2InpRel.put(proj, oldPos2Rel.get(proj.getInputNum()));
                }

                newExpPlans.add(expPlan);

                newFlattens.add(flattens[i]);
                if(newUserSchema != null)
                    newUserSchema.add(userSchema.get(i));

            }
        }

        //get mapping of LoGenerate input relation to current position
        Map<LogicalRelationalOperator, Integer> rel2pos = new HashMap<LogicalRelationalOperator, Integer>();
        List<Operator> newGenPreds = innerPlan.getPredecessors(gen);
        int numNewGenPreds = 0;
        if(newGenPreds != null)
            numNewGenPreds = newGenPreds.size();
            
        for(int i=0; i<numNewGenPreds; i++){
            rel2pos.put((LogicalRelationalOperator) newGenPreds.get(i),i);
        }
        
        //correct the input num for projects
        for(Entry<ProjectExpression, LogicalRelationalOperator> projAndInp : proj2InpRel.entrySet()){
           ProjectExpression proj = projAndInp.getKey();
           LogicalRelationalOperator rel = projAndInp.getValue();
           proj.setInputNum(rel2pos.get(rel));
        }
        
        // set the new lists
        gen.setOutputPlans(newExpPlans);
        gen.setFlattenFlags(Booleans.toArray(newFlattens));
        gen.setUserDefinedSchema(newUserSchema);
        
        gen.resetSchema();
        foreach.resetSchema();
        
    }
    
    
    static class FindProjects extends LogicalExpressionVisitor{

        private List<ProjectExpression> projs = new ArrayList<ProjectExpression>();

        protected FindProjects(LogicalExpressionPlan plan)
                throws FrontendException {
            super(plan, new DepthFirstWalker(plan));
        }
        
        @Override
        public void visit(ProjectExpression proj){
           projs.add(proj);
        }

        public List<ProjectExpression> getProjs(){
            return projs;
        }
    }

    /**
     * Find project-star in foreach statement. The LOInnerLoad corresponding
     * to the project-star also needs to have a project-star
     * @param expPlan - expression plan
     * @param oldPos2Rel - inner relational plan of foreach
     * @return ProjectExpression that satisfies the conditions
     * @throws FrontendException 
     */
    private ProjectExpression getProjectLonelyStar(LogicalExpressionPlan expPlan,
            Map<Integer, LogicalRelationalOperator> oldPos2Rel) throws FrontendException {

        //the expression plan should have just a single project
        if(expPlan.size() == 0 || expPlan.size() > 1){
            return null;
        }

        Operator outputOp = expPlan.getOperators().next();
        if(outputOp instanceof ProjectExpression){
            ProjectExpression proj = (ProjectExpression)outputOp;
            //check if ProjectExpression is projectStar
            if(proj.isProjectStar()){
                //now check if its input is a LOInnerLoad and it is projectStar 
                // or range project
                LogicalRelationalOperator inputRel = oldPos2Rel.get(proj.getInputNum());
                if(! (inputRel  instanceof LOInnerLoad)){
                    return null;
                }

                ProjectExpression innerProj = ((LOInnerLoad) inputRel).getProjection(); 
                if( innerProj.isRangeOrStarProject()){
                    return proj;
                }
            }
        }
        return null;
    }


    private void expandPlans(
            MultiMap<Integer, LogicalExpressionPlan> inpExprPlans)
    throws FrontendException {
 
        //for each input relation, expand any logical plan that has project-star
        for(int i=0; i< inpExprPlans.size() ; i++){
            List<LogicalExpressionPlan> plans = inpExprPlans.get(i);
            List<LogicalExpressionPlan> newPlans =
                new ArrayList<LogicalExpressionPlan>();
            for(LogicalExpressionPlan plan : plans){
                newPlans.addAll(expandPlan(plan, i));
            }
            inpExprPlans.removeKey(i);
            inpExprPlans.put(i, newPlans);
        }
   
    }

    /**
     * expand this plan containing project star to multiple plans 
     * each projecting a single column
     * @param expPlan
     * @param proj
     * @return
     * @throws FrontendException 
     */
    private List<LogicalExpressionPlan> expandPlan(
            LogicalExpressionPlan expPlan, ProjectExpression proj, int inputNum)
            throws FrontendException {
        
        Pair<Integer, Integer> startAndEndProjs =
            ProjectStarExpanderUtil.getProjectStartEndCols(expPlan, proj);  
        List<LogicalExpressionPlan> newPlans = new ArrayList<LogicalExpressionPlan>();

        if(startAndEndProjs == null){
            // can't expand this project
            newPlans.add(expPlan);
            return newPlans;
        }

        //expand from firstProjCol to lastProjCol 
        int firstProjCol = startAndEndProjs.first;
        int lastProjCol = startAndEndProjs.second;

        LogicalRelationalOperator relOp = proj.getAttachedRelationalOp();
        for(int i = firstProjCol; i <= lastProjCol; i++){
            newPlans.add(createExpPlanWithProj(relOp, inputNum, i));
        }

        return newPlans;

    }

    /**
     * Create new logical plan with a project that is attached to LogicalRelation
     * attachRel and projects i'th column from input
     * @param attachRel 
     * @param inputNum
     * @param colNum 
     * @return
     */
    private LogicalExpressionPlan createExpPlanWithProj(
            LogicalRelationalOperator attachRel, 
            int inputNum, int colNum) {
        LogicalExpressionPlan newExpPlan = new LogicalExpressionPlan();
        ProjectExpression newProj = 
            new ProjectExpression(newExpPlan, inputNum, colNum, attachRel);
        newExpPlan.add(newProj);
        return newExpPlan;
    }

    /**
     * if LogicalExpressionPlan argument has a project star output then
     * return it, otherwise return null
     * @param expPlan
     * @return
     * @throws FrontendException 
     */
    private ProjectExpression getProjectStar(LogicalExpressionPlan expPlan)
    throws FrontendException {
        List<Operator> outputs = expPlan.getSources();
        ProjectExpression projStar = null;
        for(Operator outputOp : outputs){
            if(outputOp instanceof ProjectExpression){
                ProjectExpression proj = (ProjectExpression)outputOp;
                if(proj.isRangeOrStarProject()){
                    if(outputs.size() > 1){
                        String msg = "More than one operator in an expression plan" +
                        " containing project star(*)/project-range (..)";
                        throw new VisitorException(proj,
                                msg,
                                2264,
                                PigException.BUG
                        );
                    }
                    projStar = proj;
                }
            }
        }
        return projStar;
    }


}