from skimage.transform import resize
from skimage.metrics import structural_similarity

from scipy.ndimage.filters import gaussian_filter as __gaussian_filter
from scipy.ndimage.filters import convolve as __convolve
from scipy.ndimage.filters import correlate as __correlate
from scipy.fftpack import fftshift as __fftshift

from skimage import img_as_ubyte, img_as_float
from skvideo.measure import msssim
import scipy.signal
import scipy.ndimage

from tqdm import tqdm
import numpy as np
import imageio
import torch 
import os

from dataclasses import dataclass


		
@dataclass
class Metrics:
	'''
		A monitor class keeps track of running averages for the PSNR, SSIM, MS-SSIM and VIF 
		between the original video and the animated video.
	'''
	total_psnr, avg_psnr = 0,0
	total_ssim, avg_ssim = 0,0
	total_vif, avg_vif = 0,0
	avg_ms_ssim = 0
	
	def update_metrics(self, original_frame: np.array, animated_frame: np.array, idx: int) -> None:
		psnr_val = psnr(img_as_ubyte(original_frame), img_as_ubyte(animated_frame))
		ssim_val = ssim(img_as_ubyte(original_frame), img_as_ubyte(animated_frame))
		vif_val = vif(img_as_ubyte(original_frame), img_as_ubyte(animated_frame))
		
		self.total_psnr += psnr_val
		self.total_ssim += ssim_val
		self.total_vif += vif_val
		
		if idx == 0:
			self.avg_psnr = self.total_psnr
			self.avg_ssim = self.total_ssim
			self.avg_vif = self.total_vif
		else:
			self.avg_psnr = self.total_psnr/(idx+1)
			self.avg_ssim = self.total_ssim/(idx+1)
			self.avg_vif = self.total_vif/(idx+1)
			
	def update_average_quality(self, frame_idx: int) -> None:
		self.avg_psnr = self.total_psnr/(frame_idx+1)
		self.avg_ssim = self.total_ssim/(frame_idx+1)
		self.avg_vif = self.total_vif/(frame_idx+1)
		
	def compute_ms_ssim(self, original_video: list, animated_video: list) -> float:
		ms_ssim = mssim( [img_as_ubyte(frame) for frame in original_video], [img_as_ubyte(frame) for frame in animated_video], seq=True)
		self.avg_ms_ssim = np.mean(ms_ssim)
		return self.avg_ms_ssim
		
		

def mse(original, decoded):
	'''Mean squared error'''
	mse = np.mean((img_as_ubyte(original)-img_as_ubyte(decoded))**2)
	return mse
		
def psnr(original, decoded, seq=False):
	'''Peak signal to noise ratio'''
	if seq:
		seq_psnr = []
		for idx in range(len(decoded)):
			error_img = original[idx]-decoded[idx]
			mse_val_r = mse(original[idx][:,:,0], decoded[idx][:,:,0])
			mse_val_g = mse(original[idx][:,:,1], decoded[idx][:,:,1])
			mse_val_b = mse(original[idx][:,:,2], decoded[idx][:,:,2])
			mse_val = (mse_val_r+mse_val_g+mse_val_b)/3
			psnr_val = 10*np.log10(255**2/mse_val)			
			seq_psnr.append(psnr_val)
		return seq_psnr
					
	else:
		mse_val = mse(original[:,:,0], decoded[:,:,0])
		mse_val_r = mse(original[:,:,0], decoded[:,:,0])
		mse_val_g = mse(original[:,:,1], decoded[:,:,1])
		mse_val_b = mse(original[:,:,2], decoded[:,:,2])
		mse_val = (mse_val_r+mse_val_g+mse_val_b)/3
		psnr_val = 10*np.log10(255**2/mse_val)
		return psnr_val
		
def ssim(original, decoded, seq=False):
	'''structual similarity index'''
	if seq:
		seq_ssim = []
		for idx in range(len(decoded)):
			ref = img_as_ubyte(original[idx][:,:,0])
			dec = img_as_ubyte(decoded[idx][:,:,0])
			ssim_val = structural_similarity(ref, dec, data_range = np.amax(dec)-np.amin(dec), size_average=False)
			seq_ssim.append(ssim_val)
		return seq_ssim
	else:
		ref = img_as_ubyte(original[:,:,0])
		dec = img_as_ubyte(decoded[:,:,0])
		ssim_val = structural_similarity(ref, dec, data_range = np.amax(dec)-np.amin(dec), size_average=False)
		return ssim_val
		

		
def mssim(original, decoded, seq=False):
	'''Multiscale structual similarity index'''
	if seq:
		frames = len(decoded)
		org = [frame[:,:,0] for frame in original[0:frames]]
		dec = [frame[:,:,0] for frame in decoded]
		seq_ssim = msssim(org, dec)
		seq_ssim = [x.astype(float) for x in seq_ssim]
		return seq_ssim
	else:
		val = msssim(original[:,:,0],decoded[:,:,0])
		return val.astype(float)
		
def vif(org=None, dec=None, seq=None):
	'''Multi-scale Visual information fidelity'''	    
	if seq:
		s_vif = []
		ref = [img_as_ubyte(frame) for frame in org]
		dec = [img_as_ubyte(frame) for frame in dec]
		for idx in tqdm(range(len(dec))):
			vif = mscale(ref[idx],dec[idx])
			s_vif.append(vif)
		return s_vif
	else:
		original = img_as_ubyte(org)
		decoded = img_as_ubyte(dec)
		vif = mscale(original, decoded)
		return vif
		
def mscale(ref, dist):
	sigma_nsq=4
	eps = 1e-10
	num = 0.0
	den = 0.0
	for scale in range(1, 5):
		N = 2**(4-scale+1) + 1
		sd = N/5.0
		if (scale > 1):
			ref = scipy.ndimage.gaussian_filter(ref, sd)
			dist = scipy.ndimage.gaussian_filter(dist, sd)
			ref = ref[::2, ::2]
			dist = dist[::2, ::2]        
		mu1 = scipy.ndimage.gaussian_filter(ref, sd)
		mu2 = scipy.ndimage.gaussian_filter(dist, sd)
		mu1_sq = mu1 * mu1
		mu2_sq = mu2 * mu2
		mu1_mu2 = mu1 * mu2
		sigma1_sq = scipy.ndimage.gaussian_filter(ref * ref, sd) - mu1_sq
		sigma2_sq = scipy.ndimage.gaussian_filter(dist * dist, sd) - mu2_sq
		sigma12 = scipy.ndimage.gaussian_filter(ref * dist, sd) - mu1_mu2
       
		sigma1_sq[sigma1_sq<0] = 0
		sigma2_sq[sigma2_sq<0] = 0
       	
		g = sigma12 / (sigma1_sq + eps)
		sv_sq = sigma2_sq - g * sigma12
		g[sigma1_sq<eps] = 0
		sv_sq[sigma1_sq<eps] = sigma2_sq[sigma1_sq<eps]
		sigma1_sq[sigma1_sq<eps] = 0
       
		g[sigma2_sq<eps] = 0
		sv_sq[sigma2_sq<eps] = 0
		sv_sq[g<0] = sigma2_sq[g<0]
		g[g<0] = 0
		sv_sq[sv_sq<=eps] = eps
       	
		num += np.sum(np.log10(1 + g * g * sigma1_sq / (sv_sq + sigma_nsq)))
		den += np.sum(np.log10(1 + sigma1_sq / sigma_nsq))
	vifp = num/den
	if np.isnan(vifp):
		return 1.0
	else:
		return vifp
		

if __name__ == "__main__":
	monitor = Metrics()
