"""Shared helpers for the nnU-Net iterative pseudo-labeling pipeline (scripts_nnu/).

Runs against the installed nnunetv2 2.8.1 build (checkpoint key `network_weights`,
probabilities saved as <case>.npz channel-first, folds in splits_final.json).
"""
import os
import sys
import time
import json
import shutil
import subprocess

ROOT = os.environ.get(
    "LONGITUDINAL_ROOT", os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, ROOT)
from src.common import DATA, d, load_jsonl, save_jsonl, read_nii_arr, head_mask_from_image, largest_cc  # noqa: E402

import numpy as np
import SimpleITK as sitk
from scipy import ndimage

DS_ID = 210
DS_NAME = f"Dataset{DS_ID:03d}_NTUH_T1C_PL"
TRAINER = "NTUHLPLTrainer"
CONFIG = "3d_fullres"
FOLD = 0
PLANS = "nnUNetPlans"

SCRIPTS_NNU = os.path.join(ROOT, "scripts_nnu")
TRAINERS_DIR = os.path.join(SCRIPTS_NNU, "trainers")


def nnu_root():
    return d(os.path.join("nnu"))


def raw_ds():
    return os.path.join(nnu_root(), "raw", DS_NAME)


def preproc_ds():
    return os.path.join(nnu_root(), "preprocessed", DS_NAME)


def results_root():
    return os.path.join(nnu_root(), "results")


def model_folder():
    return os.path.join(results_root(), DS_NAME, f"{TRAINER}__{PLANS}__{CONFIG}")


def fold_dir():
    return os.path.join(model_folder(), f"fold_{FOLD}")


def best_ckpt():
    return os.path.join(fold_dir(), "checkpoint_best.pth")


def final_ckpt():
    return os.path.join(fold_dir(), "checkpoint_final.pth")


def nnu_env(epoch=None, lr=None, warmstart=None):
    env = dict(os.environ)
    env["nnUNet_raw"] = os.path.join(nnu_root(), "raw")
    env["nnUNet_preprocessed"] = os.path.join(nnu_root(), "preprocessed")
    env["nnUNet_results"] = results_root()
    env["nnUNet_extTrainer"] = TRAINERS_DIR
    env.pop("NNU_PL_EPOCHS", None)
    env.pop("NNU_PL_LR", None)
    env.pop("NNU_PL_WARMSTART", None)
    env.pop("CUDA_VISIBLE_DEVICES", None)
    if epoch is not None:
        env["NNU_PL_EPOCHS"] = str(epoch)
    if lr is not None:
        env["NNU_PL_LR"] = repr(float(lr))
    if warmstart:
        env["NNU_PL_WARMSTART"] = str(warmstart)
    return env


def run(cmd, log, env=None, retries=2):
    os.makedirs(os.path.dirname(os.path.abspath(log)), exist_ok=True)
    cmd_s = " ".join(str(c) for c in cmd)
    for attempt in range(retries + 1):
        try:
            with open(log, "a") as f:
                f.write(f"$ (attempt {attempt + 1}) " + cmd_s + "\n")
                f.flush()
                print(f"$ (attempt {attempt + 1}) " + cmd_s, flush=True)
                subprocess.run(cmd, cwd=ROOT, stdout=f, stderr=subprocess.STDOUT,
                               env=env, check=True)
            return
        except subprocess.CalledProcessError:
            if attempt == retries:
                raise
            print(f"[nnu] command failed, retrying in 60s: {cmd_s}", flush=True)
            time.sleep(60)


def train_cmd(gpus):
    return ["nnUNetv2_train", str(DS_ID), CONFIG, str(FOLD), "-tr", TRAINER,
            "-p", PLANS, "-num_gpus", str(gpus)]


def plan_preprocess_cmd(npp=8):
    return ["nnUNetv2_plan_and_preprocess", "-d", str(DS_ID), "-c", CONFIG,
            "-np", str(npp), "--clean", "--no_pbar"]


def predict_cmd(in_dir, out_dir, gpus, part, chk="checkpoint_best.pth", tta=True,
                npp=2, nps=2):
    cmd = ["nnUNetv2_predict", "-i", str(in_dir), "-o", str(out_dir), "-d", str(DS_ID),
           "-c", CONFIG, "-tr", TRAINER, "-p", PLANS, "-f", str(FOLD), "-chk", chk,
           "--save_probabilities", "--continue_prediction", "--disable_progress_bar",
           "-num_parts", str(gpus), "-part_id", str(part),
           "-npp", str(npp), "-nps", str(nps)]
    if not tta:
        cmd.append("--disable_tta")
    return cmd


# ---------------- raw dataset (nnU-Net format) ----------------
#
# Pipeline B feeds nnU-Net the NATIVE volumes (no resampling / cropping /
# intensity normalization beforehand): nnUNetv2_plan_and_preprocess does all of
# that itself (nonzero-bbox crop, resample to the plan's target spacing,
# per-channel normalization, .b2nd storage).

def same_frame(a, b):
    return (a.GetSize() == b.GetSize() and a.GetSpacing() == b.GetSpacing()
            and a.GetOrigin() == b.GetOrigin() and a.GetDirection() == b.GetDirection())


