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

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Properties;

import org.apache.pig.ExecType;
import org.apache.pig.PigServer;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.builtin.BinStorage;
import org.apache.pig.builtin.PigStorage;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.DefaultTuple;
import org.apache.pig.data.NonSpillableDataBag;
import org.apache.pig.data.SingleTupleBag;
import org.apache.pig.data.Tuple;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.util.MultiMap;
import org.junit.Before;
import org.junit.Test;

/**
 *
 */
public class TestDataBagAccess {
    private PigServer pigServer;

    @Before
    public void setUp() throws Exception{
        pigServer = new PigServer(ExecType.LOCAL, new Properties());
    }

    @Test
    public void testSingleTupleBagAcess() throws Exception {
        Tuple inputTuple = new DefaultTuple();
        inputTuple.append("a");
        inputTuple.append("b");

        SingleTupleBag bg = new SingleTupleBag(inputTuple);
        Iterator<Tuple> it = bg.iterator();
        assertEquals(inputTuple, it.next());
        assertFalse(it.hasNext());
    }

    @Test
    public void testNonSpillableDataBag() throws Exception {
        String[][] tupleContents = new String[][] {{"a", "b"},{"c", "d" }, { "e", "f"} };
        NonSpillableDataBag bg = new NonSpillableDataBag();
        for (int i = 0; i < tupleContents.length; i++) {
            bg.add(Util.createTuple(tupleContents[i]));
        }
        Iterator<Tuple> it = bg.iterator();
        int j = 0;
        while(it.hasNext()) {
            Tuple t = it.next();
            assertEquals(Util.createTuple(tupleContents[j]), t);
            j++;
        }
        assertEquals(tupleContents.length, j);
    }

    @Test
    public void testBagConstantAccess() throws IOException, ExecException {
        File input = Util.createInputFile("tmp", "",
                new String[] {"sampledata\tnot_used"});
        pigServer.registerQuery("a = load '"
                + Util.generateURI(Util.encodeEscape(input.toString()), pigServer.getPigContext()) + "';");
        pigServer.registerQuery("b = foreach a generate {(16, 4.0e-2, 'hello', -101)} as mybag:{t:(i: int, d: double, c: chararray, e : int)};");
        pigServer.registerQuery("c = foreach b generate mybag.i, mybag.d, mybag.c, mybag.e;");
        Iterator<Tuple> it = pigServer.openIterator("c");
        Tuple t = it.next();
        Object[] results = new Object[] { new Integer(16), new Double(4.0e-2), "hello", new Integer( -101 ) };
        Class[] resultClasses = new Class[] { Integer.class, Double.class, String.class, Integer.class };
        assertEquals(results.length, t.size());
        for (int i = 0; i < results.length; i++) {
            DataBag bag = (DataBag)t.get(i);
            assertEquals(results[i], bag.iterator().next().get(0));
            assertEquals(resultClasses[i], bag.iterator().next().get(0).getClass());
        }
    }

    @Test
    public void testBagConstantAccessFailure() throws IOException, ExecException {
        File input = Util.createInputFile("tmp", "",
                new String[] {"sampledata\tnot_used"});
        boolean exceptionOccured = false;
        pigServer.setValidateEachStatement(true);
        try {
            pigServer.registerQuery("a = load '"
                    + Util.generateURI(Util.encodeEscape(input.toString()), pigServer.getPigContext()) + "';");
            pigServer.registerQuery("b = foreach a generate {(16, 4.0e-2, 'hello')} as mybag:{t:(i: int, d: double, c: chararray)};");
            pigServer.registerQuery("c = foreach b generate mybag.t;");
            pigServer.explain("c", System.out);
        } catch(FrontendException e) {
            exceptionOccured = true;
            String msg = e.getMessage();
            Util.checkStrContainsSubStr(msg, "Cannot find field t in i:int,d:double,c:chararray");
        }
        assertTrue(exceptionOccured);
    }

