import pandas as pd
from data import *
from config import CFG


def add_imageid():
    """
    modify the caption and image name file with assigning id to each image which has
    more than one caption.
    """
    df = pd.read_csv(CFG.dataset_path+"/results.csv", delimiter="|")
    df.columns = ['image', 'caption_number', 'caption']
    df['caption'] = df['caption'].str.lstrip()
    df['caption_number'] = df['caption_number'].str.lstrip()
    df.loc[19999, 'caption_number'] = "4"
    df.loc[19999, 'caption'] = "A dog runs across the grass ."
    ids = [id_ for id_ in range(len(df) // 5) for i in range(5)]
    df['id'] = ids
    df.to_csv(CFG.captions_path+"captions.csv", index=False)

class AvgMeter:
    """
    Helps in storing average of loss across a batch/epoch 
    """
    def __init__(self, name="Metric"):
        self.name = name
        self.reset()

    def reset(self):
        self.avg, self.sum, self.count = [0] * 3

    def update(self, val, count=1):
        self.count += count
        self.sum += val * count
        self.avg = self.sum / self.count

    def __repr__(self):
        text = f"{self.name}: {self.avg:.4f}"
        return text

def get_lr(optimizer):
    """
    Return paramater gourp for lr in oprimiser.
    """
    for param_group in optimizer.param_groups:
        return param_group["lr"]
