/*
 * 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.test;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;

import org.apache.pig.ExecType;
import org.apache.pig.PigServer;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.plans.PhysicalPlan;
import org.apache.pig.impl.PigContext;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
import org.apache.pig.newplan.logical.optimizer.LogicalPlanPrinter;
import org.apache.pig.newplan.logical.optimizer.ProjectionPatcher;
import org.apache.pig.newplan.logical.optimizer.SchemaPatcher;
import org.apache.pig.newplan.logical.relational.LOLoad;
import org.apache.pig.newplan.logical.rules.ColumnMapKeyPrune;
import org.apache.pig.newplan.logical.rules.MapKeysPruneHelper;
import org.apache.pig.newplan.optimizer.PlanOptimizer;
import org.apache.pig.newplan.optimizer.PlanTransformListener;
import org.apache.pig.newplan.optimizer.Rule;
import org.junit.Test;

public class TestNewPlanPruneMapKeys {
    PigContext pc = new PigContext(ExecType.LOCAL, new Properties());

    private org.apache.pig.newplan.logical.relational.LogicalPlan migratePlan(String query) throws Exception{
        PigServer pigServer = new PigServer( pc );
        org.apache.pig.newplan.logical.relational.LogicalPlan newPlan =
        Util.buildLp(pigServer, query);

        // run filter rule
        Set<Rule> s = new HashSet<Rule>();
        List<Set<Rule>> ls = new ArrayList<Set<Rule>>();
        ls.add(s);
        // Add the PruneMap Filter
        Rule r = new ColumnMapKeyPrune("PruneMapKeys");
        s.add(r);

        printPlan((org.apache.pig.newplan.logical.relational.LogicalPlan)newPlan);

        // Run the optimizer
        MyPlanOptimizer optimizer = new MyPlanOptimizer(newPlan, ls, 3);
        optimizer.addPlanTransformListener(new ProjectionPatcher());
        optimizer.addPlanTransformListener(new SchemaPatcher());
        optimizer.optimize();

        return newPlan;
    }

    public class MyPlanOptimizer extends PlanOptimizer {

        protected MyPlanOptimizer(OperatorPlan p, List<Set<Rule>> rs,
                int iterations) {
            super(p, rs, iterations);
        }

        public void addPlanTransformListener(PlanTransformListener listener) {
            super.addPlanTransformListener(listener);
        }

    }

    @Test
    @SuppressWarnings("unchecked")
    public void testSimplePlan() throws Exception {
        String query = "a =load 'd.txt' as (a:map[], b:int, c:float);" +
        "b = filter a by a#'name' == 'hello';" +
        "store b into 'empty';";

        org.apache.pig.newplan.logical.relational.LogicalPlan newLogicalPlan = migratePlan(query);

        List<Operator> sources = newLogicalPlan.getSources();
        assertEquals( 1, sources.size() );
        for( Operator source : sources ) {
            Map<Long,Set<String>> annotation =
                (Map<Long, Set<String>>) source.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
            assertTrue(annotation == null || annotation.isEmpty() );
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testSimplePlan2() throws Exception {
        String query = "a =load 'd.txt' as (a:map[], b:int, c:float);" +
        "b = filter a by a#'name' == 'hello';" +
        "c = foreach b generate b,c;" +
        "store c into 'empty';";

        org.apache.pig.newplan.logical.relational.LogicalPlan newLogicalPlan = migratePlan(query);

        assertEquals( 1, newLogicalPlan.getSources().size() );
        LOLoad load = (LOLoad) newLogicalPlan.getSources().get(0);
        Map<Long,Set<String>> annotation =
            (Map<Long, Set<String>>) load.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
        assertNotNull(annotation);
        assertEquals( 1, annotation.keySet().size() );
        Integer[] keySet = annotation.keySet().toArray( new Integer[0] );
        assertEquals( new Integer(0), keySet[0] );
        Set<String> keys = annotation.get(0);
        assertEquals( 1, keys.size() );
        assertEquals( "name", keys.toArray( new String[0] )[0] );
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testSimplePlan3() throws Exception {
        String query = "a =load 'd.txt' as (a:map[], b:int, c:float);" +
        "b = filter a by a#'name' == 'hello';" +
        "c = foreach b generate a#'age',b,c;" +
        "store c into 'empty';";

        org.apache.pig.newplan.logical.relational.LogicalPlan newLogicalPlan = migratePlan(query);

        assertEquals( 1, newLogicalPlan.getSources().size() );
        LOLoad load = (LOLoad) newLogicalPlan.getSources().get(0);
        Map<Long,Set<String>> annotation =
            (Map<Long, Set<String>>) load.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
        assertNotNull(annotation);
        assertEquals( 1, annotation.keySet().size() );
        Integer[] keySet = annotation.keySet().toArray( new Integer[0] );
        assertEquals( new Integer(0), keySet[0] );
        Set<String> keys = annotation.get(0);
        assertEquals( 2, keys.size() );
        assertTrue( keys.contains("name") );
        assertTrue( keys.contains("age"));
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testSimplePlan4() throws Exception {
        String query = "a =load 'd.txt' as (a:map[], b:int, c:float);" +
        "b = filter a by a#'name' == 'hello';" +
        "c = foreach b generate a#'age',a,b,c;" +
        "store c into 'empty';";

        org.apache.pig.newplan.logical.relational.LogicalPlan newLogicalPlan = migratePlan(query);

        List<Operator> sources = newLogicalPlan.getSources();
        assertEquals( 1, sources.size() );
        for( Operator source : sources ) {
            Map<Long,Set<String>> annotation =
                (Map<Long, Set<String>>) source.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
            assertTrue(annotation == null || annotation.isEmpty() );
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testSimplePlan5() throws Exception {
        String query = "a =load 'd.txt' as (a:chararray, b:int, c:float);" +
        "b = filter a by a == 'hello';" +
        "c = foreach b generate a,b,c;" +
        "store c into 'empty';";

        org.apache.pig.newplan.logical.relational.LogicalPlan newLogicalPlan = migratePlan(query);

        List<Operator> sources = newLogicalPlan.getSources();
        assertEquals( 1, sources.size() );
        for( Operator source : sources ) {
            Map<Long,Set<String>> annotation =
                (Map<Long, Set<String>>) source.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
            assertNull(annotation);
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testSimplePlan6() throws Exception {
        String query = "a =load 'd.txt';" +
        "b = filter a by $0 == 'hello';" +
        "c = foreach b generate $0,$1,$2;" +
        "store c into 'empty';";

        org.apache.pig.newplan.logical.relational.LogicalPlan newLogicalPlan = migratePlan(query);

        List<Operator> sources = newLogicalPlan.getSources();
        assertEquals( 1, sources.size() );
        for( Operator source : sources ) {
            Map<Long,Set<String>> annotation =
                (Map<Long, Set<String>>) source.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
            assertTrue(annotation == null );
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testSimplePlan7() throws Exception {
        String query = "a =load 'd.txt';" +
        "a1 = load 'b.txt' as (a:map[],b:int, c:float);" +
        "b = join a by $0, a1 by a#'name';" +
        "c = foreach b generate $0,$1,$2;" +
        "store c into 'empty';";

        org.apache.pig.newplan.logical.relational.LogicalPlan newLogicalPlan = migratePlan(query);
        printPlan( newLogicalPlan );

        List<Operator> sources = newLogicalPlan.getSources();
        assertEquals( 2, sources.size() );
        for( Operator source : sources ) {
            Map<Long,Set<String>> annotation =
                (Map<Long, Set<String>>) source.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
            assertTrue( annotation == null || annotation.isEmpty() );
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testSimplePlan8() throws Exception {
        String query = "a =load 'd.txt';" +
        "a1 = load 'b.txt' as (a:chararray,b:int, c:float);" +
        "b = join a by $0, a1 by a;" +
        "c = foreach b generate $0,$1,$2;" +
        "store c into 'empty';";

        org.apache.pig.newplan.logical.relational.LogicalPlan newLogicalPlan = migratePlan(query);
        printPlan(newLogicalPlan);

        List<Operator> sources = newLogicalPlan.getSources();
        assertEquals( 2, sources.size() );
        for( Operator source : sources ) {
            Map<Long,Set<String>> annotation =
                (Map<Long, Set<String>>) source.getAnnotation(MapKeysPruneHelper.REQUIRED_MAPKEYS);
            assertTrue( annotation == null || annotation.isEmpty() );
        }
    }

    public void printPlan(org.apache.pig.newplan.logical.relational.LogicalPlan logicalPlan ) throws Exception {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        PrintStream ps = new PrintStream(out);
        LogicalPlanPrinter pp = new LogicalPlanPrinter(logicalPlan,ps);
        pp.visit();
        System.err.println(out.toString());
    }

    public void printPlan(PhysicalPlan physicalPlan) throws Exception {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        PrintStream ps = new PrintStream(out);
        physicalPlan.explain(ps, "text", true);
        System.err.println(out.toString());
    }
}