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

import java.io.IOException;
import java.io.ObjectInputStream;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.pig.ComparisonFunc;
import org.apache.pig.EvalFunc;
import org.apache.pig.FuncSpec;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.MRCompiler;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.partitioners.CountingMap;
import org.apache.pig.data.BagFactory;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.DataType;
import org.apache.pig.data.InternalMap;
import org.apache.pig.data.NonSpillableDataBag;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
import org.apache.pig.impl.PigContext;


public class FindQuantiles extends EvalFunc<Map<String, Object>>{
    // keys for the weightedparts Map
    public static final String QUANTILES_LIST = "quantiles.list";
    public static final String WEIGHTED_PARTS = "weighted.parts";

    BagFactory mBagFactory = BagFactory.getInstance();
    TupleFactory mTupleFactory = TupleFactory.getInstance();

    boolean[] mAsc;
    enum State { ALL_ASC, ALL_DESC, MIXED };
    State mState;
    
    protected Integer numQuantiles = null;
    protected DataBag samples = null;
    
    private class SortComparator implements Comparator<Tuple> {
        @Override
        @SuppressWarnings("unchecked")
        public int compare(Tuple t1, Tuple t2) {
            switch (mState) {
            case ALL_ASC:
                return t1.compareTo(t2);

            case ALL_DESC:
                return t2.compareTo(t1);

            case MIXED:
                // Have to break the tuple down and compare it field to field.
                int sz1 = t1.size();
                int sz2 = t2.size();
                if (sz2 < sz1) {
                    return 1;
                } else if (sz2 > sz1) {
                    return -1;
                } else {
                    for (int i = 0; i < sz1; i++) {
                        try {
                            int c = DataType.compare(t1.get(i), t2.get(i));
                            if (c != 0) {
                                if (!mAsc[i]) c *= -1;
                                return c;
                            }
                        } catch (ExecException e) {
                            throw new RuntimeException("Unable to compare tuples", e);
                        }
                    }
                    return 0;
                }
            }
            return -1; // keep the compiler happy
        }
    }

    private Comparator<Tuple> mComparator = new SortComparator();
    private FuncSpec mUserComparisonFuncSpec;
    private ComparisonFunc mUserComparisonFunc;
    
    
    @SuppressWarnings("unchecked")
    private void instantiateFunc() {
        if(mUserComparisonFunc != null) {
            this.mUserComparisonFunc = (ComparisonFunc) PigContext.instantiateFuncFromSpec(this.mUserComparisonFuncSpec);
            this.mUserComparisonFunc.setReporter(reporter);
            this.mComparator = mUserComparisonFunc;
        }
    }
    
    // We need to instantiate any user defined comparison function 
    // on the backend when the FindQuantiles udf is deserialized
    private void readObject(ObjectInputStream is) throws IOException, ClassNotFoundException{
        is.defaultReadObject();
        instantiateFunc();
    }
    

    public FindQuantiles() {
        mState = State.ALL_ASC;
    }

    public FindQuantiles(String[] args) {
        int startIndex = 0;
        int ascFlagsLength = args.length;
        // the first argument may be the information
        // about user defined comparison function if one
        // was specified
        if(args[0].startsWith(MRCompiler.USER_COMPARATOR_MARKER)) {
            mUserComparisonFuncSpec = new FuncSpec(
                    args[0].substring(MRCompiler.USER_COMPARATOR_MARKER.length()));
            // skip the first argument now that we used it
            startIndex++;
            ascFlagsLength--;
        }
        
        mAsc = new boolean[ascFlagsLength];
        boolean sawAsc = false;
        boolean sawDesc = false;
        for (int i = startIndex; i < ascFlagsLength; i++) {
            mAsc[i] = Boolean.parseBoolean(args[i]);
            if (mAsc[i]) sawAsc = true;
            else sawDesc = true;
        }
        if (sawAsc && sawDesc) mState = State.MIXED;
        else if (sawDesc) mState = State.ALL_DESC;
        else mState = State.ALL_ASC; // In cast they gave us no args this
                                     // defaults to all ascending.
    }

    /**
     * first field in the input tuple is the number of quantiles to generate
     * second field is the *sorted* bag of samples
     */
    
