# Tai Sakuma <tai.sakuma@cern.ch>

##__________________________________________________________________||
import numpy as np

##__________________________________________________________________||
class Sum(object):
    def __init__(self, initial_val = (0, )):
        self._results = { }
        self.initial_val = initial_val

    def add(self, key, val, weight = 1):
        self.add_key(key)
        self._results[key] = self._results[key] + np.array(val*weight)

    def add_key(self, key):
        if key not in self._results:
            self._results[key] = np.array(self.initial_val)

    def keys(self):
        return self._results.keys()

    def copy_from(self, src):
        self._results.clear()
        self._results.update(src._results)

    def results(self):
        return self._results

    def __add__(self, other):
        ret = Sum()
        results = {k: v.copy() for k, v in self._results.iteritems()}
        if not other == 0: # other is 0 when e.g. sum([obj1, obj2])
            self._add_results_inplace(results, other._results)
        ret._results.clear()
        ret._results.update(results)
        return ret

    def __iadd__(self, other):
        self._add_results_inplace(self._results, other._results)
        return self

    def __radd__(self, other):
        return self.__add__(other)

    def _add_results_inplace(self, res1, res2):
        # res1 += res2, modify res1
        for k, v in res2.iteritems():
            if k not in self._results:
                res1[k] = v.copy()
            else:
                res1[k] = self._results[k] + v

##__________________________________________________________________||

@@ -119,16 +119,383 @@
  as np%0A%0A
+class SumImp(object):%0A%0A    def __init__(self, val = None, weight = 1, contents = None):%0A%0A        if contents is not None:%0A            self.contents = np.array(contents)%0A            return%0A%0A        self.contents = np.array(val*weight)%0A%0A    def __add__(self, other):%0A        contents = self.contents + other.contents%0A        return self.__class__(contents = contents)%0A%0A
 ##______
@@ -662,27 +662,62 @@
    self.
-initial_val
+Summary = SumImp%0A        self.initial_contents
  = initi
@@ -841,29 +841,34 @@
 %5Bkey%5D + 
-np.arra
+self.Summar
 y(val
-*
+, 
 weight)%0A
@@ -966,25 +966,40 @@
 %5Bkey%5D = 
-np.array(
+self.Summary(contents = 
 self.ini
@@ -1003,19 +1003,24 @@
 initial_
-val
+contents
 )%0A%0A    d
@@ -1197,32 +1197,35 @@
 s(self):%0A       
+ ##
  return self._re
@@ -1229,16 +1229,85 @@
 _results
+%0A        return %7Bk: v.contents for k, v in self._results.iteritems()%7D
 %0A%0A    de
@@ -1344,19 +1344,30 @@
   ret = 
-Sum
+self.__class__
 ()%0A     
@@ -1383,23 +1383,31 @@
 s = %7Bk: 
-v.
+copy.deep
 copy(
+v
 ) for k,
@@ -2031,15 +2031,23 @@
 %5D = 
-v.
+copy.deep
 copy(
+v
 )%0A  

make Sum similar to Count