"""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", "/mnt/b4/xfr/git26/longitudinal")
sys.path.insert(0, ROOT)
from src.common import load_jsonl, save_jsonl, read_nii_arr  # 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 d(name):
    p = os.path.join(ROOT, name)
    os.makedirs(p, exist_ok=True)
    return p


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) ----------------

def make_raw_dataset(rows, channel="0000"):
    """rows: list of {key, pimg, label}. Rebuilds imagesTr/labelsTr symlinks + dataset.json."""
    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["pimg"], 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"))
        os.symlink(os.path.abspath(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) ----------------

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):
    """Returns (mask, cc_frac, vol_mm3) if the positive gates pass, else None."""
    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())
    if cc_frac >= min_cc_frac and vol_lo <= vol_mm3 <= vol_hi:
        return big, cc_frac, vol_mm3
    return None


def neg_frac_bg(p_tumor, vol, tau_neg, min_frac, min_interior=5000):
    interior = vol > 0.02
    if interior.sum() < min_interior:
        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:
#   * tumor centroid relative to the head centroid (crop space, 1mm voxels)
#   * tumor volume ratio
# 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 = vol > 0.02
    if m.sum() < 5000:
        return None
    lab, n = ndimage.label(m)
    if n == 0:
        return None
    sizes = ndimage.sum(m, lab, range(1, n + 1))
    return np.array(ndimage.center_of_mass(lab == (int(np.argmax(sizes)) + 1)))


def rel_tumor_features(mask, vol):
    """Returns (centroid_rel_to_head [z,y,x] in mm, vol_mm3) or None."""
    if int(mask.sum()) == 0:
        return None
    tc = np.array(ndimage.center_of_mass(mask))
    hc = head_centroid(vol)
    rel = tc - hc if hc is not None else None
    return rel, int(mask.sum())


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):
    """Greedy removal of positive timepoints inconsistent with accepted neighbors.

    masks from <out_dir>/<key>_label.nii.gz, volumes from row["pimg"]. 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["pimg"]).astype("float32")
            feats[i] = rel_tumor_features(mask, vol)
        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(ROOT, "data/vols.json")
    return json.load(open(p)) if os.path.exists(p) else {}