# Created by BVH, Apr 2026.
# Metrics computation for A4D training / validation.
# Separated from a4d_visuals.py to keep file sizes manageable.

import torch

from custom.any4d.a4d_logistics import detach_dict_recursive
from custom.eval.metrics import calculate_mae, calculate_mse, calculate_psnr, calculate_ssim, calculate_lpips
from custom.eval.visuals import (
    cached_decode_latent_lowdim_auto,
    cached_decode_latent_video_auto,
)


def calculate_metrics_a4d(
    phase: str = 'train',
    train_step: int = -1,
    data_batch: dict[str, torch.Tensor] = None,
    output_dict: dict[str, torch.Tensor] = None,
    a4d_vae=None,
    config=None,
):
    '''
    NOTE: This function can be called on ALL ranks, and processes ALL samples within the batch,
        as long as they have at least one predicted token.
    :param phase (str): train / val / test.
    :param train_step (int): Batch index, like training step but changes along with gradient
        accumulation. Starts from 1.
    :param data_batch (dict): Input batch, contains anydata, fps, prompt, etc.
    :param output_dict (dict): Output batch, contains both encoded (latent) and decoded (raw)
        input/output/GT entries.
    :param a4d_vae (VAE).
    :param config (Any4DConfig).
    :return metrics_dict (dict): Maps metric name to list of size B of float.
        Subkey friendly maps metric name to list of size B of str.
    '''
    output_dict = detach_dict_recursive(output_dict)

    gts_for_metrics = dict()
    preds_for_metrics = dict()
    has_output = dict()
    has_gt = dict()

    # TODO(bvh): Apply supervise_mask to metrics to accommodate partial GT availability?
    # TODO(bvh): Check per batch index if each entry actually has output, or is just fully conditional.

    for k in config.all_highdim_entries:
        if k.endswith('_mask'):
            continue

        if k in config.track_metrics.keys():
            if k not in output_dict['x0_entries'] or k not in data_batch['a4d_latent']:
                continue

            has_output[k] = (data_batch['a4d_latent'][k + '_output_mask'].sum(dim=(1, 2, 3, 4)) > 0.0)
            # ^ (B) tensor of bool.

            assert k in output_dict['y0_pred_entries'], f'{k} (output) not in y0_pred_entries'
            assert k not in gts_for_metrics, f'{k} already in gts_for_metrics'
            assert k not in preds_for_metrics, f'{k} already in preds_for_metrics'

            # NOTE(bvh): The unsupervised input regions are replaced with HY / conditioning information.
            # This slightly inflates metrics, but is preferred anyway due to simplicity and standard practice.
            gts_for_metrics[k] = cached_decode_latent_video_auto(
                output_dict['x0_entries'], k, a4d_vae)
            preds_for_metrics[k] = cached_decode_latent_video_auto(
                output_dict['y0_pred_entries'], k, a4d_vae)

    for k in config.all_lowdim_entries:
        if k.endswith('_mask'):
            continue

        if k in config.track_metrics.keys():
            if k not in output_dict['x0_entries'] or k not in data_batch['a4d_latent']:
                continue

            has_output[k] = (data_batch['a4d_latent'][k + '_output_mask'].sum(dim=(1, 2)) > 0.0)
            # ^ (B) tensor of bool.

            assert k in output_dict['y0_pred_entries'], f'{k} (output) not in y0_pred_entries'
            assert k not in gts_for_metrics, f'{k} already in gts_for_metrics'
            assert k not in preds_for_metrics, f'{k} already in preds_for_metrics'

            gts_for_metrics[k] = cached_decode_latent_lowdim_auto(
                output_dict['x0_entries'], k, a4d_vae)
            preds_for_metrics[k] = cached_decode_latent_lowdim_auto(
                output_dict['y0_pred_entries'], k, a4d_vae)

    # WARNING(bvh): These metrics are based on the ENCODED + DECODED ground truth videos!
    # This might seem weird but I think it's preferable because then we can disentangle
    # errors arising from the VAE versus the diffusion model itself.
    psnr_dict = dict()
    ssim_dict = dict()
    mse_dict = dict()
    mae_dict = dict()
    lpips_dict = dict()

    for (k, v) in config.track_metrics.items():
        if k not in gts_for_metrics:
            continue

        assert k in preds_for_metrics, f'{k} not in preds_for_metrics'
        B = gts_for_metrics[k].shape[0]

        if 'psnr' in v:
            psnr_dict[k] = [calculate_psnr(
                gts_for_metrics[k][b], preds_for_metrics[k][b])
                if has_output[k][b] else None
                for b in range(B)]
        if 'ssim' in v:
            ssim_dict[k] = [calculate_ssim(
                gts_for_metrics[k][b], preds_for_metrics[k][b])
                if has_output[k][b] else None
                for b in range(B)]
        if 'mse' in v:
            mse_dict[k] = [calculate_mse(
                gts_for_metrics[k][b], preds_for_metrics[k][b])
                if has_output[k][b] else None
                for b in range(B)]
        if 'mae' in v:
            mae_dict[k] = [calculate_mae(
                gts_for_metrics[k][b], preds_for_metrics[k][b])
                if has_output[k][b] else None
                for b in range(B)]
        if 'lpips' in v:
            lpips_dict[k] = [calculate_lpips(
                gts_for_metrics[k][b], preds_for_metrics[k][b])
                if has_output[k][b] else None
                for b in range(B)]

    psnr_friendly = ['  '.join([f'{k}: {m[b]:.2f}'
                                for (k, m) in psnr_dict.items()
                                if m[b] is not None]) for b in range(B)]
    ssim_friendly = ['  '.join([f'{k}: {m[b]:.3f}'
                                for (k, m) in ssim_dict.items()
                                if m[b] is not None]) for b in range(B)]
    mse_friendly = ['  '.join([f'{k}: {m[b]:.3f}'
                                for (k, m) in mse_dict.items()
                                if m[b] is not None]) for b in range(B)]
    mae_friendly = ['  '.join([f'{k}: {m[b]:.3f}'
                                for (k, m) in mae_dict.items()
                                if m[b] is not None]) for b in range(B)]
    lpips_friendly = ['  '.join([f'{k}: {m[b]:.4f}'
                                for (k, m) in lpips_dict.items()
                                if m[b] is not None]) for b in range(B)]

    # Keep track of these such that it can be aggregated correctly later:
    dset_name = data_batch['dset_name']
    dl_idx = [int(x) for x in data_batch['dl_idx']]
    dl_key = data_batch.get('dl_key', None) or ''
    sample_id = data_batch.get('sample_id', ['unknown'] * len(dset_name))

    metrics_dict = {
        'psnr': psnr_dict,
        'ssim': ssim_dict,
        'mse': mse_dict,
        'mae': mae_dict,
        'lpips': lpips_dict,
        'friendly': {
            'psnr': psnr_friendly,
            'ssim': ssim_friendly,
            'mse': mse_friendly,
            'mae': mae_friendly,
            'lpips': lpips_friendly,
        },
        'dset_name': dset_name,
        'dl_idx': dl_idx,
        'dl_key': dl_key,
        'sample_id': sample_id,
        'true_val_size': data_batch.get('true_val_size', None),
    }

    return metrics_dict