# Created by BVH, Jul - Aug 2025.
# Helper methods for metrics.

import os
import time

import lovely_tensors
import matplotlib
import numpy as np
import skimage.metrics
import torch
from einops import rearrange
from lovely_numpy import lo
from rich import print

try:
    import lpips as lpips_lib
    _LPIPS_NET = None  # lazy singleton
except ImportError:
    lpips_lib = None
    _LPIPS_NET = None

from imaginaire.utils import log

lovely_tensors.monkey_patch()
np.set_printoptions(precision=3, suppress=True)
torch.set_printoptions(precision=3, sci_mode=False, threshold=1000)



def calculate_psnr(gt, pred):
    '''
    :param gt (3, T, H, W) tensor of float in [0, 1].
    :param pred (3, T, H, W) tensor of float in [0, 1].
    :return psnr: single float.
    '''
    (C, T, H, W) = gt.shape

    assert not gt.isnan().any(), f'NaN in gt: {gt.isnan().sum()} values'
    assert not pred.isnan().any(), f'NaN in pred: {pred.isnan().sum()} values'
    assert gt.min() >= 0.0 and gt.max() <= 1.0, f'gt out of [0,1]: min={gt.min():.4f} max={gt.max():.4f}'
    assert pred.min() >= 0.0 and pred.max() <= 1.0, f'pred out of [0,1]: min={pred.min():.4f} max={pred.max():.4f}'

    # Temporary fix: Check for all black frames (e.g. DyCheck)
    valid_frames = (gt.mean(dim=(0, 2, 3)) > 0.01)  # (T) tensor of bool.

    gt = gt.clamp(0.0, 1.0).to(torch.float64)
    pred = pred.clamp(0.0, 1.0).to(torch.float64)

    # psnr = -10.0 * torch.log10(torch.mean(torch.square(gt - pred))).item()
    psnr_values = [-10.0 * torch.log10(torch.mean(torch.square(gt[:, t] - pred[:, t]))).item()
                   for t in range(T) if valid_frames[t]]

    psnr = np.mean(psnr_values)

    return psnr


def calculate_ssim(gt, pred):
    '''
    :param gt (3, T, H, W) tensor of float in [0, 1].
    :param pred (3, T, H, W) tensor of float in [0, 1].
    :return ssim: single float.
    '''
    (C, T, H, W) = gt.shape

    assert not gt.isnan().any(), f'NaN in gt: {gt.isnan().sum()} values'
    assert not pred.isnan().any(), f'NaN in pred: {pred.isnan().sum()} values'
    assert gt.min() >= 0.0 and gt.max() <= 1.0, f'gt out of [0,1]: min={gt.min():.4f} max={gt.max():.4f}'
    assert pred.min() >= 0.0 and pred.max() <= 1.0, f'pred out of [0,1]: min={pred.min():.4f} max={pred.max():.4f}'

    # Temporary fix: Check for all black frames (e.g. DyCheck)
    valid_frames = (gt.mean(dim=(0, 2, 3)) > 0.01)  # (T) tensor of bool.

    gt = gt.clamp(0.0, 1.0).to(torch.float64)
    pred = pred.clamp(0.0, 1.0).to(torch.float64)

    gt_np = gt.detach().cpu().numpy()    # (3, T, H, W)
    pred_np = pred.detach().cpu().numpy()    # (3, T, H, W)

    ssim_values = []
    for t in range(T):
        if not valid_frames[t]:
            continue
        
        gt_frame = np.transpose(gt_np[:, t], (1, 2, 0))     # (H, W, 3)
        pred_frame = np.transpose(pred_np[:, t], (1, 2, 0))  # (H, W, 3)
        ssim_val = skimage.metrics.structural_similarity(
            gt_frame, pred_frame, data_range=1.0, channel_axis=2)
            # gaussian_weights=True, sigma=1.5, use_sample_covariance=False)
        ssim_values.append(ssim_val)
    
    ssim = np.mean(ssim_values)

    return ssim


