"""医学图像预处理模块。

预处理步骤：
1. NIfTI 读取。
2. 重采样到各向同性分辨率，例如 1x1x1 mm。
3. 使用颈部软组织/气道窗口进行 HU 截断并归一化到 [0, 1]。
4. 中心裁剪或填充到 VoxelMorph 兼容尺寸，例如 160x192x224。
"""

from __future__ import annotations

import argparse
import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterable, Sequence, Tuple

import numpy as np
from scipy import ndimage

from config import (
    DEFAULT_TARGET_SHAPE,
    DEFAULT_TARGET_SPACING,
    DEFAULT_WINDOW_LEVEL,
    DEFAULT_WINDOW_WIDTH,
)


@dataclass
class PreprocessMeta:
    input_path: str
    output_path: str
    original_shape_xyz: Tuple[int, int, int]
    output_shape_xyz: Tuple[int, int, int]
    original_spacing_xyz: Tuple[float, float, float]
    target_spacing_xyz: Tuple[float, float, float]
    window_width: float
    window_level: float
    crop_or_pad_offset_xyz: Tuple[int, int, int]


def _require_nibabel():
    try:
        import nibabel as nib  # type: ignore
    except Exception as exc:  # pragma: no cover
        raise RuntimeError(
            "缺少 nibabel，无法读取/保存 NIfTI。请先运行: pip install -r requirements.txt"
        ) from exc
    return nib


def ensure_multiple_of_16(shape: Sequence[int]) -> Tuple[int, int, int]:
    shape = tuple(int(v) for v in shape)
    if len(shape) != 3:
        raise ValueError("target_shape 必须包含 3 个维度。")
    bad = [v for v in shape if v <= 0 or v % 16 != 0]
    if bad:
        raise ValueError(f"VoxelMorph 建议三维尺寸均为 16 的倍数，当前非法维度: {bad}")
    return shape  # type: ignore[return-value]


def load_nifti(path: str | Path) -> Tuple[np.ndarray, np.ndarray, Tuple[float, float, float]]:
    """读取 NIfTI，返回 data_xyz、affine 和 spacing_xyz。"""

    nib = _require_nibabel()
    img = nib.load(str(path), mmap=True)
    if len(img.shape) < 3:
        raise ValueError(f"NIfTI 至少应为 3D: {path}")
    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)
    spacing = tuple(float(v) for v in img.header.get_zooms()[:3])
    return data, img.affine.copy(), spacing  # type: ignore[return-value]


def estimate_memory_mb(shape: Sequence[int], dtype=np.float32, copies: int = 3) -> float:
    """估算中间数组内存，copies 用于覆盖 scipy zoom 等临时副本。"""

    voxels = int(np.prod(tuple(int(v) for v in shape)))
    return voxels * np.dtype(dtype).itemsize * copies / (1024**2)


def window_and_normalize(
    data: np.ndarray,
    window_width: float = DEFAULT_WINDOW_WIDTH,
    window_level: float = DEFAULT_WINDOW_LEVEL,
) -> np.ndarray:
    """HU 截断并 Min-Max 归一化。

    W=400/L=40 对颈部软组织和气道边界比较友好：既不会被骨窗拉爆，
    也能保留软组织灰度差异供 NCC 训练使用。
    """

    low = float(window_level) - float(window_width) / 2.0
    high = float(window_level) + float(window_width) / 2.0
    clipped = np.clip(data.astype(np.float32, copy=False), low, high)
    normalized = (clipped - low) / max(high - low, 1e-6)
    return normalized.astype(np.float32, copy=False)


def affine_with_spacing(affine: np.ndarray, target_spacing: Sequence[float]) -> np.ndarray:
    """保持方向不变，只把 affine 三个轴的向量长度改成目标 spacing。"""

    output = affine.copy()
    for axis, spacing in enumerate(target_spacing):
        vector = output[:3, axis]
        length = float(np.linalg.norm(vector))
        if length > 1e-8:
            output[:3, axis] = vector / length * float(spacing)
        else:
            output[:3, axis] = 0
            output[axis, axis] = float(spacing)
    return output