def native_label_target(img_path, lab_path, link_path):
    """Label target for a raw-dataset row: 0/1 on the image's native grid.

    Symlinks the label when it already sits on the image grid with values
    {0, 1}; otherwise writes a converted copy to link_path (nearest-neighbor
    warp to the image grid + binarization to 0/1). This is label fix-up only —
    no common-grid resampling, cropping, or normalization.
    """
    img_itk = sitk.ReadImage(img_path)
    lab_itk = sitk.ReadImage(lab_path)
    if same_frame(img_itk, lab_itk):
        vals = set(np.unique(sitk.GetArrayFromImage(lab_itk)).tolist())
        if vals <= {0, 1}:
            os.symlink(os.path.abspath(lab_path), link_path)
            return
    warped = sitk.Resample(lab_itk, img_itk, sitk.Transform(), sitk.sitkNearestNeighbor, 0.0)
    out = sitk.GetImageFromArray((sitk.GetArrayFromImage(warped) > 0).astype(np.uint8))
    out.CopyInformation(img_itk)
    os.makedirs(os.path.dirname(link_path), exist_ok=True)
    sitk.WriteImage(out, link_path, True)


def make_raw_dataset(rows, channel="0000"):
    """rows: list of {key, img, label} — img/label are native (unprocessed) NIfTI
    paths. Rebuilds imagesTr/labelsTr + dataset.json. Images are symlinked
    as-is; labels via native_label_target. See section note."""
    img_dir = os.path.join(raw_ds(), "imagesTr")
    lab_dir = os.path.join(raw_ds(), "labelsTr")
    shutil.rmtree(img_dir, ignore_errors=True)
    shutil.rmtree(lab_dir, ignore_errors=True)
    os.makedirs(img_dir, exist_ok=True)
    os.makedirs(lab_dir, exist_ok=True)
    seen = set()
    for r in rows:
        key, img, lab = r["key"], r["img"], r["label"]
        if key in seen:
            raise ValueError(f"duplicate case id in dataset: {key}")
        seen.add(key)
        if not os.path.exists(img):
            raise FileNotFoundError(f"image missing for {key}: {img}")
        if not os.path.exists(lab):
            raise FileNotFoundError(f"label missing for {key}: {lab}")
        os.symlink(os.path.abspath(img), os.path.join(img_dir, f"{key}_{channel}.nii.gz"))
        native_label_target(img, lab, os.path.join(lab_dir, f"{key}.nii.gz"))
    ds_json = {"channel_names": {"0": "t1c"},
               "labels": {"background": 0, "tumor": 1},
               "numTraining": len(seen),
               "file_ending": ".nii.gz"}
    with open(os.path.join(raw_ds(), "dataset.json"), "w") as f:
        json.dump(ds_json, f, indent=1)
    return len(seen)


def dataset_case_keys():
    img_dir = os.path.join(raw_ds(), "imagesTr")
    out = []
    if os.path.isdir(img_dir):
        for fn in sorted(os.listdir(img_dir)):
            if fn.endswith("_0000.nii.gz"):
                out.append(fn[:-len("_0000.nii.gz")])
    return out


def write_splits(train_keys, val_keys):
    val_keys = sorted(set(val_keys))
    train_keys = sorted(set(k for k in train_keys if k not in val_keys))
    folds = [{"train": train_keys, "val": val_keys}] * 5
    with open(os.path.join(preproc_ds(), "splits_final.json"), "w") as f:
        json.dump(folds, f)
    return len(train_keys), len(val_keys)


# ---------------- label writing / selection (same gates as scripts/06) ----------------
#
# Prediction outputs are ALREADY on each case's native grid: nnUNetv2_predict
# (--save_probabilities) resamples tumor probabilities and the argmax seg back
# to the original input shape. Gates below therefore run directly on the native
# volume; physical mm is recovered from the native voxel spacing.

def write_label(arr_u01, path, ref_itk):
    img = sitk.GetImageFromArray(arr_u01.astype(np.uint8))
    img.CopyInformation(ref_itk)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    sitk.WriteImage(img, path, True)


def pos_mask(p_tumor, tau_pos, min_cc_frac, vol_lo, vol_hi, voxel_vol_mm3=1.0):
    """Returns (mask, cc_frac, vol_mm3) if the positive gates pass, else None.

    voxel_vol_mm3 is the native voxel volume in mm3 (Product of NIfTI spacing);
    1.0 for 1mm processed volumes (Pipelines A/C)."""
    m = p_tumor >= tau_pos
    if m.sum() == 0:
        return None
    m = ndimage.median_filter(m, size=(3, 3, 3))
    if m.sum() == 0:
        return None
    lab, n = ndimage.label(m)
    sizes = ndimage.sum(m, lab, range(1, n + 1))
    big = (lab == (int(np.argmax(sizes)) + 1)).astype(np.uint8)
    cc_frac = float(big.sum()) / float(m.sum())
    vol_mm3 = int(big.sum()) * voxel_vol_mm3
    if cc_frac >= min_cc_frac and vol_lo <= vol_mm3 <= vol_hi:
        return big, cc_frac, vol_mm3
    return None


