package shaman.hive.udf;

import java.util.Map;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.parse.SemanticException;
import org.apache.hadoop.hive.ql.udf.generic.AbstractGenericUDAFResolver;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator.AggregationBuffer;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFParameterInfo;
import org.apache.hadoop.hive.serde2.objectinspector.MapObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorUtils;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorUtils.ObjectInspectorCopyOption;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.Text;

import com.google.common.collect.Maps;

@Description(value="_FUNC_ calculate the cdn bytes loaded", name="shaman_cdnbytesloaded")
public class GenericUDAFCdnBytesLoaded extends AbstractGenericUDAFResolver {

    @Override
    public GenericUDAFEvaluator getEvaluator(TypeInfo[] parameters)
            throws SemanticException {
        throw new SemanticException(
              "This UDAF does not support the deprecated getEvaluator() method.");
    }

    @Override
    public GenericUDAFEvaluator getEvaluator(GenericUDAFParameterInfo info)
            throws SemanticException {
        ObjectInspector[] ois = info.getParameterObjectInspectors();
        if (ois.length != 1) {
            throw new UDFArgumentLengthException("_FUNC_ take only one parameter");
        }
        return new GenericUDAFCdnBytesLoadedEvaluator();
    }

    public static class CdnAggregationBuffer implements AggregationBuffer {
        Map<String, Long> bytesLoaded;
    }
    
    public static class GenericUDAFCdnBytesLoadedEvaluator extends GenericUDAFEvaluator {
        
        private MapObjectInspector inputOI;
        private MapObjectInspector outputOI;
        private Map<String, String> cdnPatterns;
        
        {
            cdnPatterns = Maps.newHashMap();
            /* Do we really need to convert the CDN info here?
            cdnPatterns.put("cdnl3nl", "Level3");
            cdnPatterns.put("cdnak", "Akamai");
            cdnPatterns.put("cdnllnwnl", "LimeLight");
            cdnPatterns.put("cdncd", "CDNetworks");
            */
        }
        
        @Override
        public ObjectInspector init(Mode m, ObjectInspector[] parameters) throws HiveException {
            super.init(m, parameters);
            if (parameters.length != 1) {
                throw new UDFArgumentLengthException("In GenericUDAFCdnBytesLoadedEvaluator "
                        + "parameters length is not 1");
            }
            inputOI = (MapObjectInspector) parameters[0];
            System.out.println(inputOI.getTypeName());
            outputOI = (MapObjectInspector) ObjectInspectorUtils
                    .getStandardObjectInspector(inputOI, ObjectInspectorCopyOption.JAVA);
            System.out.println(outputOI.getTypeName());
            return outputOI;
        }
        
        @Override
        public AggregationBuffer getNewAggregationBuffer() throws HiveException {
            CdnAggregationBuffer buffer = new CdnAggregationBuffer();
            buffer.bytesLoaded = Maps.newHashMap();
            return buffer;
        }

        @Override
        public void reset(AggregationBuffer agg) throws HiveException {
            CdnAggregationBuffer buffer = (CdnAggregationBuffer) agg;
            buffer.bytesLoaded.clear();
        }

        @Override
        public void iterate(AggregationBuffer agg, Object[] parameters)
                throws HiveException {
            if (parameters.length == 1 && parameters[0] != null) {
                try {
                    CdnAggregationBuffer buffer = (CdnAggregationBuffer) agg;
                    @SuppressWarnings("unchecked")
                    Map<String, Long> bytesLoaded = (Map<String, Long>) inputOI.getMap(parameters[0]);
                    for (String key : bytesLoaded.keySet()) {
                        boolean matched = false;
                        for (String pattern : cdnPatterns.keySet()) {
                            if (key.contains(pattern)) {
                                String cdnName = cdnPatterns.get(pattern);
                                Long loaded = buffer.bytesLoaded.get(cdnName);
                                if (loaded == null) {
                                    loaded = bytesLoaded.get(key);
                                } else {
                                    loaded = loaded + bytesLoaded.get(key);
                                }
                                buffer.bytesLoaded.put(cdnName, loaded);
                                matched = true;
                                break;
                            }
                        }
                        if (matched == false) {
                            Long loaded = buffer.bytesLoaded.get(key);
                            if (loaded == null) {
                                loaded = bytesLoaded.get(key);
                            } else {
                                loaded = loaded + bytesLoaded.get(key);
                            }
                            buffer.bytesLoaded.put(key, loaded);
                        }
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }

        @Override
        public Object terminatePartial(AggregationBuffer agg) throws HiveException {
            return ((CdnAggregationBuffer) agg).bytesLoaded;
        }

        /**
         * merge is called at reduce phase.
         * But why inputOI.getMap change to return Map<Text, LongWritable>
         */
        @Override
        public void merge(AggregationBuffer agg, Object partial) throws HiveException {
            if (partial == null) {
                return;
            }
            try {
                CdnAggregationBuffer buffer = (CdnAggregationBuffer) agg;
                @SuppressWarnings("unchecked")
                Map<Text, LongWritable> bytesLoaded = (Map<Text, LongWritable>) inputOI.getMap(partial);
                for (Text key : bytesLoaded.keySet()) {
                    boolean matched = false;
                    for (String pattern : cdnPatterns.keySet()) {
                        if (key.toString().contains(pattern)) {
                            String cdnName = cdnPatterns.get(pattern);
                            Long loaded = buffer.bytesLoaded.get(cdnName);
                            if (loaded == null) {
                                loaded = bytesLoaded.get(key).get();
                            } else {
                                loaded = loaded + bytesLoaded.get(key).get();
                            }
                            buffer.bytesLoaded.put(cdnName, loaded);
                            matched = true;
                            break;
                        }
                    }
                    if (matched == false) {
                        Long loaded = buffer.bytesLoaded.get(key);
                        if (loaded == null) {
                            loaded = bytesLoaded.get(key).get();
                        } else {
                            loaded = loaded + bytesLoaded.get(key).get();
                        }
                        buffer.bytesLoaded.put(key.toString(), loaded);
                    }
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }

        @Override
        public Object terminate(AggregationBuffer agg) throws HiveException {
            return ((CdnAggregationBuffer) agg).bytesLoaded;
        }
        
    }
}

@@ -4303,32 +4303,196 @@
 ded.keySet()) %7B%0A
+                        Long bytes = bytesLoaded.get(key);%0A                        if (bytes %3C 0) %7B%0A                            continue;%0A                        %7D%0A
                 
@@ -4905,39 +4905,24 @@
 aded = bytes
-Loaded.get(key)
 ;%0A          
@@ -5007,39 +5007,24 @@
 aded + bytes
-Loaded.get(key)
 ;%0A          
@@ -5474,39 +5474,24 @@
 aded = bytes
-Loaded.get(key)
 ;%0A          
@@ -5568,39 +5568,24 @@
 aded + bytes
-Loaded.get(key)
 ;%0A          

fix negative issue