from functools import partial

import torch
import skimage
import lpips as lpips_lib


class SSIM:
    def __init__(self):
        self.criterion = partial(skimage.metrics.structural_similarity, data_range=1.0, channel_axis=2)
        #     data_range=1.0, channel_axis=2, gaussian_weights=True, sigma=1.5, use_sample_covariance=False)

    def __call__(self, pred, gt):
        ssim = self.criterion(
            pred.permute(1, 2, 0).cpu().numpy(),
            gt.permute(1, 2, 0).cpu().numpy(),
        )
        return ssim

class PSNR:
    def __init__(self):
        self.criterion = partial(skimage.metrics.peak_signal_noise_ratio, data_range=1.0)

    def __call__(self, pred, gt):
        # return -10.0 * torch.log10(torch.mean(torch.square(gt.permute(1, 2, 0) - pred.permute(1, 2, 0)))).item()
        psnr = self.criterion(
            pred.permute(1, 2, 0).cpu().numpy(),
            gt.permute(1, 2, 0).cpu().numpy(),
        )
        return psnr

class LPIPS:
    def __init__(self):
        self.criterion = lpips_lib.LPIPS(net='vgg', verbose=False).cuda()

    def __call__(self, pred, gt):
        pred = pred.float()
        gt = gt.float()
        self.criterion = self.criterion.to(pred.device).to(pred.dtype)
        val = self.criterion(pred * 2 - 1, gt * 2 - 1)
        return val


class RGBEvaluation:
    def __init__(self, cfg):

        self.metrics = ('PSNR', 'SSIM', 'LPIPS')

        self.ssim = SSIM()
        self.psnr = PSNR()
        self.lpips = LPIPS()

        self.crop_edges = getattr(cfg, 'crop_edges', None)

        if cfg.has('resize'):
            from vidar.utils.tensor import interpolate
            self.resize = partial(interpolate, mode='bilinear', size=cfg.resize)
        else:
            self.resize = None

    def compute(self, gt, pred):
        # For each batch sample
        metrics = []
        for pred_i, gt_i in zip(pred, gt):

            gt_i = gt_i.unsqueeze(0).clone().to(torch.float64)
            pred_i = pred_i.unsqueeze(0).clone().to(torch.float64)
            # pred_i = self.interp_bilinear(pred_i, gt_i)

            if self.resize is not None:
                gt_i = self.resize(gt_i)
                pred_i = self.resize(pred_i)

            gt_i = gt_i.clamp(min=0.0, max=1.0)
            pred_i = pred_i.clamp(min=0.0, max=1.0)

            if self.crop_edges:
                h, w = gt_i.shape[-2:]
                crop_h = int(self.crop_edges * h)
                crop_w = int(self.crop_edges * w)
                if crop_h > 0 and crop_w > 0:
                    gt_i = gt_i[:, :, crop_h:-crop_h, crop_w:-crop_w]
                    pred_i = pred_i[:, :, crop_h:-crop_h, crop_w:-crop_w]

            ssim = self.ssim(pred_i[0], gt_i[0])
            psnr = self.psnr(pred_i[0], gt_i[0])
            lpips_val = self.lpips(pred_i[0], gt_i[0])

            if gt_i.sum() < 0.01:
                psnr, ssim, lpips_val = -999, -999, -999

            metrics.append([psnr, ssim, lpips_val])

        # Return metrics
        return torch.tensor(metrics, dtype=gt.dtype, device=gt.device)

    def evaluate(self, gt, pred):
        return self.compute(gt, pred)