def calculate_mse(gt, pred):
    '''
    :param gt (*) tensor of float.
    :param pred (*) tensor of float.
    :return mse: single tensor of float.
    '''
    gt = gt.to(torch.float64)
    pred = pred.to(torch.float64)
    mse = torch.mean(torch.square(gt - pred)).item()
    return mse


def calculate_mae(gt, pred):
    '''
    :param gt (*) tensor of float.
    :param pred (*) tensor of float.
    :return mae: single tensor of float.
    '''
    gt = gt.to(torch.float64)
    pred = pred.to(torch.float64)
    mae = torch.mean(torch.abs(gt - pred)).item()
    return mae


def _get_lpips_net(device):
    global _LPIPS_NET
    if lpips_lib is None:
        raise ImportError('lpips package not installed')
    if _LPIPS_NET is None:
        _LPIPS_NET = lpips_lib.LPIPS(net='vgg', verbose=False)
    _LPIPS_NET = _LPIPS_NET.to(device)
    return _LPIPS_NET


def calculate_lpips(gt, pred):
    '''
    :param gt (3, T, H, W) tensor of float in [0, 1].
    :param pred (3, T, H, W) tensor of float in [0, 1].
    :return lpips_val: single float.
    '''
    (C, T, H, W) = gt.shape

    assert not gt.isnan().any(), f'NaN in gt: {gt.isnan().sum()} values'
    assert not pred.isnan().any(), f'NaN in pred: {pred.isnan().sum()} values'
    assert gt.min() >= 0.0 and gt.max() <= 1.0, f'gt out of [0,1]: min={gt.min():.4f} max={gt.max():.4f}'
    assert pred.min() >= 0.0 and pred.max() <= 1.0, f'pred out of [0,1]: min={pred.min():.4f} max={pred.max():.4f}'

    valid_frames = (gt.mean(dim=(0, 2, 3)) > 0.01)  # (T) tensor of bool.

    gt = gt.clamp(0.0, 1.0).float()
    pred = pred.clamp(0.0, 1.0).float()

    # LPIPS expects (N, 3, H, W) in [-1, 1]
    net = _get_lpips_net(gt.device)

    lpips_values = []
    for t in range(T):
        if not valid_frames[t]:
            continue
        # gt[:, t] is (3, H, W), unsqueeze to (1, 3, H, W)
        val = net(pred[:, t].unsqueeze(0) * 2 - 1,
                  gt[:, t].unsqueeze(0) * 2 - 1)
        lpips_values.append(val.item())

    return float(np.mean(lpips_values))


def flatten_metrics_dict(metrics_dict):
    flattened_metrics = {}
    for metric, modalities in metrics_dict.items():
        for modality, value in modalities.items():
            flattened_metrics[f'{metric}_{modality}'] = value
    return flattened_metrics


