from __future__ import annotations

import torch


CSI_THRESHOLDS = torch.tensor([16, 74, 133, 160, 181, 219], dtype=torch.float32) / 255.0


def mse(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.mean((prediction - target) ** 2)


def _ssim_per_sample(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    prediction, target = prediction.clamp(0.0, 1.0), target.clamp(0.0, 1.0)
    mu_x, mu_y = prediction.mean((2, 3)), target.mean((2, 3))
    centered_x = prediction - mu_x[:, :, None, None, :]
    centered_y = target - mu_y[:, :, None, None, :]
    var_x = centered_x.square().mean((2, 3))
    var_y = centered_y.square().mean((2, 3))
    covariance = (centered_x * centered_y).mean((2, 3))
    c1, c2 = 0.01**2, 0.03**2
    score = ((2 * mu_x * mu_y + c1) * (2 * covariance + c2)) / (
        (mu_x.square() + mu_y.square() + c1) * (var_x + var_y + c2)
    )
    return score.mean((1, 2))


def metric_sums(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Return additive sample-level MSE/MAE/SSIM sums and CSI event counts."""
    prediction, target = prediction.float(), target.float()
    reduce_dims = tuple(range(1, prediction.ndim))
    values = [
        (prediction - target).square().mean(reduce_dims).sum(),
        (prediction - target).abs().mean(reduce_dims).sum(),
        _ssim_per_sample(prediction, target).sum(),
        prediction.new_tensor(prediction.shape[0]),
    ]
    for threshold in CSI_THRESHOLDS.to(prediction.device):
        predicted, observed = prediction >= threshold, target >= threshold
        values.extend([(predicted & observed).sum(), (predicted & ~observed).sum(), (~predicted & observed).sum()])
    return torch.stack(values).to(torch.float64)


def metrics_from_sums(sums: torch.Tensor) -> dict[str, float]:
    count = sums[3].clamp_min(1)
    csi = []
    for index in range(len(CSI_THRESHOLDS)):
        hits, false_alarms, misses = sums[4 + index * 3 : 7 + index * 3]
        csi.append((hits / (hits + false_alarms + misses).clamp_min(1)).item())
    return {
        "mse": (sums[0] / count).item(),
        "mae": (sums[1] / count).item(),
        "ssim": (sums[2] / count).item(),
        "mean_csi": sum(csi) / len(csi),
    }


def metric_sums_light(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Additive MSE/MAE sums for cheap per-epoch validation: [mse_sum, mae_sum, count]."""
    prediction, target = prediction.float(), target.float()
    reduce_dims = tuple(range(1, prediction.ndim))
    return torch.stack(
        [
            (prediction - target).square().mean(reduce_dims).sum(),
            (prediction - target).abs().mean(reduce_dims).sum(),
            prediction.new_tensor(prediction.shape[0]),
        ]
    ).to(torch.float64)


def metrics_from_light_sums(sums: torch.Tensor) -> dict[str, float]:
    count = sums[2].clamp_min(1)
    return {"mse": (sums[0] / count).item(), "mae": (sums[1] / count).item()}


def compute_metrics(prediction: torch.Tensor, target: torch.Tensor) -> dict[str, float]:
    return metrics_from_sums(metric_sums(prediction.detach(), target.detach()))