def _interior_mask(vol, min_interior=5000):
    """Head-interior mask. Processed [0,1] percentile-normalized volumes keep
    the legacy `vol > 0.02` rule; native raw-intensity volumes use the robust
    percentile / largest-CC head mask (see src/common.head_mask_from_image)."""
    if float(vol.max()) <= 1.0 + 1e-6:
        m = vol > 0.02
    else:
        m = head_mask_from_image(vol)
    return m if m.sum() >= min_interior else None


def neg_frac_bg(p_tumor, vol, tau_neg, min_frac, min_interior=5000):
    interior = _interior_mask(vol, min_interior)
    if interior is None:
        return None
    frac = float(((1.0 - p_tumor)[interior] >= tau_neg).mean())
    return frac if frac >= min_frac else None


def tumor_probs_from_npz(npz_path):
    """<case>.npz probabilities: channel-first (C,H,W,D), C=2 (bg, tumor)."""
    probs = np.load(npz_path, allow_pickle=False)["probabilities"]
    if probs.shape[0] != 2:
        raise ValueError(f"unexpected probability channels {probs.shape} in {npz_path}")
    return probs[1]


# ---------------- longitudinal consistency (head-relative) ----------------
#
# The absolute patient-space resampling used by scripts/06 cannot work for this
# data: each acquisition has its own scanner/patient coordinate system (table
# offset, head orientation), so true labels from two visits of the same subject
# resample to ~0 dice. Instead we gate on frame-invariant features, computed in
# each case's own grid with physical mm (native spacing for Pipeline B, 1mm for
# A/C):
#   * tumor centroid relative to the head centroid (offset in mm)
#   * tumor volume ratio (mm3)
# Two timepoints are compatible if they agree with at least one accepted
# neighbor on both; greedy removal until stable (same scheme as scripts/06).

def head_centroid(vol):
    m = _interior_mask(vol)
    if m is None:
        return None
    m = largest_cc(m)
    if int(m.sum()) == 0:
        return None
    return np.array(ndimage.center_of_mass(m))


def rel_tumor_features(mask, vol, spacing=(1.0, 1.0, 1.0)):
    """Returns (centroid_rel_to_head [z,y,x] in mm, vol_mm3) or None.

    spacing is the NIfTI voxel spacing in array order (z, y, x); (1, 1, 1) for
    1mm processed volumes (Pipelines A/C)."""
    if int(mask.sum()) == 0:
        return None
    spacing = np.asarray(spacing, dtype=np.float64)
    tc = np.array(ndimage.center_of_mass(mask))
    hc = head_centroid(vol)
    rel = (tc - hc) * spacing if hc is not None else None
    return rel, int(mask.sum()) * float(spacing.prod())


def _compatible(fa, fb, max_rel_dist_mm, vol_ratio_max):
    ra, va = fa
    rb, vb = fb
    if ra is not None and rb is not None and np.linalg.norm(ra - rb) > max_rel_dist_mm:
        return False
    if va > 0 and vb > 0:
        ratio = max(va / vb, vb / va)
        if ratio > vol_ratio_max:
            return False
    return True


def consistency_filter(rows, out_dir, max_rel_dist_mm=40.0, vol_ratio_max=10.0, image_key="pimg"):
    """Greedy removal of positive timepoints inconsistent with accepted neighbors.

    masks from <out_dir>/<key>_label.nii.gz, volumes + native spacing from
    row[image_key] ("pimg": 1mm processed, Pipelines A/C; "img": native,
    Pipeline B). Returns n_rejected.
    """
    by_subj = {}
    for r in rows:
        if r["role"] == "pos":
            by_subj.setdefault(r["subject"], []).append(r)
    rejected = 0
    for subj, tps in by_subj.items():
        if len(tps) < 2:
            continue
        tps.sort(key=lambda x: (x.get("date") or ""))
        feats = {}
        for i, r in enumerate(tps):
            mask = read_nii_arr(os.path.join(out_dir, r["key"] + "_label.nii.gz")).astype(bool)
            vol = read_nii_arr(r[image_key]).astype("float32")
            if mask.shape != vol.shape:
                raise ValueError(f"mask shape {mask.shape} != volume {r[image_key]} shape {vol.shape} for {r['key']}")
            sp = np.asarray(sitk.ReadImage(r[image_key]).GetSpacing()[::-1], dtype=np.float64)
            feats[i] = rel_tumor_features(mask, vol, sp)
        accepted = list(range(len(tps)))
        while True:
            changed = False
            for ai in list(accepted):
                neigh = [b for b in (ai - 1, ai + 1) if b in accepted]
                if not neigh or feats[ai] is None:
                    continue
                if not any(_compatible(feats[ai], feats[b], max_rel_dist_mm, vol_ratio_max)
                           for b in neigh if feats[b] is not None):
                    tps[ai]["role"] = "rejected"
                    accepted.remove(ai)
                    rejected += 1
                    changed = True
                    break
            if not changed:
                break
    return rejected


def load_voxel_stats():
    p = os.path.join(DATA, "vols.json")
    return json.load(open(p)) if os.path.exists(p) else {}