def aggregate_all_metrics(
    all_metrics: list[dict[str, dict[str, list[float]]]]
) -> dict[str, dict[str, list[float]]]:
    '''
    :param all_metrics (list): List -> metric -> entry -> list of float/None.
    :return agg_metrics (dict): Metric -> entry -> long list of (dset, idx, float/None).
    '''
    agg_metrics = dict()
    agg_metrics['dset_name'] = []
    agg_metrics['dl_idx'] = []
    true_val_sizes = dict()  # dl_key -> true dataset size (before DistributedSampler padding)

    for metrics_dict in all_metrics:
        if metrics_dict is None:
            continue

        # Track true val sizes per dl_key
        _dl_key = metrics_dict.get('dl_key', '')
        _tvs = metrics_dict.get('true_val_size', None)
        if _dl_key and _tvs is not None and _dl_key not in true_val_sizes:
            true_val_sizes[_dl_key] = _tvs

        dset_name = metrics_dict['dset_name']  # list of str
        dl_idx = metrics_dict['dl_idx']  # list of int
        dl_key = metrics_dict.get('dl_key', '')  # str
        sample_id = metrics_dict.get('sample_id', ['unknown'] * len(dset_name))  # list of str
        true_val_size = metrics_dict.get('true_val_size', None)  # int or None

        for metric_name, entries_dict in metrics_dict.items():
            if metric_name in ('friendly', 'dset_name', 'dl_idx', 'dl_key', 'sample_id', 'true_val_size'):
                continue
            if 'friendly' in metric_name:
                continue

            if metric_name not in agg_metrics:
                agg_metrics[metric_name] = dict()

            for entry_name, values_list in entries_dict.items():
                if entry_name not in agg_metrics[metric_name]:
                    agg_metrics[metric_name][entry_name] = []

                # NOTE(bvh): keep metadata close to values because different batches might have
                # different subsets of metrics/entries, which would cause lists to diverge.
                tagged_values = list(zip(dset_name, dl_idx, values_list,
                                         [dl_key] * len(values_list), sample_id))
                agg_metrics[metric_name][entry_name].extend(tagged_values)
    
        # agg_metrics['dset_name'].extend(metrics_dict['dset_name'])
        # agg_metrics['dl_idx'].extend(metrics_dict['dl_idx'])

    agg_metrics['_true_val_sizes'] = true_val_sizes
    return agg_metrics


def average_metrics_per_dataset(
    agg_metrics: dict[str, dict[str, list[float]]]
) -> dict[str, dict[str, float]]:
    '''
    :param agg_metrics (dict): Metric -> entry -> long list of (dset, idx, float/None).
    :return avg_metrics (dict): Metric -> entry -> dset -> (avg, count).
    '''
    avg_metrics = dict()
    true_val_sizes = agg_metrics.pop('_true_val_sizes', {})

    for metric_name, entries_dict in agg_metrics.items():
        if metric_name.startswith('_') or 'friendly' in metric_name or 'dset_name' in metric_name or 'dl_idx' in metric_name:
            continue
        
        avg_metrics[metric_name] = dict()
        
        for entry_name, (agg_tuples) in entries_dict.items():
            avg_metrics[metric_name][entry_name] = dict()

            dset_name_list = [x[0] for x in agg_tuples]
            dl_key_list = [x[3] if len(x) > 3 else '' for x in agg_tuples]
            values_list = [x[2] for x in agg_tuples]

            # Truncate per dl_key to true_val_size (DistributedSampler padding creates extra samples).
            # NOTE(bvh): we keep only the first `true_val_size` samples per dl_key per val round.
            if true_val_sizes:
                key_counts = {}
                truncated_dset = []
                truncated_vals = []
                for i in range(len(dset_name_list)):
                    dk = dl_key_list[i]
                    max_n = true_val_sizes.get(dk, float('inf'))
                    key_counts[dk] = key_counts.get(dk, 0) + 1
                    if key_counts[dk] <= max_n:
                        truncated_dset.append(dset_name_list[i])
                        truncated_vals.append(values_list[i])
                dset_name_list = truncated_dset
                values_list = truncated_vals

            # Calculate total statistics across all datasets
            valid_values = [x for x in values_list
                            if x is not None and not np.isnan(x) and not np.isinf(x)]
            avg_metrics[metric_name][entry_name]['all'] = \
                (np.mean(valid_values) if len(valid_values) > 0 else -1.0, len(valid_values))

            # Calculate statistics for each dataset
            for dset_name in set(dset_name_list):
                subset_indices = [i for i, dset in enumerate(dset_name_list)
                                  if dset == dset_name]
                subset_values = [values_list[i] for i in subset_indices]
                valid_values = [x for x in subset_values
                                if x is not None and not np.isnan(x) and not np.isinf(x)]
                avg_metrics[metric_name][entry_name][dset_name] = \
                    (np.mean(valid_values) if len(valid_values) > 0 else -1.0, len(valid_values))

    return avg_metrics
    


