# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE_inria.md file.
#
# For inquiries contact george.drettakis@inria.fr

from pathlib import Path
import os
from PIL import Image
import torch
import torchvision.transforms.functional as tf
from utils.loss_utils import ssim
from lpipsPyTorch import lpips
import json
from tqdm import tqdm
from utils.image_utils import psnr
from argparse import ArgumentParser

def readImages(renders_dir, gt_dir):
 renders = []
 gts = []
 image_names = []
 for fname in os.listdir(renders_dir):
 render = Image.open(renders_dir / fname)
 gt = Image.open(gt_dir / fname)
 renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda())
 gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda())
 image_names.append(fname)
 return renders, gts, image_names

def evaluate(model_paths):

 full_dict = {}
 per_view_dict = {}
 full_dict_polytopeonly = {}
 per_view_dict_polytopeonly = {}
 print("")

 for scene_dir in model_paths:
 try:
 print("Scene:", scene_dir)
 full_dict[scene_dir] = {}
 per_view_dict[scene_dir] = {}
 full_dict_polytopeonly[scene_dir] = {}
 per_view_dict_polytopeonly[scene_dir] = {}

 test_dir = Path(scene_dir) / "test"

 for method in os.listdir(test_dir):
 print("Method:", method)

 full_dict[scene_dir][method] = {}
 per_view_dict[scene_dir][method] = {}
 full_dict_polytopeonly[scene_dir][method] = {}
 per_view_dict_polytopeonly[scene_dir][method] = {}

 method_dir = test_dir / method
 gt_dir = method_dir/ "gt"
 renders_dir = method_dir / "renders"
 renders, gts, image_names = readImages(renders_dir, gt_dir)

 ssims = []
 psnrs = []
 lpipss = []

 for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"):
 ssims.append(ssim(renders[idx], gts[idx]))
 psnrs.append(psnr(renders[idx], gts[idx]))
 lpipss.append(lpips(renders[idx], gts[idx], net_type='vgg'))

 print(" SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5"))
 print(" PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5"))
 print(" LPIPS: {:>12.7f}".format(torch.tensor(lpipss).mean(), ".5"))
 print("")

 full_dict[scene_dir][method].update({"SSIM": torch.tensor(ssims).mean().item(),
 "PSNR": torch.tensor(psnrs).mean().item(),
 "LPIPS": torch.tensor(lpipss).mean().item()})
 per_view_dict[scene_dir][method].update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)},
 "PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)},
 "LPIPS": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}})

 with open(scene_dir + "/results.json", 'w') as fp:
 json.dump(full_dict[scene_dir], fp, indent=True)
 with open(scene_dir + "/per_view.json", 'w') as fp:
 json.dump(per_view_dict[scene_dir], fp, indent=True)
 except:
 print("Unable to compute metrics for model", scene_dir)

if __name__ == "__main__":
 device = torch.device("cuda:0")
 torch.cuda.set_device(device)

 # Set up command line argument parser
 parser = ArgumentParser(description="Training script parameters")
 parser.add_argument('--model_paths', '-m', required=True, nargs="+", type=str, default=[])
 args = parser.parse_args()
 evaluate(args.model_paths)