    @Test
    public void testBagConstantFlatten1() throws IOException, ExecException {
        File input = Util.createInputFile("tmp", "",
                new String[] {"sampledata\tnot_used"});
        pigServer.registerQuery("A = load '"
                + Util.generateURI(Util.encodeEscape(input.toString()), pigServer.getPigContext()) + "';");
        pigServer.registerQuery("B = foreach A generate {(('p1-t1-e1', 'p1-t1-e2'),('p1-t2-e1', 'p1-t2-e2'))," +
                "(('p2-t1-e1', 'p2-t1-e2'), ('p2-t2-e1', 'p2-t2-e2'))};");
        pigServer.registerQuery("C = foreach B generate $0 as pairbag : { pair: ( t1: (e1, e2), t2: (e1, e2) ) };");
        pigServer.registerQuery("D = foreach C generate FLATTEN(pairbag);");
        pigServer.registerQuery("E = foreach D generate t1.e2 as t1e2, t2.e1 as t2e1;");
        Iterator<Tuple> it = pigServer.openIterator("E");
        // We should get the following two tuples as the result:
        // (p1-t1-e2,p1-t2-e1)
        // (p2-t1-e2,p2-t2-e1)
        Tuple t = it.next();
        assertEquals("p1-t1-e2", t.get(0));
        assertEquals("p1-t2-e1", t.get(1));
        t = it.next();
        assertEquals("p2-t1-e2", t.get(0));
        assertEquals("p2-t2-e1", t.get(1));
        assertFalse(it.hasNext());
    }

    @Test
    public void testBagConstantFlatten2() throws IOException, ExecException {
        File input = Util.createInputFile("tmp", "",
                new String[] {"somestring\t10\t{(a,10),(b,20)}"});
        pigServer.registerQuery("a = load '"
                + Util.generateURI(Util.encodeEscape(input.toString()), pigServer.getPigContext())
                + "' " + "as (str:chararray, intval:int, bg:bag{t:tuple(s:chararray, i:int)});");
        pigServer.registerQuery("b = foreach a generate str, intval, flatten(bg);");
        pigServer.registerQuery("c = foreach b generate str, intval, s, i;");
        Iterator<Tuple> it = pigServer.openIterator("c");
        int i = 0;
        Object[][] results = new Object[][] { {"somestring", new Integer(10), "a", new Integer(10)},
                {"somestring", new Integer(10), "b", new Integer(20) }};
        Class[] resultClasses = new Class[] { String.class, Integer.class, String.class, Integer.class };
        while(it.hasNext()) {
            Tuple t = it.next();
            for (int j = 0; j < resultClasses.length; j++) {
                assertEquals(results[i][j], t.get(j));
                assertEquals(resultClasses[j], t.get(j).getClass());
            }
            i++;
        }
        assertEquals(results.length, i);

        pigServer.registerQuery("c = foreach b generate str, intval, bg::s, bg::i;");
        it = pigServer.openIterator("c");
        i = 0;
        while(it.hasNext()) {
            Tuple t = it.next();
            for (int j = 0; j < resultClasses.length; j++) {
                assertEquals(results[i][j], t.get(j));
                assertEquals(resultClasses[j], t.get(j).getClass());
            }
            i++;
        }
        assertEquals(results.length, i);
    }