def make_masks(a4d_lat):
    # TODO @vitor documentation
    
    # NOTE: only video output_masks (5D) -- skip lowdim mask companions.
    masks = {key.split('_')[0]: val for key, val in a4d_lat.items()
             if key.endswith('output_mask') and val.ndim == 5}

    output = {}
    for key, val in masks.items():

        valid = val.max() == 0
        b, _, t, h, w = val.shape
        red = val.view(b * t, h * w)
        red = torch.min(red, 1)[0]
        red = red.view(b, t)

        if red.min() == 1:
            output[key] = 'valid'
        elif red.max() == 0:
            output[key] = 'invalid'
        else:
            idx = torch.argmax(red)
            idx = (idx - 1) * 4 + 1
            output[key] = idx

    return output


def mask_entries(entries, masks):
    # TODO(vitor): documentation

    if masks is None:
        return entries
    entries_masked = {}
    for key, val in masks.items():
        if key not in entries:
            continue
        if masks[key] == 'valid':
            entries_masked[key] = entries[key]
        elif masks[key] == 'invalid':
            pass
        else:
            idx = masks[key]
            entries_masked[key] = entries[key][:, :, idx:]
    return entries_masked


def unroll_a4d(data_batch, val_dict, remove_cond=True):
    """Convert A4D data to gt/rec/pred dictionaries for metrics evaluation.

    Only video (5D) entries flow through. Lowdim entries (action/proprio/traj/
    relcams/...) share the `_cdec` cache key with video ones but are 3D
    (B, T, C) -- they would break the per-(time, cam) iteration below.
    """

    gt_entries = data_batch['a4d_raw']

    # Get used modalities and cameras to unroll back to original dicts
    used_modals = [m for m in data_batch['meta']['used_modals'].split(',')]
    used_cams = [int(c) for c in data_batch['meta']['used_cams'].split(',')]

    # Strip _cdec suffix; only keep video (5D) entries.
    # NOTE: a4d_raw can contain non-tensor values
    rec_entries = {key[:-5]: val for key, val in val_dict['x0_entries'].items()
                   if key.endswith('_cdec') and torch.is_tensor(val) and val.ndim == 5}
    pred_entries = {key[:-5]: val for key, val in val_dict['y0_pred_entries'].items()
                    if key.endswith('_cdec') and torch.is_tensor(val) and val.ndim == 5}
    gt_entries = {k: v for k, v in gt_entries.items()
                  if torch.is_tensor(v) and v.ndim == 5}

    masks = make_masks(data_batch['a4d_latent']) if remove_cond else None

    gt_entries = mask_entries(gt_entries, masks)
    rec_entries = mask_entries(rec_entries, masks)
    pred_entries = mask_entries(pred_entries, masks)

    # Unroll each specific dictionary
    gt_dict = unroll_a4d_single(gt_entries, used_modals, used_cams, masks)
    rec_dict = unroll_a4d_single(rec_entries, used_modals, used_cams, masks)
    pred_dict = unroll_a4d_single(pred_entries, used_modals, used_cams, masks)

    return dict(gt=gt_dict, rec=rec_dict, pred=pred_dict, info=data_batch['meta'])


def unroll_a4d_single(entries, used_modals, used_cams, masks=None):
    """Helper function to unroll specific A4D dictionaries (see above)"""
    out = {}
    for modal in used_modals:
        for key, val in entries.items():
            if key.startswith(modal):
                view = int(key[len(modal):])
                if modal not in out:
                    out[modal] = {}
                if masks is None or masks[key] == 'valid':
                    st, fn = 0, val.shape[2]
                else:
                    st, fn = masks[key], masks[key] + val.shape[2]
                for i, time in enumerate(range(st, fn)):
                    out[modal][(time, view)] = val[:, :, i]
    return out