"""
Post-processing of PANSEGX probability maps (pan-cancer CT segmentation).

Context
-------
The network (nnU-Net ResEnc-M, `3d_fullres`, patch 64x192x192, TTA disabled)
suffers from a dominant, documented failure mode: **scattered false positives**
(91% of the <0.5 DSC cases are fragmented, 74% over-segmented). There is *no*
official post-processing (no `postprocessing.pkl` in the checkpoint): this step
must therefore be added downstream, on the **full reassembled volume**, never
patch by patch.

Guiding idea
------------
A plain `argmax` (prob >= 0.5) keeps the whole "fog" of false positives, which
precisely crossed 0.5. The right answer is not a higher global threshold (which
erodes the boundary of true lesions -> hurts NSD), but a **connected-component
hysteresis**:

    - a LOW  threshold `t_low`  defines the candidate extent (preserves boundaries);
    - a HIGH threshold `t_high` defines confident "seeds";
    - a component of the low mask is kept only if it contains at least one seed
      (a confident core).

True lesions almost always have a highly confident core; the false-positive fog
does not. Two physical guards are added: minimum volume (in mm³, via the spacing)
and minimum core confidence. No single-lesion assumption is made (multiple
metastatic lesions are legitimate) -> multi-lesion by default.

This module depends only on numpy + scipy (SimpleITK/cc3d optional for I/O and
acceleration). It is deterministic and stateless.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Optional, Sequence

import numpy as np
from scipy import ndimage as ndi

# cc3d speeds up labelling on large volumes; scipy is the universal fallback.
try:
    import cc3d  # type: ignore
    _HAS_CC3D = True
except Exception:  # pragma: no cover
    _HAS_CC3D = False


# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #
@dataclass(frozen=True)
class PostProcConfig:
    """Post-processing parameters.

    Defaults = reasonable starting points, **to be calibrated** on the 50
    validation-public cases (see `calibrate`). All probabilities are those of the
    tumor (foreground) class in [0, 1].
    """
    t_low: float = 0.50          # low threshold: candidate extent (preserves boundaries)
    t_high: float = 0.90         # high threshold: confident seed core (hysteresis)
    min_volume_mm3: float = 50.0 # smaller components -> removed
    min_core_prob: float = 0.90  # minimum peak probability required within a component
    min_mean_prob: float = 0.0   # minimum mean probability (0 = disabled)
    connectivity: int = 26       # 6 | 18 | 26 (3D neighbourhood)
    fill_holes: bool = True      # fills internal holes (DSC gain, boundary-neutral)
    keep_top_k: Optional[int] = None       # keep only the k largest (None = all)
    drop_border_components: bool = False   # drop components touching the FOV border

    def __post_init__(self):
        if not (0.0 <= self.t_low <= self.t_high <= 1.0):
            raise ValueError(f"Requires 0 <= t_low ({self.t_low}) <= t_high ({self.t_high}) <= 1.")
        if self.connectivity not in (6, 18, 26):
            raise ValueError("connectivity must be 6, 18 or 26.")


# --------------------------------------------------------------------------- #
# Low-level helpers
# --------------------------------------------------------------------------- #
def foreground_prob(logits_or_prob: np.ndarray) -> np.ndarray:
    """Extract the tumor probability (Z, Y, X) as float32.

    Accepts:
      - an already-reduced probability map, shape (Z, Y, X);
      - 2-class logits, shape (2, Z, Y, X) or (1, 2, Z, Y, X).
    The 2-class softmax reduces to a stable sigmoid on (l1 - l0).
    """
    a = np.asarray(logits_or_prob)
    if a.ndim == 5 and a.shape[0] == 1:
        a = a[0]
    if a.ndim == 4:
        if a.shape[0] != 2:
            raise ValueError(f"Expected logits with 2 channels, got shape {a.shape}.")
        diff = a[1].astype(np.float64) - a[0].astype(np.float64)
        return _sigmoid(diff).astype(np.float32)
    if a.ndim == 3:
        p = a.astype(np.float32)
        if p.min() < 0.0 or p.max() > 1.0:
            raise ValueError("3D map provided but outside [0,1]: pass logits (2,Z,Y,X) instead.")
        return p
    raise ValueError(f"Unsupported shape: {a.shape}.")


def _sigmoid(x: np.ndarray) -> np.ndarray:
    return np.where(x >= 0, 1.0 / (1.0 + np.exp(-x)),
                    np.exp(x) / (1.0 + np.exp(x)))


def _label(mask: np.ndarray, connectivity: int) -> tuple[np.ndarray, int]:
    """3D connected-component labelling (cc3d if available, else scipy)."""
    if _HAS_CC3D:
        lab = cc3d.connected_components(mask.astype(np.uint8), connectivity=connectivity)
        return lab, int(lab.max())
    # scipy: structure = connectivity rank derived from the requested connectivity
    rank = {6: 1, 18: 2, 26: 3}[connectivity]
    structure = ndi.generate_binary_structure(3, rank)
    lab, n = ndi.label(mask, structure=structure)
    return lab, n


# --------------------------------------------------------------------------- #
# Core: per-component hysteresis + guards
# --------------------------------------------------------------------------- #
def postprocess(
    logits_or_prob: np.ndarray,
    spacing_zyx: Sequence[float],
    config: PostProcConfig = PostProcConfig(),
    return_report: bool = False,
):
    """Apply post-processing and return a binary uint8 mask (Z, Y, X).

    If `return_report=True`, returns `(mask, report)` where `report` is a dict per
    candidate component (volume mm³, peak, mean, kept/rejected + reason) — useful
    for analysis and threshold tuning.
    """
    prob = foreground_prob(logits_or_prob)
    spacing = np.asarray(spacing_zyx, dtype=np.float64)
    if spacing.shape != (3,):
        raise ValueError("spacing_zyx must contain 3 values (z, y, x) in mm.")
    voxel_vol = float(np.prod(spacing))  # mm³ / voxel

    low = prob >= config.t_low
    labels, n = _label(low, config.connectivity)

    if n == 0:
        mask = np.zeros(prob.shape, dtype=np.uint8)
        return (mask, {"components": [], "kept": 0, "config": asdict(config)}) if return_report else mask

    idx = np.arange(1, n + 1)
    # Per-component statistics, fully vectorized.
    counts = np.bincount(labels.ravel(), minlength=n + 1)[1:]          # voxels
    volumes = counts * voxel_vol                                       # mm³
    peak = ndi.maximum(prob, labels, index=idx)                        # core probability
    mean = ndi.mean(prob, labels, index=idx)                           # mean probability

    keep = (
        (peak >= config.t_high)
        & (peak >= config.min_core_prob)
        & (volumes >= config.min_volume_mm3)
        & (mean >= config.min_mean_prob)
    )

    if config.drop_border_components:
        keep &= ~_border_labels_mask(labels, n)

    if config.keep_top_k is not None and keep.sum() > config.keep_top_k:
        # score = "probability mass" (sum of probabilities): favours lesions that
        # are both large and confident, without assuming a single lesion.
        mass = ndi.sum(prob, labels, index=idx)
        order = np.argsort(mass)[::-1]
        allowed = np.zeros(n, dtype=bool)
        allowed[order[keep[order]][: config.keep_top_k]] = True
        keep &= allowed

    keep_ids = idx[keep]
    mask = np.isin(labels, keep_ids)

    if config.fill_holes and mask.any():
        mask = ndi.binary_fill_holes(mask)

    mask = mask.astype(np.uint8)

    if not return_report:
        return mask

    report = {"config": asdict(config), "voxel_vol_mm3": voxel_vol,
              "n_candidates": int(n), "kept": int(keep.sum()), "components": []}
    for i in range(n):
        report["components"].append({
            "label": int(idx[i]),
            "volume_mm3": round(float(volumes[i]), 2),
            "peak_prob": round(float(peak[i]), 4),
            "mean_prob": round(float(mean[i]), 4),
            "kept": bool(keep[i]),
            "reason": _reject_reason(peak[i], mean[i], volumes[i], config) if not keep[i] else "kept",
        })
    return mask, report


def _reject_reason(peak, mean, vol, cfg: PostProcConfig) -> str:
    if peak < max(cfg.t_high, cfg.min_core_prob):
        return "low_core_confidence"
    if vol < cfg.min_volume_mm3:
        return "too_small"
    if mean < cfg.min_mean_prob:
        return "low_mean_confidence"
    return "trimmed_by_top_k_or_border"


def _border_labels_mask(labels: np.ndarray, n: int) -> np.ndarray:
    """Boolean mask (length n) of labels touching a face of the volume."""
    faces = np.concatenate([
        labels[0].ravel(), labels[-1].ravel(),
        labels[:, 0].ravel(), labels[:, -1].ravel(),
        labels[:, :, 0].ravel(), labels[:, :, -1].ravel(),
    ])
    border_ids = np.unique(faces)
    border_ids = border_ids[border_ids > 0]
    out = np.zeros(n, dtype=bool)
    out[border_ids - 1] = True
    return out


# --------------------------------------------------------------------------- #
# Metrics (for calibration)
# --------------------------------------------------------------------------- #
def dice(pred: np.ndarray, gt: np.ndarray) -> float:
    pred = pred.astype(bool); gt = gt.astype(bool)
    denom = pred.sum() + gt.sum()
    if denom == 0:
        return 1.0  # two empty masks: perfect agreement (convention)
    return 2.0 * np.logical_and(pred, gt).sum() / denom


def surface_dice(pred: np.ndarray, gt: np.ndarray, spacing_zyx: Sequence[float],
                 tol_mm: float = 2.0) -> float:
    """NSD (Normalized Surface Dice) at tolerance `tol_mm`.

    Reference implementation (boundaries via erosion, distances via anisotropic
    EDT). Verify against the official FLARE scorer before publishing numbers.
    """
    pred = pred.astype(bool); gt = gt.astype(bool)
    sp = np.asarray(spacing_zyx, dtype=np.float64)
    if pred.sum() == 0 and gt.sum() == 0:
        return 1.0
    if pred.sum() == 0 or gt.sum() == 0:
        return 0.0
    s_pred = pred ^ ndi.binary_erosion(pred)
    s_gt = gt ^ ndi.binary_erosion(gt)
    # distance to the GT boundary / to the pred boundary (EDT on the surface complement)
    dt_gt = ndi.distance_transform_edt(~s_gt, sampling=sp)
    dt_pred = ndi.distance_transform_edt(~s_pred, sampling=sp)
    pred_ok = (dt_gt[s_pred] <= tol_mm).sum()
    gt_ok = (dt_pred[s_gt] <= tol_mm).sum()
    return (pred_ok + gt_ok) / (s_pred.sum() + s_gt.sum())


# --------------------------------------------------------------------------- #
# Calibration on the 50 validation-public cases
# --------------------------------------------------------------------------- #
def calibrate(
    prob_maps: Sequence[np.ndarray],
    gts: Sequence[np.ndarray],
    spacings: Sequence[Sequence[float]],
    grid: Optional[dict] = None,
    metric: str = "dsc",          # "dsc" | "nsd" | "mean"
    base: PostProcConfig = PostProcConfig(),
):
    """Sweep a threshold grid and return (best_config, scores_table).

    `prob_maps` are probability maps (or logits) aligned with `gts` (binary masks)
    and `spacings`. Maximizes the mean metric over the cohort.
    """
    if not (len(prob_maps) == len(gts) == len(spacings)):
        raise ValueError("prob_maps, gts and spacings must have the same length.")
    grid = grid or {
        "t_high": [0.70, 0.80, 0.90, 0.95],
        "min_volume_mm3": [0.0, 50.0, 100.0, 250.0, 500.0],
        "min_core_prob": [0.80, 0.90, 0.95],
    }
    keys = list(grid.keys())
    from itertools import product
    results = []
    best = None
    probs_fg = [foreground_prob(p) for p in prob_maps]  # compute once
    for combo in product(*(grid[k] for k in keys)):
        overrides = dict(zip(keys, combo))
        # min_core_prob cannot exceed t_high in our logic -> align it
        cfg = _with(base, **overrides)
        dscs, nsds = [], []
        for p, gt, sp in zip(probs_fg, gts, spacings):
            m = postprocess(p, sp, cfg)
            dscs.append(dice(m, gt))
            if metric in ("nsd", "mean"):
                nsds.append(surface_dice(m, gt, sp, tol_mm=2.0))
        d = float(np.mean(dscs))
        s = float(np.mean(nsds)) if nsds else float("nan")
        score = {"dsc": d, "nsd": s, "mean": np.nanmean([d, s])}[metric]
        row = {**overrides, "dsc": round(d, 4), "nsd": round(s, 4), "score": round(score, 4)}
        results.append(row)
        if best is None or score > best[0]:
            best = (score, cfg)
    results.sort(key=lambda r: r["score"], reverse=True)
    return best[1], results


def _with(cfg: PostProcConfig, **overrides) -> PostProcConfig:
    d = asdict(cfg); d.update(overrides)
    # consistency: the required core never exceeds the high threshold
    d["min_core_prob"] = min(d["min_core_prob"], d["t_high"])
    return PostProcConfig(**d)


# --------------------------------------------------------------------------- #
# NIfTI I/O + CLI
# --------------------------------------------------------------------------- #
def _read_nifti(path: str):
    """Return (array Z,Y,X ; spacing z,y,x). Handles a 3D mask or prob/logits."""
    import SimpleITK as sitk
    img = sitk.ReadImage(str(path))
    arr = sitk.GetArrayFromImage(img)              # (Z, Y, X) for a scalar image
    spacing = tuple(reversed(img.GetSpacing()))    # sitk returns (x, y, z)
    return arr, spacing, img


def _write_nifti(mask: np.ndarray, ref_img, out_path: str):
    import SimpleITK as sitk
    out = sitk.GetImageFromArray(mask.astype(np.uint8))
    out.CopyInformation(ref_img)
    sitk.WriteImage(out, str(out_path))


def main():
    p = argparse.ArgumentParser(description="PANSEGX post-processing (hysteresis + physical guards).")
    p.add_argument("--input", required=True, help="NIfTI: tumor probability map (Z,Y,X) OR logits.")
    p.add_argument("--output", required=True, help="Output NIfTI (binary mask).")
    p.add_argument("--t-low", type=float, default=PostProcConfig.t_low)
    p.add_argument("--t-high", type=float, default=PostProcConfig.t_high)
    p.add_argument("--min-volume-mm3", type=float, default=PostProcConfig.min_volume_mm3)
    p.add_argument("--min-core-prob", type=float, default=PostProcConfig.min_core_prob)
    p.add_argument("--min-mean-prob", type=float, default=PostProcConfig.min_mean_prob)
    p.add_argument("--connectivity", type=int, choices=(6, 18, 26), default=PostProcConfig.connectivity)
    p.add_argument("--no-fill-holes", action="store_true")
    p.add_argument("--keep-top-k", type=int, default=None)
    p.add_argument("--drop-border", action="store_true")
    p.add_argument("--report", action="store_true", help="Print the per-component report.")
    args = p.parse_args()

    cfg = PostProcConfig(
        t_low=args.t_low, t_high=args.t_high, min_volume_mm3=args.min_volume_mm3,
        min_core_prob=args.min_core_prob, min_mean_prob=args.min_mean_prob,
        connectivity=args.connectivity, fill_holes=not args.no_fill_holes,
        keep_top_k=args.keep_top_k, drop_border_components=args.drop_border,
    )
    arr, spacing, ref = _read_nifti(args.input)
    result = postprocess(arr, spacing, cfg, return_report=args.report)
    mask, report = result if args.report else (result, None)
    _write_nifti(mask, ref, args.output)
    kept = int(mask.max() and _label(mask, cfg.connectivity)[1])
    print(f"[postprocess] spacing(z,y,x)={spacing} -> {kept} component(s) kept. Written: {args.output}")
    if report is not None:
        import json
        print(json.dumps(report, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()