    @Test
    public void testBagStoreLoad() throws IOException, ExecException {
        File input = Util.createInputFile("tmp", "",
                new String[] {"a\tid1", "a\tid2", "a\tid3", "b\tid4", "b\tid5", "b\tid6"});
        pigServer.registerQuery("a = load '"
                + Util.generateURI(Util.encodeEscape(input.toString()), pigServer.getPigContext())
                + "' " + "as (s:chararray, id:chararray);");
        pigServer.registerQuery("b = group a by s;");
        Class[] loadStoreClasses = new Class[] { BinStorage.class, PigStorage.class };
        for (int i = 0; i < loadStoreClasses.length; i++) {
            String output = "TestDataBagAccess-testBagStoreLoad-" +
                             loadStoreClasses[i].getName() + ".txt";
            pigServer.deleteFile(output);
            pigServer.store("b", output, loadStoreClasses[i].getName());
            pigServer.registerQuery("c = load '" + output + "' using " + loadStoreClasses[i].getName() + "() AS " +
                    "(gp: chararray, bg:bag { t: tuple (sReLoaded: chararray, idReLoaded: chararray)});;");
            Iterator<Tuple> it = pigServer.openIterator("c");
            MultiMap<Object, Object> results = new MultiMap<Object, Object>();
            results.put("a", "id1");
            results.put("a", "id2");
            results.put("a", "id3");
            results.put("b", "id4");
            results.put("b", "id5");
            results.put("b", "id6");
            int j = 0;
            while(it.hasNext()) {
                Tuple t = it.next();
                Object groupKey = t.get(0);
                DataBag groupBag = (DataBag)t.get(1);
                Iterator<Tuple> bgIt = groupBag.iterator();
                int k = 0;
                while(bgIt.hasNext()) {
                    // a hash to make sure we don't see the
                    // same "ids" twice
                    HashMap<Object, Boolean> seen = new HashMap<Object, Boolean>();
                    Tuple bgt = bgIt.next();
                    // the first col is the group by key
                    assertTrue(bgt.get(0).equals(groupKey));
                    Collection<Object> values = results.get(groupKey);
                    // check that the second column is one
                    // of the "id" values associated with this
                    // group by key
                    assertTrue(values.contains(bgt.get(1)));
                    // check that we have not seen the same "id" value
                    // before
                    if(seen.containsKey(bgt.get(1)))
                        fail("LoadStoreClass used : " + loadStoreClasses[i].getName() + " " +
                        		", duplicate value (" + bgt.get(1) + ")");
                    else
                        seen.put(bgt.get(1), true);
                    k++;
                }
                // check that we saw 3 tuples in each group bag
                assertEquals(3, k);
                j++;
            }
            // make sure we saw the right number of high
            // level tuples
            assertEquals(results.keySet().size(), j);

            pigServer.registerQuery("d = foreach c generate gp, flatten(bg);");
            // results should be
            // a a id1
            // a a id2
            // a a id3
            // b b id4
            // b b id5
            // b b id6
            // However order is not guaranteed
            List<Tuple> resultTuples = new ArrayList<Tuple>();
            resultTuples.add(Util.createTuple(new String[] { "a", "a", "id1"}));
            resultTuples.add(Util.createTuple(new String[] { "a", "a", "id2"}));
            resultTuples.add(Util.createTuple(new String[] { "a", "a", "id3"}));
            resultTuples.add(Util.createTuple(new String[] { "b", "b", "id4"}));
            resultTuples.add(Util.createTuple(new String[] { "b", "b", "id5"}));
            resultTuples.add(Util.createTuple(new String[] { "b", "b", "id6"}));
            it = pigServer.openIterator("d");
            j = 0;
            HashMap<Tuple, Boolean> seen = new HashMap<Tuple, Boolean>();
            while(it.hasNext()) {
                Tuple t = it.next();
                assertTrue(resultTuples.contains(t));
                if(seen.containsKey(t)) {
                    fail("LoadStoreClass used : " + loadStoreClasses[i].getName() + " " +
                                ", duplicate tuple (" + t + ") encountered.");
                } else {
                    seen.put(t, true);
                }
                j++;
            }
            // check we got expected number of tuples
            assertEquals(resultTuples.size(), j);

            // same output as above - but projection based on aliases
            pigServer.registerQuery("e = foreach d generate gp, sReLoaded, idReLoaded;");
            it = pigServer.openIterator("e");
            j = 0;
            seen = new HashMap<Tuple, Boolean>();
            while(it.hasNext()) {
                Tuple t = it.next();
                assertTrue(resultTuples.contains(t));
                if(seen.containsKey(t)) {
                    fail("LoadStoreClass used : " + loadStoreClasses[i].getName() + " " +
                                ", duplicate tuple (" + t + ") encountered.");
                } else {
                    seen.put(t, true);
                }
                j++;
            }
            // check we got expected number of tuples
            assertEquals(resultTuples.size(), j);

            // same result as above but projection based on position specifiers
            pigServer.registerQuery("f = foreach d generate $0, $1, $2;");
            it = pigServer.openIterator("f");
            j = 0;
            seen = new HashMap<Tuple, Boolean>();
            while(it.hasNext()) {
                Tuple t = it.next();
                assertTrue(resultTuples.contains(t));
                if(seen.containsKey(t)) {
                    fail("LoadStoreClass used : " + loadStoreClasses[i].getName() + " " +
                                ", duplicate tuple (" + t + ") encountered.");
                } else {
                    seen.put(t, true);
                }
                j++;
            }
            // check we got expected number of tuples
            assertEquals(resultTuples.size(), j);


        }
    }
}