"""官方 VoxelMorph 训练适配器。

本项目以 https://github.com/voxelmorph/voxelmorph 的 PyTorch 核心为运行核心：
- 网络: ``voxelmorph.nn.models.VxmPairwise``
- 空间变换: 官方 ``voxelmorph.nn.modules.SpatialTransformer``
- 平滑正则: 官方推荐的 ``neurite.nn.modules.SpatialGradient``
- 相似度: 默认使用带 epsilon 的本地稳定 NCC；可切换到 ``neurite.nn.modules.NCC``

输入张量约定：
- NIfTI 数组轴顺序为 (X, Y, Z)。
- 官方 PyTorch VoxelMorph 输入为 (B, C, X, Y, Z)。
- DDF 通道顺序与空间轴一致，即 (X, Y, Z)。
"""

from __future__ import annotations

import argparse
from contextlib import nullcontext
import json
import os
from pathlib import Path
from typing import Dict, Iterable, List, Sequence, Tuple

import numpy as np

from config import DEFAULT_CHECKPOINT

OFFICIAL_VXM_REPO = "https://github.com/voxelmorph/voxelmorph"
OFFICIAL_VXM_COMMIT = "db73f34b910bcefcb520f7f40a1bc4a3e0b6401d"

try:  # 让缺 torch 的环境仍能 py_compile。
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
except Exception as _torch_exc:  # pragma: no cover
    torch = None  # type: ignore
    nn = None  # type: ignore
    F = None  # type: ignore
    _TORCH_IMPORT_ERROR = _torch_exc
else:
    _TORCH_IMPORT_ERROR = None


def configure_voxelmorph_backend() -> None:
    """显式启用官方 VoxelMorph/Neurite 的 PyTorch 后端。"""

    os.environ.setdefault("NEURITE_BACKEND", "pytorch")
    os.environ.setdefault("VXM_BACKEND", "pytorch")


def require_torch():
    if torch is None:
        raise RuntimeError(
            "缺少 PyTorch，无法训练或推理。建议使用 conda 创建 CUDA 环境："
            "conda env create -f environment.yml"
        ) from _TORCH_IMPORT_ERROR
    return torch


def require_official_voxelmorph():
    """导入官方 VoxelMorph 与 Neurite，并检查关键 API。"""

    configure_voxelmorph_backend()
    require_torch()
    try:
        import neurite as ne  # type: ignore
        import voxelmorph as vxm  # type: ignore
    except Exception as exc:  # pragma: no cover
        raise RuntimeError(
            "缺少官方 VoxelMorph/Neurite。请运行: conda env create -f environment.yml "
            "或 pip install -r requirements.txt"
        ) from exc

    if not hasattr(vxm.nn.models, "VxmPairwise"):
        raise RuntimeError("当前 voxelmorph 版本缺少 vxm.nn.models.VxmPairwise，请安装官方 dev 版本。")
    return vxm, ne


def _require_nibabel():
    try:
        import nibabel as nib  # type: ignore
    except Exception as exc:  # pragma: no cover
        raise RuntimeError("缺少 nibabel，请先安装项目依赖。") from exc
    return nib


if torch is not None:

    class LocalNCCLoss(nn.Module):
        """本地 NCC 兜底。

        官方 VoxelMorph 当前把 NCC 放在 Neurite 中；若用户安装的 Neurite
        接口短期变动，本类保证训练脚本仍可运行，但网络核心仍是官方 VoxelMorph。
        """

        def __init__(self, window_size: int = 9, eps: float = 1e-5):
            super().__init__()
            self.window_size = int(window_size)
            self.eps = float(eps)

        def forward(self, fixed, warped):
            win = [self.window_size] * 3
            padding = self.window_size // 2
            filt = torch.ones((1, 1, *win), dtype=fixed.dtype, device=fixed.device)
            win_volume = float(np.prod(win))

            fixed_sum = F.conv3d(fixed, filt, padding=padding)
            warped_sum = F.conv3d(warped, filt, padding=padding)
            fixed2_sum = F.conv3d(fixed * fixed, filt, padding=padding)
            warped2_sum = F.conv3d(warped * warped, filt, padding=padding)
            cross_sum = F.conv3d(fixed * warped, filt, padding=padding)

            fixed_mean = fixed_sum / win_volume
            warped_mean = warped_sum / win_volume
            cross = cross_sum - warped_mean * fixed_sum - fixed_mean * warped_sum + fixed_mean * warped_mean * win_volume
            fixed_var = fixed2_sum - 2 * fixed_mean * fixed_sum + fixed_mean * fixed_mean * win_volume
            warped_var = warped2_sum - 2 * warped_mean * warped_sum + warped_mean * warped_mean * win_volume

            cc = (cross * cross) / (fixed_var * warped_var + self.eps)
            return -torch.mean(cc)


    class NegativeSimilarity(nn.Module):
        """把 Neurite NCC 分数转换为可最小化的 loss。"""

        def __init__(self, module: nn.Module):
            super().__init__()
            self.module = module

        def forward(self, fixed, warped):
            return -torch.mean(self.module(fixed, warped))


