# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import time


def human_format(num):
    magnitude = 0
    while num >= 1000:
        magnitude += 1
        num /= 1000.0
    # add more suffixes if you need them
    return '{:3}{}'.format(num, [' ', 'K', 'M', 'G', 'T', 'P'][magnitude])


def cal_hit(pred, label, pad_id=None):
    bz = pred.size(0)
    if pad_id is None:
        return (pred.argmax(dim=-1) == label).sum().item()

    pred = pred.argmax(dim=-1).tolist()
    label = label.tolist()
    hit = 0
    for batch_p, batch_l in zip(pred, label):
        for p, l in zip(batch_p, batch_l):
            if p == l and p != pad_id:
                hit += 1
    return hit


class Timer():
    ''' Timer for recording training time distribution. '''
    def __init__(self):
        self.prev_t = time.time()
        self.clear()

    def set(self):
        self.prev_t = time.time()

    def cnt(self,mode):
        self.time_table[mode] += time.time()-self.prev_t
        self.set()
        if mode =='bw':
            self.click += 1

    def show(self):
        total_time = sum(self.time_table.values())
        self.time_table['avg'] = total_time/self.click
        self.time_table['rd'] = 100*self.time_table['rd']/total_time
        self.time_table['fw'] = 100*self.time_table['fw']/total_time
        self.time_table['bw'] = 100*self.time_table['bw']/total_time
        msg  = '{avg:.3f} sec/step (read_data {rd:.1f}% | forward {fw:.1f}% | backward {bw:.1f}%)'.format(**self.time_table)
        self.clear()
        return msg

    def clear(self):
        self.time_table = {'rd':0,'fw':0,'bw':0}
        self.click = 0

