from src.methods.brice_convers.dataHandler import DataHandler
from src.utils import psnr, ssim
from sklearn.metrics import f1_score, mean_squared_error
import numpy as np

class DataEvaluation:
    def __init__(self, DataHandler: DataHandler):
        DataEvaluation.DataHandler = DataHandler

    def print_metrics(self, indexImage, method):
        DataEvaluation.DataHandler.indexImageExists(indexImage)

        img = DataEvaluation.DataHandler.load_image(indexImage)
        res = DataEvaluation.DataHandler.get_reconstructed_image(indexImage, method)

        ssimMetric = ssim(img, res)
        psnrMetrc = psnr(img, res)
        mse = mean_squared_error(img.flatten(), res.flatten())
        mseRedPixels = mean_squared_error(img[:,:,0], res[:,:,0])
        mseGreenPixels = mean_squared_error(img[:,:,1], res[:,:,1])
        mseBluePixels = mean_squared_error(img[:,:,2], res[:,:,2])
        miMetric = DataEvaluation.MI(img, res)
        ccMetric = DataEvaluation.CC(img, res)
        sadMetric = DataEvaluation.SAD(img, res)
        lsMetric = DataEvaluation.LS(img, res)

        print("[INFO] Metrics for image {}".format(indexImage))
        print("#" * 30)
        print(" SSIM: {:.6}  ".format(ssimMetric))
        print(" PSNR: {:.6}   ".format(psnrMetrc))
        print(" MSE : {:.3e}    ".format(mse))
        print(" MSE (R): {:.3e}    ".format(mseRedPixels))
        print(" MSE (G): {:.3e}    ".format(mseGreenPixels))
        print(" MSE (B): {:.3e}    ".format(mseBluePixels))
        print(" MI: {:.6}     ".format(miMetric))
        print(" CC: {:.6}    ".format(ccMetric))
        print(" SAD: {:.6} ".format(sadMetric))
        print(" LS: {:.3e} ".format(lsMetric))
        print("#" * 30)

    #Mutual Information
    def MI(img_mov, img_ref):
        hgram, x_edges, y_edges = np.histogram2d(img_mov.ravel(), img_ref.ravel(), bins=20)
        pxy = hgram / float(np.sum(hgram))
        px = np.sum(pxy, axis=1) # marginal for x over y
        py = np.sum(pxy, axis=0) # marginal for y over x
        px_py = px[:, None] * py[None, :] # Broadcast to multiply marginals
        # Now we can do the calculation using the pxy, px_py 2D arrays
        nzs = pxy > 0 # Only non-zero pxy values contribute to the sum
        return np.sum(pxy[nzs] * np.log(pxy[nzs] / px_py[nzs]))

    # Cross Correlation
    def CC(img_mov, img_ref):
        # Vectorized versions of c,d,e
        a = img_mov.astype('float64')
        b = img_ref.astype('float64')

        # Calculating mean values
        AM = np.mean(a)
        BM = np.mean(b)

        c_vect = (a - AM) * (b - BM)
        d_vect = (a - AM) ** 2
        e_vect = (b - BM) ** 2

        # Finally get r using those vectorized versions
        r_out = np.sum(c_vect) / float(np.sqrt(np.sum(d_vect) * np.sum(e_vect)))
        return r_out

    #Sum of Absolute Differences
    def SAD(img_mov, img_ref):
        img1 = img_mov.astype('float64')
        img2 = img_ref.astype('float64')
        ab = np.abs(img1 - img2)
        sav = np.sum(ab.ravel())
        sav /= ab.ravel().shape[0]
        return sav

    #Sum of Least Squared Errors
    def LS(img_mov, img_ref):
        img1 = img_mov.astype('float64')
        img2 = img_ref.astype('float64')
        r = (img1 - img2)**2
        sse = np.sum(r.ravel())
        sse /= r.ravel().shape[0]
        return sse