def resample_to_spacing(
    data: np.ndarray,
    affine: np.ndarray,
    current_spacing: Sequence[float],
    target_spacing: Sequence[float] = DEFAULT_TARGET_SPACING,
    order: int = 1,
    max_memory_mb: int = 4096,
) -> Tuple[np.ndarray, np.ndarray]:
    """重采样到目标体素间距。

    data 轴顺序为 (X, Y, Z)。zoom factor = 当前 spacing / 目标 spacing。
    """

    current_spacing = tuple(float(v) for v in current_spacing)
    target_spacing = tuple(float(v) for v in target_spacing)
    zoom_factors = tuple(current_spacing[i] / target_spacing[i] for i in range(3))
    new_shape = tuple(max(1, int(round(data.shape[i] * zoom_factors[i]))) for i in range(3))

    expected_mb = estimate_memory_mb(new_shape, dtype=np.float32, copies=4)
    if expected_mb > max_memory_mb:
        raise MemoryError(
            f"重采样预计需要约 {expected_mb:.1f} MB，超过限制 {max_memory_mb} MB。"
            "可提高 --max-memory-mb，或降低目标尺寸/分辨率。"
        )

    resampled = ndimage.zoom(
        data.astype(np.float32, copy=False),
        zoom=zoom_factors,
        order=order,
        mode="nearest",
        prefilter=(order > 1),
        output=np.float32,
    )
    return resampled.astype(np.float32, copy=False), affine_with_spacing(affine, target_spacing)


def center_crop_or_pad(
    data: np.ndarray,
    target_shape: Sequence[int] = DEFAULT_TARGET_SHAPE,
    pad_value: float = 0.0,
) -> Tuple[np.ndarray, Tuple[int, int, int]]:
    """中心裁剪/填充到固定尺寸。

    返回 offset_xyz，用于修正 affine 原点：
    original_index = output_index + offset。
    """

    target_shape = ensure_multiple_of_16(target_shape)
    output = np.full(target_shape, pad_value, dtype=np.float32)
    input_slices = []
    output_slices = []
    offsets = []

    for axis, target in enumerate(target_shape):
        size = int(data.shape[axis])
        if size >= target:
            start = (size - target) // 2
            input_slices.append(slice(start, start + target))
            output_slices.append(slice(0, target))
            offsets.append(start)
        else:
            pad_before = (target - size) // 2
            input_slices.append(slice(0, size))
            output_slices.append(slice(pad_before, pad_before + size))
            offsets.append(-pad_before)

    output[tuple(output_slices)] = data[tuple(input_slices)]
    return output, tuple(int(v) for v in offsets)  # type: ignore[return-value]


def shift_affine_origin(affine: np.ndarray, offset_xyz: Sequence[int]) -> np.ndarray:
    """根据裁剪/填充偏移修正 NIfTI 原点。"""

    output = affine.copy()
    offset = np.asarray(offset_xyz, dtype=np.float64)
    output[:3, 3] = output[:3, 3] + output[:3, :3] @ offset
    return output


def save_nifti(data_xyz: np.ndarray, affine: np.ndarray, output_path: str | Path) -> None:
    nib = _require_nibabel()
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    img = nib.Nifti1Image(data_xyz.astype(np.float32, copy=False), affine)
    img.header.set_xyzt_units("mm")
    nib.save(img, str(output_path))


def preprocess_array(
    data_xyz: np.ndarray,
    affine: np.ndarray,
    spacing_xyz: Sequence[float],
    target_spacing: Sequence[float] = DEFAULT_TARGET_SPACING,
    target_shape: Sequence[int] = DEFAULT_TARGET_SHAPE,
    window_width: float = DEFAULT_WINDOW_WIDTH,
    window_level: float = DEFAULT_WINDOW_LEVEL,
    normalize_mode: str = "window",
    max_memory_mb: int = 4096,
) -> Tuple[np.ndarray, np.ndarray, Tuple[int, int, int]]:
    """数组级预处理，供 CLI 和 infer.py 共同复用。"""

    if normalize_mode == "window":
        data_xyz = window_and_normalize(data_xyz, window_width=window_width, window_level=window_level)
    elif normalize_mode == "auto":
        finite = data_xyz[np.isfinite(data_xyz)]
        if finite.size and float(np.nanmin(finite)) >= -0.1 and float(np.nanmax(finite)) <= 1.1:
            data_xyz = np.clip(data_xyz.astype(np.float32, copy=False), 0.0, 1.0)
        else:
            data_xyz = window_and_normalize(data_xyz, window_width=window_width, window_level=window_level)
    elif normalize_mode == "none":
        data_xyz = data_xyz.astype(np.float32, copy=False)
    else:
        raise ValueError("normalize_mode 只能是 window/auto/none。")

    resampled, resampled_affine = resample_to_spacing(
        data_xyz,
        affine,
        current_spacing=spacing_xyz,
        target_spacing=target_spacing,
        order=1,
        max_memory_mb=max_memory_mb,
    )
    cropped, offset = center_crop_or_pad(resampled, target_shape=target_shape, pad_value=0.0)
    output_affine = shift_affine_origin(resampled_affine, offset)
    return cropped, output_affine, offset