else:

    class LocalNCCLoss:  # type: ignore[no-redef]
        def __init__(self, *args, **kwargs):
            require_torch()


    class NegativeSimilarity:  # type: ignore[no-redef]
        def __init__(self, *args, **kwargs):
            require_torch()


def load_nifti_tensor(path: str | Path, device: str | "torch.device" = "cpu") -> Tuple["torch.Tensor", np.ndarray]:
    """读取 NIfTI 并转换为官方 VoxelMorph 输入张量。

    NIfTI data: (X, Y, Z)
    Torch tensor: (B, C, X, Y, Z)
    """

    require_torch()
    nib = _require_nibabel()
    img = nib.load(str(path), mmap=True)
    data = np.asanyarray(img.dataobj, dtype=np.float32)
    if data.ndim > 3:
        data = data[..., 0]
    data = np.nan_to_num(data.astype(np.float32, copy=False), copy=False)
    tensor = torch.from_numpy(data.copy())[None, None].to(device=device, dtype=torch.float32)
    return tensor, img.affine.copy()


def resolve_device(device: str):
    require_torch()
    if device == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(device)


def build_vxm_model(
    nb_features: Sequence[int],
    integration_steps: int,
    device,
):
    """构建官方 VoxelMorph Pairwise 模型。"""

    vxm, _ = require_official_voxelmorph()
    model = vxm.nn.models.VxmPairwise(
        ndim=3,
        source_channels=1,
        target_channels=1,
        nb_features=list(int(v) for v in nb_features),
        integration_steps=int(integration_steps),
        device=str(device),
    )
    return model.to(device)


def build_similarity_loss(loss_name: str, ncc_window: int, ncc_impl: str = "local"):
    """构建相似度 loss。

    颈部 CT 预处理后通常有大面积零填充背景，Neurite NCC 在低方差窗口中可能
    返回非有限值。因此默认使用带 epsilon 的本地 NCC；模型核心仍为官方
    VoxelMorph。需要对照官方依赖时，可传入 ``ncc_impl="neurite"``。
    """

    _, ne = require_official_voxelmorph()
    loss_name = loss_name.lower()
    if loss_name == "ncc":
        ncc_impl = ncc_impl.lower()
        if ncc_impl == "neurite":
            return NegativeSimilarity(ne.nn.modules.NCC(window_size=int(ncc_window)))
        if ncc_impl == "local":
            return LocalNCCLoss(window_size=int(ncc_window))
        raise ValueError("ncc_impl 只能是 local 或 neurite。")
    if loss_name == "mse":
        return ne.nn.modules.MSE()
    raise ValueError("loss_name 只能是 ncc 或 mse。")


def build_gradient_loss():
    """官方平滑项：Neurite SpatialGradient。"""

    _, ne = require_official_voxelmorph()
    return ne.nn.modules.SpatialGradient(penalty="l2")


