import sys
import tqdm
import torch
import argparse
import numpy as np
from omegaconf import OmegaConf

sys.path.append('.')
from utils.build_utils import build_from_cfg
from datasets.gopro_datasets import GoPro_Test_Dataset
from metrics.psnr_ssim import calculate_psnr, calculate_ssim

parser = argparse.ArgumentParser(
                prog = 'AMT',
                description = 'GOPRO evaluation',
                )
parser.add_argument('-c', '--config', default='cfgs/AMT-S_gopro.yaml') 
parser.add_argument('-p', '--ckpt', default='pretrained/gopro_amt-s.pth',) 
parser.add_argument('-r', '--root', default='data/GOPRO',) 
args = parser.parse_args()

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
cfg_path = args.config
ckpt_path = args.ckpt
root = args.root

network_cfg = OmegaConf.load(cfg_path).network
network_name = network_cfg.name
model = build_from_cfg(network_cfg)
ckpt = torch.load(ckpt_path)
model.load_state_dict(ckpt['state_dict'])
model = model.to(device)
model.eval()

dataset = GoPro_Test_Dataset(dataset_dir=root)

psnr_list = []
ssim_list = []
pbar = tqdm.tqdm(dataset, total=len(dataset))
for data in pbar:
    input_dict = {}
    for k, v in data.items():
        input_dict[k] = v.to(device).unsqueeze(0)
    with torch.no_grad():
        imgt_pred = model(**input_dict)['imgt_pred']
        psnr = calculate_psnr(imgt_pred, input_dict['imgt'])
        ssim = calculate_ssim(imgt_pred, input_dict['imgt'])
    psnr_list.append(psnr)
    ssim_list.append(ssim)
    avg_psnr = np.mean(psnr_list)
    avg_ssim = np.mean(ssim_list)
    desc_str = f'[{network_name}/GOPRO] psnr: {avg_psnr:.02f}, ssim: {avg_ssim:.04f}'
    pbar.set_description_str(desc_str)