import caffe2.contrib.playground.meter as Meter
from caffe2.python import workspace


class ComputeLoss(Meter.Meter):
    def __init__(self, opts=None, blob_name=''):
        self.blob_name = blob_name
        self.opts = opts
        self.iter = 0
        self.value = 0

    def Reset(self):
        self.iter = 0
        self.value = 0

    def Add(self):
        """Average values of a blob on each gpu"""
        value = 0
        for idx in range(self.opts['distributed']['first_xpu_id'],
                         self.opts['distributed']['first_xpu_id'] +
                         self.opts['distributed']['num_xpus']):
            value += workspace.FetchBlob('{}_{}/{}'.
                format(self.opts['distributed']['device'], idx, self.blob_name))
        self.value += value
        self.iter += 1

    def Compute(self):
        result = self.opts['distributed']['num_shards'] * self.value / self.iter
        self.Reset()
        return result