    @Override
    public Map<String, Object> exec(Tuple in) throws IOException {
        Map<String, Object> output = new HashMap<String, Object>();
        if(in==null || in.size()==0)
            return null;
        
        ArrayList<Tuple> quantilesList = new ArrayList<Tuple>();
        InternalMap weightedParts = new InternalMap();
        // the sample file has a tuple as under:
        // (numQuantiles, bag of samples) 
        // numQuantiles here is the reduce parallelism
        try{
            if (numQuantiles == null) {
                numQuantiles = (Integer)in.get(0);
            }
            if (samples == null) {
                samples = (DataBag)in.get(1);
            }
            long numSamples = samples.size();
            double toSkip = (double)numSamples / numQuantiles;
            if(toSkip < 1) {
                // numSamples is < numQuantiles;
                // set numQuantiles to numSamples
                numQuantiles = (int)numSamples;
                toSkip = 1;
            }
            
            long ind=0, j=-1;
            double nextQuantile = toSkip-1;
            for (Tuple it : samples) {
                if (ind==(long)nextQuantile){
                    ++j;
                    quantilesList.add(it);
                    nextQuantile+=toSkip;
                    if(j==numQuantiles-1)
                        break;
                }
                ind++;
                if (ind % 1000 == 0) progress();
            }
            long i=-1;
            Map<Tuple,CountingMap<Integer>> contribs = new HashMap<Tuple, CountingMap<Integer>>();
            for (Tuple it : samples){
                ++i;
                if (i % 1000 == 0) progress();
                int partInd = (int)(i/toSkip); // which partition
                if(partInd==numQuantiles) break;
                // the quantiles array has the element from the sample which is the
                // last element for a given partition. For example: if numQuantiles 
                // is 5 and number of samples is 100, then toSkip = 20 
                // quantiles[0] = sample[19] // the 20th element
                // quantiles[1] = sample[39] // the 40th element
                // and so on. For any element in the sample between 0 and 19, partInd
                // will be 0. We want to check if a sample element which is
                // present between 0 and 19 is also the 19th (quantiles[0] element).
                // This would mean that element might spread over the 0th and 1st 
                // partition. We are looking for contributions to a partition
                // from such elements. 
                
                // First We only check for sample elements in partitions other than the last one
                // < numQuantiles -1 (partInd is 0 indexed). 
                if(partInd<numQuantiles-1 && areEqual(it,quantilesList.get(partInd))){
                    if(!contribs.containsKey(it)){
                        CountingMap<Integer> cm = new CountingMap<Integer>();
                        cm.put(partInd, 1);
                        contribs.put(it, cm);
                    }
                    else
                        contribs.get(it).put(partInd, 1);
                }
                else{ 
                    // we are either in the last partition (last quantile)
                    // OR the sample element we are currently processing is not
                    // the same as the element in the quantile array for this partition
                    // if we haven't seen this sample item earlier, this is not an
                    // element which crosses partitions - so ignore
                    if(!contribs.containsKey(it))
                        continue;
                    else
                        // we have seen this sample before (in a previous partInd), 
                        // add to the contribution associated with this sample - if we had 
                        // not seen this sample in a previous partInd, then we would have not
                        // had this in the contribs map! (because of the if above).This 
                        // "key" (represented by the sample item) can either go to the 
                        // previous partInd or this partInd in the final sort reduce stage. 
                        // That is where the amount of contribution to each partInd will
                        // matter and influence the choice.
                        contribs.get(it).put(partInd, 1);
                }
            }
            int k = 0;
            for(Entry<Tuple, CountingMap<Integer>> ent : contribs.entrySet()){
                if (k % 1000 == 0) progress();
                Tuple key = ent.getKey(); // sample item which repeats
                
                // this map will have the contributions of the sample item to the different partitions
                CountingMap<Integer> value = ent.getValue(); 
                
                long total = value.getTotalCount();
                Tuple probVec =  mTupleFactory.newTuple(numQuantiles.intValue());
                // initialize all contribution fractions for different
                // partitions to 0.0
                for (int l = 0; l < numQuantiles; l++) {
                    probVec.set(l, new Float(0.0));
                }
                // for each partition that this sample item is present in,
                // compute the fraction of the total occurrences for that
                // partition - this will be the probability with which we
                // will pick this partition in the final sort reduce job
                // for this sample item
                for (Entry<Integer,Integer> valEnt : value.entrySet()) {
                    probVec.set(valEnt.getKey(), (float)valEnt.getValue()/total);
                }
                weightedParts.put(key, probVec);
            }
            output.put(QUANTILES_LIST, new NonSpillableDataBag(quantilesList));
            output.put(WEIGHTED_PARTS, weightedParts);
            return output;
        }catch (Exception e){
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    private boolean areEqual(Tuple it, Tuple tuple) {
        return mComparator.compare(it, tuple)==0;
    }
}