def train_pair(
    moving_path: str | Path,
    fixed_path: str | Path,
    checkpoint_path: str | Path = DEFAULT_CHECKPOINT,
    epochs: int = 200,
    learning_rate: float = 1e-4,
    smooth_weight: float = 0.01,
    image_loss: str = "ncc",
    ncc_window: int = 9,
    ncc_impl: str = "local",
    nb_features: Sequence[int] = (16, 16, 16, 16, 16),
    integration_steps: int = 0,
    device: str = "auto",
    use_amp: bool = True,
    save_every: int = 50,
) -> List[Dict[str, float]]:
    """使用一对 moving/fixed 进行无监督 VoxelMorph 训练。"""

    require_official_voxelmorph()
    device_obj = resolve_device(device)
    moving, _ = load_nifti_tensor(moving_path, device=device_obj)
    fixed, _ = load_nifti_tensor(fixed_path, device=device_obj)
    if moving.shape != fixed.shape:
        raise ValueError(f"Moving 与 Fixed 形状不一致: {tuple(moving.shape)} vs {tuple(fixed.shape)}")

    model = build_vxm_model(nb_features=nb_features, integration_steps=integration_steps, device=device_obj)
    similarity_loss = build_similarity_loss(image_loss, ncc_window=ncc_window, ncc_impl=ncc_impl).to(device_obj)
    grad_loss_fn = build_gradient_loss().to(device_obj)
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    amp_enabled = bool(use_amp and device_obj.type == "cuda")
    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)

    checkpoint_path = Path(checkpoint_path)
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
    history: List[Dict[str, float]] = []

    for epoch in range(1, epochs + 1):
        model.train()
        optimizer.zero_grad(set_to_none=True)

        try:
            autocast_context = torch.amp.autocast("cuda", enabled=True) if amp_enabled else nullcontext()
            with autocast_context:
                displacement, warped = model(
                    moving,
                    fixed,
                    return_warped_source=True,
                    return_field_type="displacement",
                )

            loss_sim = similarity_loss(fixed.float(), warped.float())
            loss_smooth = grad_loss_fn(displacement.float())
            loss = loss_sim + smooth_weight * loss_smooth

            if not torch.isfinite(loss):
                raise RuntimeError(
                    "训练 loss 出现非有限值。可尝试使用 --ncc-impl local、"
                    "--image-loss mse，或检查预处理后是否有大面积常值区域。"
                )

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        except RuntimeError as exc:
            if "out of memory" in str(exc).lower():
                raise RuntimeError(
                    "GPU/内存不足。可尝试减小 --target-shape，或降低 --nb-features，"
                    "再重新预处理与训练。"
                ) from exc
            raise

        row = {
            "epoch": float(epoch),
            "loss": float(loss.detach().cpu()),
            "image_loss": float(loss_sim.detach().cpu()),
            "smooth_loss": float(loss_smooth.detach().cpu()),
        }
        history.append(row)
        print(
            f"Epoch {epoch:04d}/{epochs} | "
            f"loss={row['loss']:.6f} image={row['image_loss']:.6f} smooth={row['smooth_loss']:.6f}"
        )

        if epoch == epochs or (save_every > 0 and epoch % save_every == 0):
            save_checkpoint(
                checkpoint_path,
                model=model,
                epoch=epoch,
                input_shape_xyz=tuple(int(v) for v in moving.shape[2:]),
                nb_features=nb_features,
                integration_steps=integration_steps,
                config={
                    "moving_path": str(moving_path),
                    "fixed_path": str(fixed_path),
                    "learning_rate": learning_rate,
                    "smooth_weight": smooth_weight,
                    "image_loss": image_loss,
                    "ncc_window": ncc_window,
                    "ncc_impl": ncc_impl,
                },
                history=history,
            )

    history_path = checkpoint_path.with_suffix(".history.json")
    history_path.write_text(json.dumps(history, ensure_ascii=False, indent=2), encoding="utf-8")
    return history


def save_checkpoint(
    path: str | Path,
    model,
    epoch: int,
    input_shape_xyz: Sequence[int],
    nb_features: Sequence[int],
    integration_steps: int,
    config: Dict,
    history: List[Dict[str, float]],
) -> None:
    require_torch()
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "epoch": int(epoch),
            "input_shape_xyz": tuple(int(v) for v in input_shape_xyz),
            "nb_features": list(int(v) for v in nb_features),
            "integration_steps": int(integration_steps),
            "official_core": {
                "library": "voxelmorph",
                "repo": OFFICIAL_VXM_REPO,
                "commit": OFFICIAL_VXM_COMMIT,
                "model_class": "voxelmorph.nn.models.VxmPairwise",
            },
            "config": config,
            "history": history,
        },
        path,
    )


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="用官方 VoxelMorph 训练 3D 颈部 CT 形变配准模型。")
    parser.add_argument("--moving", required=True, help="预处理后的 moving NIfTI。")
    parser.add_argument("--fixed", required=True, help="预处理后的 fixed NIfTI。")
    parser.add_argument("--checkpoint", default=str(DEFAULT_CHECKPOINT), help="模型权重保存路径。")
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--smooth-weight", type=float, default=0.01)
    parser.add_argument("--image-loss", choices=["ncc", "mse"], default="ncc")
    parser.add_argument("--ncc-window", type=int, default=9)
    parser.add_argument("--ncc-impl", choices=["local", "neurite"], default="local")
    parser.add_argument("--nb-features", type=int, nargs="+", default=[16, 16, 16, 16, 16])
    parser.add_argument("--integration-steps", type=int, default=0, help="0 为普通 dense flow；>0 使用 scaling-and-squaring。")
    parser.add_argument("--device", default="auto")
    parser.add_argument("--no-amp", action="store_true", help="关闭 CUDA AMP 混合精度。")
    parser.add_argument("--save-every", type=int, default=50)
    return parser


def main(argv: Iterable[str] | None = None) -> None:
    args = build_arg_parser().parse_args(argv)
    train_pair(
        moving_path=args.moving,
        fixed_path=args.fixed,
        checkpoint_path=args.checkpoint,
        epochs=args.epochs,
        learning_rate=args.lr,
        smooth_weight=args.smooth_weight,
        image_loss=args.image_loss,
        ncc_window=args.ncc_window,
        ncc_impl=args.ncc_impl,
        nb_features=args.nb_features,
        integration_steps=args.integration_steps,
        device=args.device,
        use_amp=not args.no_amp,
        save_every=args.save_every,
    )


if __name__ == "__main__":
    main()