import os
import re
import json
import glob
import numpy as np
import SimpleITK as sitk
from scipy import ndimage

# Repo root: next to this file's parent; overridable for scratch checkouts.
ROOT = os.environ.get(
    "LONGITUDINAL_ROOT", os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def _load_paths():
    """All path config from config/paths.json ({} when absent/unreadable)."""
    p = os.path.join(ROOT, "config", "paths.json")
    if os.path.exists(p):
        try:
            with open(p) as f:
                v = json.load(f)
            if isinstance(v, dict):
                return v
        except (OSError, ValueError):
            pass
    return {}


PATHS = _load_paths()


def path(name):
    """Named external path from config/paths.json (e.g. 'data', 'lee', 'm6')."""
    try:
        return PATHS[name]
    except KeyError:
        raise KeyError(
            f"missing path {name!r} in config/paths.json (have: {sorted(PATHS)})")


DATA = PATHS.get("data") or os.path.join(ROOT, "data")

# series from non-head (non-brain/non-skull) exams must not enter the dataset
# (includes "head and neck" exams — neck studies are excluded wholesale)
HEAD_EXCL_RE = re.compile(
    r"spine|breast|abdomen|pelvis|chest|extremit|brachial|urograph|prostat|"
    r"\bbody\b|thigh|knee|wrist|ankle|foot\b|hand\b|elbow|shoulder|carotid|neck|"
    r"(?<![a-z])cca(?![a-z])", re.I)


def is_head_series(study_desc, series_desc):
    s = (study_desc or "") + " " + (series_desc or "")
    return not HEAD_EXCL_RE.search(s)


def d(name):
    if name == "data":
        p = DATA
    elif name.startswith("data/"):
        p = os.path.join(DATA, name[len("data/"):])
    else:
        p = os.path.join(ROOT, name)
    os.makedirs(p, exist_ok=True)
    return p


def load_npz_or_nii(path):
    if path.endswith(".npz"):
        return np.load(path, allow_pickle=False)
    return sitk.GetArrayFromImage(sitk.ReadImage(path))


def read_nii(path):
    return sitk.ReadImage(path)


def write_nii(itk_img, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    sitk.WriteImage(itk_img, path, True)


def npy_to_itk(arr, origin=(0, 0, 0), spacing=(1, 1, 1), direction=(1, 0, 0, 0, 1, 0, 0, 0, 1)):
    img = sitk.GetImageFromArray(arr.astype(np.float32 if arr.dtype != np.uint8 else np.uint8))
    img.SetOrigin(tuple(origin))
    img.SetSpacing(tuple(spacing))
    img.SetDirection(tuple(direction))
    return img


def resample_to_spacing(itk_img, spacing, interpol=sitk.sitkLinear):
    arr = sitk.GetArrayFromImage(itk_img)
    out = np.zeros(arr.shape, dtype=np.float32)
    if max(spacing) > 0:
        out = sitk.GetArrayFromImage(
            sitk.Resample(
                itk_img,
                itk_img,
                sitk.Transform(),
                interpol,
                0.0,
                tuple(float(s) for s in spacing),
            )
        )
    return out


def resample_array(arr, spacing_from, spacing_to, interp="linear"):
    h, w, d = arr.shape
    ref = sitk.GetImageFromArray(np.zeros((1, 1, 1), dtype=np.float32))
    src = sitk.GetImageFromArray(arr.astype(np.float32))
    src.SetSpacing(tuple(spacing_from))
    ref2 = sitk.GetImageFromArray(arr.astype(np.float32))
    ref2.SetSpacing(tuple(spacing_from))
    new = sitk.Resample(
        ref2,
        ref2,
        sitk.Transform(),
        sitk.sitkLinear if interp == "linear" else sitk.sitkNearestNeighbor,
        0.0,
        tuple(spacing_to),
    )
    return sitk.GetArrayFromImage(new)


def largest_cc(mask):
    lab, n = ndimage.label(mask)
    if n == 0:
        return mask
    sizes = ndimage.sum(mask, lab, range(1, n + 1))
    return lab == (int(np.argmax(sizes)) + 1)


def head_mask_from_image(arr):
    """Foreground (head + tumor) mask from a T1c volume. Robust to 8-bit or raw MR scaling."""
    a = arr.ravel()
    a = a[a > 0]
    if a.size == 0:
        a = arr.ravel()
    thr = np.percentile(a if a.size else arr.ravel(), 2)
    m = arr > max(thr, np.finfo(arr.dtype).eps)
    m = largest_cc(m)
    return m


def _axis_range(mask, i):
    proj = np.any(mask, axis=tuple(j for j in range(mask.ndim) if j != i))
    idx = np.where(proj)[0]
    return int(idx.min()), int(idx.max()) + 1


def crop_box(arr, label=None, margin=10, cap=216, spacing=1.0):
    """Compute (start, end) per axis for a head crop. Includes label with margin when given."""
    m = head_mask_from_image(arr)
    m = ndimage.binary_dilation(m, iterations=max(1, int(margin / spacing)))
    sl = [_axis_range(m, i) for i in range(m.ndim)]
    sl = [slice(lo, hi) for lo, hi in sl]
    if label is not None and (label > 0).any():
        l = ndimage.binary_dilation(label > 0, iterations=max(1, int(margin / spacing)))
        for i in range(arr.ndim):
            lo, hi = _axis_range(l, i)
            sl[i] = slice(min(sl[i].start, lo), max(sl[i].stop, hi))
    start = [s.start for s in sl]
    end = [s.stop for s in sl]
    # cap each dim to `cap`, centered on head centroid (or label centroid if given)
    cent = ndimage.center_of_mass(m)
    if label is not None and (label > 0).any():
        cent = ndimage.center_of_mass((label > 0) & (m | label))
    for i in range(arr.ndim):
        size = end[i] - start[i]
        if size > cap:
            c = int(round(cent[i]))
            s0 = max(0, min(c - cap // 2, arr.shape[i] - cap))
            start[i], end[i] = s0, s0 + cap
    return [(s, e) for s, e in zip(start, end)]


def normalize_volume(arr, mask):
    a = arr[mask]
    lo, hi = np.percentile(a, 2), np.percentile(a, 98)
    if hi - lo < 1e-6:
        hi = lo + 1.0
    out = (arr - lo) / (hi - lo)
    return np.clip(out, 0.0, 1.0).astype(np.float32), (lo, hi)


def slices(a, b):
    return slice(a, b)


def load_jsonl(path):
    with open(path) as f:
        return [json.loads(l) for l in f if l.strip()]


def save_jsonl(rows, path):
    p = str(path)
    os.makedirs(os.path.dirname(p), exist_ok=True)
    with open(p, "w") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")


def read_nii_arr(path):
    return sitk.GetArrayFromImage(sitk.ReadImage(path))


def read_nii_img(path):
    return sitk.DICOMOrient(sitk.ReadImage(path))


def write_arr(arr, path, itk_img=None):
    """Write array to nii.gz, copying geometry from itk_img when provided (else identity 1mm)."""
    if itk_img is None:
        itk_img = sitk.GetImageFromArray(arr.astype(np.float32))
        itk_img.SetSpacing((1, 1, 1))
    else:
        if arr.ndim == 3 and arr.shape == tuple(itk_img.GetSize()[::-1]):
            img = sitk.GetImageFromArray(arr.astype(np.float32))
            img.CopyInformation(itk_img)
            itk_img = img
        else:
            raise ValueError(f"shape mismatch {arr.shape} vs {itk_img.GetSize()}")
    os.makedirs(os.path.dirname(path), exist_ok=True)
    sitk.WriteImage(itk_img, path, True)


def dice(a, b, eps=1e-7):
    a = a > 0
    b = b > 0
    if a.sum() == 0 and b.sum() == 0:
        return 1.0
    return float(2 * (a & b).sum() / (a.sum() + b.sum() + eps))


def bbox3(mask):
    sl = ndimage.find_objects(ndimage.label(mask)[0])
    idx = np.argwhere(mask)
    if idx.size == 0:
        return None
    lo, hi = idx.min(0), idx.max(0)
    return lo, hi