"""`metrics` module contains useful functions for image quality assessment."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/P4_HPD_cg/91_metrics.ipynb.

# %% auto 0
__all__ = ['epsilon', 'mse', 'rmse', 'nmse', 'nrmse', 'psnr', 'ssim', 'print_metrics']

# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 2
from nbdev.showdoc import *
from fastcore.test import *
from fastcore.utils import *

# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 4
import numpy as np
from prettytable import PrettyTable

# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 5
from skimage.metrics import peak_signal_noise_ratio, structural_similarity
from typing import Optional

# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 6
epsilon = 1e-10

# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 8
def mse(
    gt: np.ndarray,     # first image (ground truth)
    pred: np.ndarray,   # second image (predicted) 
) -> np.ndarray:        # MSE (mean squared error)
    "Mean Squared Error (MSE). Broadcasting rules will be applied if different dimensions."
    gt, pred = np.squeeze(gt), np.squeeze(pred)
    if gt.ndim != pred.ndim: print("Warning: gt's and pred's dims are mismatched!")
    if gt.ndim <2: print("Warning: gt's dims < 2")
    if pred.ndim <2: print("Warning: pred's dims < 2")
    return np.array(np.mean((gt - pred) ** 2))

# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 12
def rmse(
    gt: np.ndarray,     # first image (ground truth)
    pred: np.ndarray,   # second image (predicted) 
) -> np.ndarray:        # RMSE (root-mean-square error)
    "Root-Mean-Square Error (RMSE). Broadcasting rules will be applied if different dimensions."
    gt, pred = np.squeeze(gt), np.squeeze(pred)
    if gt.ndim != pred.ndim: print("Warning: gt's and pred's dims are mismatched!")
    if gt.ndim <2: print("Warning: gt's dims < 2")
    if pred.ndim <2: print("Warning: pred's dims < 2")
    return np.array(np.sqrt(np.mean((gt - pred) ** 2)))

# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 16
def nmse(
    gt: np.ndarray,     # first image (ground truth)
    pred: np.ndarray,   # second image (predicted) 
) -> np.ndarray:        # NMSE (normalized mean-square error)
    "Normalized Mean-Square Error (NMSE). Broadcasting rules will be applied if different dimensions."
    gt, pred = np.squeeze(gt), np.squeeze(pred)
    if gt.ndim != pred.ndim: print("Warning: gt's and pred's dims are mismatched!")
    if gt.ndim <2: print("Warning: gt's dims < 2")
    if pred.ndim <2: print("Warning: pred's dims < 2")
    return np.array(np.linalg.norm(gt - pred) ** 2 / np.linalg.norm(gt) ** 2)


# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 20
def nrmse(
    gt: np.ndarray,     # first image (ground truth)
    pred: np.ndarray,   # second image (predicted) 
) -> np.ndarray:        # NRMSE
    "Normalized Root-Mean-Square Error (NRMSE). Broadcasting rules will be applied if different dimensions."
    gt, pred = np.squeeze(gt), np.squeeze(pred)
    if gt.ndim != pred.ndim: print("Warning: gt's and pred's dims are mismatched!")
    if gt.ndim <2: print("Warning: gt's dims < 2")
    if pred.ndim <2: print("Warning: pred's dims < 2")
    return np.array(np.sqrt(np.linalg.norm(gt - pred) ** 2 / np.linalg.norm(gt) ** 2))


# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 25
def psnr(
    gt: np.ndarray,     # first image (ground truth)
    pred: np.ndarray,   # second image (predicted)
    drange:float=None   # distance between min & max possible values
) -> np.ndarray:        # PSNR (peak signal-to-noise ratio)
    "Peak Signal-to-Noise Ratio (PSNR)"
    gt, pred = np.squeeze(gt), np.squeeze(pred)
    if gt.ndim != pred.ndim: print("Warning: gt's and pred's dims are mismatched!")
    if gt.ndim <2: print("Warning: gt's dims < 2")
    if pred.ndim <2: print("Warning: pred's dims < 2")
    if drange is None:
        drange = gt.max() - gt.min()
    return np.array(peak_signal_noise_ratio(gt, pred, data_range=drange))

# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 28
def ssim(
    gt: np.ndarray,     # first image (ground truth)
    pred: np.ndarray,   # second image (predicted)
    drange: float=None  # distance between min & max values
) -> np.ndarray:        # SSIM
    "Structural Similarity Index Metric (SSIM)"
    gt, pred = np.squeeze(gt), np.squeeze(pred)
    if not gt.ndim == pred.ndim:
        raise ValueError("Ground truth dimensions does not match pred.")
    if not gt.ndim == 2:
        raise ValueError("Ground truth dimensions has to be 2.")
    drange = (gt.max()-gt.min()) if drange is None else drange
    return np.array(structural_similarity(gt, pred, data_range=drange))

# %% ../nbs/P4_HPD_cg/91_metrics.ipynb 31
def print_metrics(
    gt: np.ndarray,     # first image (ground truth)
    pred: np.ndarray,   # second image (predicted)
):
    "Print MSE, NMSE, RMSE, NRMSE, PSNR, and SSIM"
    gt, pred = np.squeeze(gt), np.squeeze(pred)
    if not gt.ndim == pred.ndim:
        raise ValueError("Ground truth dimensions does not match pred.")
    if not gt.ndim == 2:
        raise ValueError("Ground truth dimensions has to be 2.")
    
    # set the table up
    ptab = PrettyTable(['Metrics', 'Values'])
    # add rows
    ptab.add_row(['MSE', f"{mse(gt, pred):1.3e}"])
    ptab.add_row(['NMSE', f"{nmse(gt, pred):1.3e}"])
    ptab.add_row(['RMSE', f"{rmse(gt, pred):1.3e}"])
    ptab.add_row(['NRMSE', f"{nrmse(gt, pred):1.3e}"])
    ptab.add_row(['PSNR', f"{psnr(gt, pred):.5}"])
    ptab.add_row(['SSIM', f"{ssim(gt, pred):.5}"])
    # display
    print(ptab)
