"""配准质量评估工具。

这里刻意只依赖 NumPy，便于在训练、推理和 Streamlit 三处复用。
"""

from __future__ import annotations

from typing import Dict, Iterable, Tuple

import numpy as np


def crop_to_common_shape(*arrays: np.ndarray) -> Tuple[np.ndarray, ...]:
    """把多个 3D/4D 数组中心裁剪到共同尺寸，避免形状不一致导致评估失败。"""

    if not arrays:
        return tuple()

    spatial_shapes = [arr.shape[:3] for arr in arrays]
    common = tuple(min(shape[axis] for shape in spatial_shapes) for axis in range(3))
    cropped = []

    for arr in arrays:
        slices = []
        for axis, target in enumerate(common):
            start = max((arr.shape[axis] - target) // 2, 0)
            slices.append(slice(start, start + target))
        if arr.ndim > 3:
            slices.append(slice(None))
        cropped.append(arr[tuple(slices)])

    return tuple(cropped)


def global_ncc(a: np.ndarray, b: np.ndarray, eps: float = 1e-8) -> float:
    """计算全局 NCC。训练时使用局部 NCC，这里用于快速量化展示。"""

    a = np.asarray(a, dtype=np.float32)
    b = np.asarray(b, dtype=np.float32)
    a = a - float(np.mean(a))
    b = b - float(np.mean(b))
    denom = float(np.sqrt(np.sum(a * a) * np.sum(b * b)) + eps)
    return float(np.sum(a * b) / denom)


def mse(a: np.ndarray, b: np.ndarray) -> float:
    diff = np.asarray(a, dtype=np.float32) - np.asarray(b, dtype=np.float32)
    return float(np.mean(diff * diff))


def mae(a: np.ndarray, b: np.ndarray) -> float:
    diff = np.abs(np.asarray(a, dtype=np.float32) - np.asarray(b, dtype=np.float32))
    return float(np.mean(diff))


def registration_metrics(
    fixed: np.ndarray,
    moving: np.ndarray,
    warped: np.ndarray,
) -> Dict[str, float]:
    """输出配准前后可比较的常用指标。"""

    fixed, moving, warped = crop_to_common_shape(fixed, moving, warped)
    before_mse = mse(fixed, moving)
    after_mse = mse(fixed, warped)
    before_mae = mae(fixed, moving)
    after_mae = mae(fixed, warped)
    before_ncc = global_ncc(fixed, moving)
    after_ncc = global_ncc(fixed, warped)

    return {
        "before_mse": before_mse,
        "after_mse": after_mse,
        "before_mae": before_mae,
        "after_mae": after_mae,
        "before_ncc": before_ncc,
        "after_ncc": after_ncc,
        "mse_improvement": before_mse - after_mse,
        "mae_improvement": before_mae - after_mae,
        "ncc_improvement": after_ncc - before_ncc,
    }


def slice_metric_curve(
    fixed: np.ndarray,
    moving: np.ndarray,
    warped: np.ndarray,
    axis: int = 2,
    metric: str = "mse",
) -> Dict[str, Iterable[float]]:
    """逐切片计算配准前后指标曲线。"""

    fixed, moving, warped = crop_to_common_shape(fixed, moving, warped)
    metric_funcs = {
        "mse": mse,
        "mae": mae,
        "ncc": global_ncc,
    }
    if metric not in metric_funcs:
        raise ValueError(f"不支持的逐切片指标: {metric}")
    metric_fn = metric_funcs[metric]

    before = []
    after = []

    for index in range(fixed.shape[axis]):
        selector = [slice(None)] * 3
        selector[axis] = index
        selector = tuple(selector)
        before.append(metric_fn(fixed[selector], moving[selector]))
        after.append(metric_fn(fixed[selector], warped[selector]))

    return {
        "slice_index": list(range(fixed.shape[axis])),
        f"before_{metric}": before,
        f"after_{metric}": after,
    }


def ddf_summary(ddf_xyz: np.ndarray) -> Dict[str, float]:
    """统计形变场向量大小。输入应为 X/Y/Z/3，单位可为 voxel 或 mm。"""

    if ddf_xyz.ndim != 4 or ddf_xyz.shape[-1] != 3:
        raise ValueError("DDF 必须是形状为 (X, Y, Z, 3) 的 4D 数组。")

    mag = np.linalg.norm(ddf_xyz.astype(np.float32), axis=-1)
    return {
        "ddf_mean": float(np.mean(mag)),
        "ddf_std": float(np.std(mag)),
        "ddf_p95": float(np.percentile(mag, 95)),
        "ddf_max": float(np.max(mag)),
    }