def preprocess_nifti(
    input_path: str | Path,
    output_path: str | Path,
    target_spacing: Sequence[float] = DEFAULT_TARGET_SPACING,
    target_shape: Sequence[int] = DEFAULT_TARGET_SHAPE,
    window_width: float = DEFAULT_WINDOW_WIDTH,
    window_level: float = DEFAULT_WINDOW_LEVEL,
    normalize_mode: str = "window",
    max_memory_mb: int = 4096,
    metadata_json: str | Path | None = None,
) -> PreprocessMeta:
    """完整 NIfTI 预处理入口。"""

    input_path = Path(input_path)
    output_path = Path(output_path)
    data, affine, spacing = load_nifti(input_path)
    original_shape = tuple(int(v) for v in data.shape[:3])

    processed, output_affine, offset = preprocess_array(
        data,
        affine,
        spacing,
        target_spacing=target_spacing,
        target_shape=target_shape,
        window_width=window_width,
        window_level=window_level,
        normalize_mode=normalize_mode,
        max_memory_mb=max_memory_mb,
    )
    save_nifti(processed, output_affine, output_path)

    meta = PreprocessMeta(
        input_path=str(input_path),
        output_path=str(output_path),
        original_shape_xyz=original_shape,
        output_shape_xyz=tuple(int(v) for v in processed.shape),
        original_spacing_xyz=tuple(float(v) for v in spacing),
        target_spacing_xyz=tuple(float(v) for v in target_spacing),
        window_width=float(window_width),
        window_level=float(window_level),
        crop_or_pad_offset_xyz=offset,
    )

    if metadata_json is None:
        metadata_json = Path(str(output_path).replace(".nii.gz", ".json").replace(".nii", ".json"))
    metadata_json = Path(metadata_json)
    metadata_json.parent.mkdir(parents=True, exist_ok=True)
    metadata_json.write_text(json.dumps(asdict(meta), ensure_ascii=False, indent=2), encoding="utf-8")
    return meta


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="NIfTI 重采样、窗宽窗位归一化、裁剪/填充。")
    parser.add_argument("--input", required=True, help="输入 .nii.gz。")
    parser.add_argument("--output", required=True, help="输出预处理后的 .nii.gz。")
    parser.add_argument("--target-spacing", type=float, nargs=3, default=DEFAULT_TARGET_SPACING, metavar=("X", "Y", "Z"))
    parser.add_argument("--target-shape", type=int, nargs=3, default=DEFAULT_TARGET_SHAPE, metavar=("X", "Y", "Z"))
    parser.add_argument("--window-width", type=float, default=DEFAULT_WINDOW_WIDTH)
    parser.add_argument("--window-level", type=float, default=DEFAULT_WINDOW_LEVEL)
    parser.add_argument("--normalize-mode", choices=["window", "auto", "none"], default="window")
    parser.add_argument("--metadata-json", default=None)
    parser.add_argument("--max-memory-mb", type=int, default=4096)
    return parser


def main(argv: Iterable[str] | None = None) -> None:
    args = build_arg_parser().parse_args(argv)
    meta = preprocess_nifti(
        input_path=args.input,
        output_path=args.output,
        target_spacing=args.target_spacing,
        target_shape=args.target_shape,
        window_width=args.window_width,
        window_level=args.window_level,
        normalize_mode=args.normalize_mode,
        metadata_json=args.metadata_json,
        max_memory_mb=args.max_memory_mb,
    )
    print(json.dumps(asdict(meta), ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()