import spaces  # must be imported before torch

import os
import tempfile
import threading
import time
import warnings
from collections import OrderedDict
from pathlib import Path

import gradio as gr
import numpy as np
import SimpleITK as sitk
import torch
from huggingface_hub import snapshot_download
from PIL import Image, ImageDraw

from batchgenerators.utilities.file_and_folder_operations import load_json
from longiseg.inference.predict_from_raw_data_longi import LongiSegPredictor
from longiseg.inference.tracking_inference import predict_patch

warnings.filterwarnings("ignore", category=UserWarning)

MODEL_ID = "ykirchhoff/LongiTrack"
VERSIONS = ["1.1", "1.0"]
DEVICE = torch.device("cuda")

# --------------------------------------------------------------------------------------
# Model loading (module scope, eager on "cuda" — ZeroGPU intercepts and streams weights)
# --------------------------------------------------------------------------------------
repo_dir = Path(snapshot_download(MODEL_ID, allow_patterns=[f"LongiTrack_v{v}/*" for v in VERSIONS]))

PREDICTORS = {}
DATASET_JSONS = {}
for v in VERSIONS:
    model_path = repo_dir / f"LongiTrack_v{v}"
    predictor = LongiSegPredictor(tile_step_size=0.5, use_gaussian=True, use_mirroring=False,
                                  perform_everything_on_device=True, device=DEVICE,
                                  verbose=False, verbose_preprocessing=False, allow_tqdm=False)
    predictor.initialize_from_trained_model_folder(str(model_path), use_folds=None)
    # identical to longiseg.inference.tracking_inference.predict(): torch resampling on device
    for fn, is_seg, memefficient in (("data", False, False), ("seg", True, True), ("probabilities", False, False)):
        predictor.configuration_manager.configuration[f"resampling_fn_{fn}"] = "resample_torch_fornnunet"
        predictor.configuration_manager.configuration[f"resampling_fn_{fn}_kwargs"] = {
            "is_seg": is_seg, "force_separate_z": False,
            "memefficient_seg_resampling": memefficient, "device": DEVICE,
        }
    predictor.network = predictor.network.to(DEVICE).eval()
    PREDICTORS[v] = predictor
    DATASET_JSONS[v] = load_json(str(model_path / "dataset.json"))
print("LongiTrack models loaded:", list(PREDICTORS))

# --------------------------------------------------------------------------------------
# Volume helpers (CPU)
# --------------------------------------------------------------------------------------
_VOL_CACHE: "OrderedDict[str, np.ndarray]" = OrderedDict()
_VOL_LOCK = threading.Lock()
WL, WW = 40, 400  # abdominal soft-tissue window


def _path(f):
    if f is None:
        return None
    return f if isinstance(f, str) else getattr(f, "name", None) or getattr(f, "path", None)


def load_volume(path: str) -> np.ndarray:
    """Return the volume as a numpy array in (z, y, x) order (cached)."""
    with _VOL_LOCK:
        if path in _VOL_CACHE:
            _VOL_CACHE.move_to_end(path)
            return _VOL_CACHE[path]
    arr = sitk.GetArrayFromImage(sitk.ReadImage(path))
    if arr.ndim == 4:
        arr = arr[..., 0]
    with _VOL_LOCK:
        _VOL_CACHE[path] = arr
        while len(_VOL_CACHE) > 6:
            _VOL_CACHE.popitem(last=False)
    return arr


def ct_to_rgb(slice2d: np.ndarray) -> np.ndarray:
    lo, hi = WL - WW / 2, WL + WW / 2
    g = np.clip((slice2d.astype(np.float32) - lo) / (hi - lo), 0, 1)
    g = (g * 255).astype(np.uint8)
    return np.stack([g, g, g], -1)


def parse_point(s: str):
    try:
        vals = [float(v) for v in str(s).replace("(", " ").replace(")", " ").replace(",", " ").split()]
        if len(vals) != 3:
            return None
        return [int(round(v)) for v in vals]
    except Exception:
        return None


def render_slice(ct_path, z, point_str="", mask_path=None, color=(255, 60, 60), mask_color=(0, 220, 255)):
    ct_path = _path(ct_path)
    if ct_path is None:
        return None
    vol = load_volume(ct_path)
    z = int(np.clip(int(z), 0, vol.shape[0] - 1))
    rgb = ct_to_rgb(vol[z]).astype(np.float32)
    mask_path = _path(mask_path)
    if mask_path:
        m = load_volume(mask_path)
        if m.shape == vol.shape:
            ms = m[z] > 0
            rgb[ms] = 0.55 * rgb[ms] + 0.45 * np.array(mask_color, np.float32)
    img = Image.fromarray(rgb.astype(np.uint8))
    p = parse_point(point_str)
    if p is not None:
        d = ImageDraw.Draw(img)
        x, y, pz = p
        r = 7 if pz == z else 4
        w = 2 if pz == z else 1
        d.line([(x - r - 4, y), (x - 2, y)], fill=color, width=w)
        d.line([(x + 2, y), (x + r + 4, y)], fill=color, width=w)
        d.line([(x, y - r - 4), (x, y - 2)], fill=color, width=w)
        d.line([(x, y + 2), (x, y + r + 4)], fill=color, width=w)
    return img


def on_ct_upload(ct_file, point_str):
    """When a CT volume is loaded, set up the slice slider and show a slice."""
    p = _path(ct_file)
    if p is None:
        return gr.update(value=None), gr.update(maximum=1, value=0)
    vol = load_volume(p)
    pt = parse_point(point_str)
    z = pt[2] if pt is not None and 0 <= pt[2] < vol.shape[0] else vol.shape[0] // 2
    return render_slice(p, z, point_str), gr.update(minimum=0, maximum=vol.shape[0] - 1, value=z, step=1)


def on_click(ct_file, z, evt: gr.SelectData):
    """Clicking on the slice places a point (x, y, z) in voxel coordinates."""
    if _path(ct_file) is None:
        return gr.update(), gr.update()
    x, y = int(evt.index[0]), int(evt.index[1])
    s = f"{x}, {y}, {int(z)}"
    return s, render_slice(ct_file, z, s)


def on_point_text(ct_file, point_str):
    """When a point is typed (or set by an example), jump to its slice."""
    p = _path(ct_file)
    pt = parse_point(point_str)
    if p is None:
        return gr.update(), gr.update()
    vol = load_volume(p)
    if pt is None or not (0 <= pt[2] < vol.shape[0]):
        return gr.update(), gr.update()
    return render_slice(p, pt[2], point_str), gr.update(value=pt[2])


# --------------------------------------------------------------------------------------
# Inference
# --------------------------------------------------------------------------------------
def make_mesh(mask: np.ndarray, spacing_zyx, out_path: str):
    from skimage import measure
    import trimesh

    m = np.pad(mask.astype(np.uint8), 1)
    verts, faces, normals, _ = measure.marching_cubes(m.astype(np.float32), 0.5, spacing=tuple(spacing_zyx))
    verts = verts[:, ::-1].copy()  # (z,y,x) -> (x,y,z) in mm
    verts -= verts.mean(0)
    faces = faces[:, ::-1]
    mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=True)
    try:
        trimesh.smoothing.filter_taubin(mesh, iterations=10)
    except Exception:
        pass
    # rotate so that the head direction (-z in LPS index space) points up in the viewer
    mesh.apply_transform(trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0]))
    mesh.visual.face_colors = [0, 200, 255, 255]
    mesh.export(out_path)


def _estimate_duration(baseline_ct=None, followup_ct=None, *args, **kwargs):
    """GPU time scales with the voxel count (reading + resampling); ~7 s measured for 2x47M voxels."""
    total = 0
    for f in (baseline_ct, followup_ct):
        try:
            r = sitk.ImageFileReader()
            r.SetFileName(_path(f))
            r.ReadImageInformation()
            total += int(np.prod(r.GetSize()))
        except Exception:
            total += 100_000_000
    return int(min(120, 8 + 0.1 * total / 1e6))


@spaces.GPU(duration=_estimate_duration)
def track_lesion(baseline_ct, followup_ct, baseline_mask=None, baseline_point: str = "",
                 followup_point: str = "", version: str = "1.1"):
    """Segment one lesion in a follow-up CT, using its baseline appearance as longitudinal context.

    Runs LongiTrack (verified / point-prompted tracking): the lesion is identified by a point in the
    baseline scan and a (clinician-verified) point in the follow-up scan; an optional baseline lesion
    mask provides the baseline lesion appearance.

    Args:
        baseline_ct: baseline CT volume (NIfTI / any SimpleITK-readable 3D image).
        followup_ct: follow-up CT volume of the same patient.
        baseline_mask: optional baseline lesion segmentation (same geometry as baseline CT).
            The label under the baseline point is used; if that is background, all foreground is used.
        baseline_point: "x, y, z" voxel coordinates of the lesion in the baseline CT.
        followup_point: "x, y, z" voxel coordinates of the lesion in the follow-up CT.
        version: LongiTrack model version ("1.1" or "1.0").

    Returns:
        Overlay image of the central lesion slice, NIfTI mask file, 3D mesh (.glb) and a summary.
    """
    t0 = time.perf_counter()
    bl_path, fu_path, mask_path = _path(baseline_ct), _path(followup_ct), _path(baseline_mask)
    if bl_path is None or fu_path is None:
        raise gr.Error("Please provide both a baseline and a follow-up CT volume.")
    bl_pt, fu_pt = parse_point(baseline_point), parse_point(followup_point)
    if bl_pt is None or fu_pt is None:
        raise gr.Error("Please place a point in both scans (click on the slices), formatted as 'x, y, z'.")
    if version not in PREDICTORS:
        version = "1.1"
    predictor = PREDICTORS[version]
    dataset_json = DATASET_JSONS[version]
    cm, pm = predictor.configuration_manager, predictor.plans_manager

    # optional baseline mask -> binary mask of the selected lesion (label 1)
    tmpdir = tempfile.mkdtemp()
    seg_file = None
    bl_lesion = 1
    if mask_path:
        mimg = sitk.ReadImage(mask_path)
        marr = sitk.GetArrayFromImage(mimg)
        x, y, z = bl_pt
        if marr.ndim == 3 and all(0 <= c < s for c, s in zip((z, y, x), marr.shape)):
            lab = marr[z, y, x]
            sel = (marr == lab) if lab != 0 else (marr > 0)
            out = sitk.GetImageFromArray(sel.astype(np.uint8))
            out.CopyInformation(mimg)
            seg_file = os.path.join(tmpdir, "bl_seg.nii.gz")
            sitk.WriteImage(out, seg_file)

    preprocessor = cm.preprocessor_class(verbose=False)
    bl_data, bl_seg, bl_props = preprocessor.run_case([bl_path], seg_file, pm, cm, dataset_json)
    fu_data, _, fu_props = preprocessor.run_case([fu_path], None, pm, cm, dataset_json)
    bl_data, fu_data = torch.from_numpy(bl_data), torch.from_numpy(fu_data)
    bl_seg = torch.from_numpy(bl_seg) if bl_seg is not None else None

    target_spacing = cm.spacing
    tf = pm.transpose_forward
    bl_spacing = [bl_props["spacing"][i] for i in tf]
    fu_spacing = [fu_props["spacing"][i] for i in tf]
    # points are (x, y, z) in original voxel space -> resampled (z, y, x), as in tracking_inference.py
    bl_p = [int(p * bl_spacing[i] / target_spacing[i]) for i, p in enumerate(bl_pt[::-1])]
    fu_p = [int(p * fu_spacing[i] / target_spacing[i]) for i, p in enumerate(fu_pt[::-1])]
    bl_p = [int(np.clip(p, 0, s - 1)) for p, s in zip(bl_p, bl_data.shape[1:])]
    fu_p = [int(np.clip(p, 0, s - 1)) for p, s in zip(fu_p, fu_data.shape[1:])]

    with torch.no_grad():
        prediction = predict_patch(bl_data, bl_seg, bl_p, bl_lesion, fu_data, fu_p, predictor,
                                   cm.patch_size, DEVICE)
        prediction = cm.resampling_fn_probabilities(prediction, fu_props["shape_after_cropping_and_before_resampling"],
                                                    target_spacing, fu_spacing)
        mask = (prediction[1] > 0.5).to(torch.uint8).cpu().numpy()
    t_inf = time.perf_counter() - t0

    out_nii = os.path.join(tmpdir, "longitrack_followup_lesion.nii.gz")
    pm.image_reader_writer_class().write_seg(mask, out_nii, fu_props)

    # visualisation
    sp_xyz = fu_props["sitk_stuff"]["spacing"]
    spacing_zyx = sp_xyz[::-1]
    nvox = int(mask.sum())
    vol_ml = nvox * float(np.prod(sp_xyz)) / 1000.0
    if nvox > 0:
        zc = int(np.round(np.nonzero(mask.any(axis=(1, 2)))[0].mean()))
        if mask[fu_pt[2]].any() if 0 <= fu_pt[2] < mask.shape[0] else False:
            zc = fu_pt[2]
    else:
        zc = int(np.clip(fu_pt[2], 0, mask.shape[0] - 1))
    vol = load_volume(fu_path)
    rgb = ct_to_rgb(vol[zc]).astype(np.float32)
    ms = mask[zc] > 0
    rgb[ms] = 0.55 * rgb[ms] + 0.45 * np.array([0, 220, 255], np.float32)
    overlay = Image.fromarray(rgb.astype(np.uint8))

    mesh_path = None
    if nvox > 10:
        mesh_path = os.path.join(tmpdir, "longitrack_lesion.glb")
        try:
            make_mesh(mask, spacing_zyx, mesh_path)
        except Exception as e:  # noqa
            print("mesh export failed:", e)
            mesh_path = None
    if mesh_path is None:
        import trimesh
        mesh_path = os.path.join(tmpdir, "longitrack_lesion.glb")
        trimesh.creation.icosphere(radius=0.5).export(mesh_path)

    info = (f"**Follow-up lesion volume:** {vol_ml:.2f} mL ({nvox} voxels)  \n"
            f"**Model:** LongiTrack v{version} · baseline mask: {'yes' if seg_file else 'no (points only)'}  \n"
            f"**Shown slice:** z = {zc} · inference time: {t_inf:.1f} s")
    if nvox == 0:
        info = "⚠️ No lesion was segmented at the given follow-up point.  \n" + info
    print(f"[longitrack] v{version} voxels={nvox} vol={vol_ml:.2f}mL time={t_inf:.2f}s")
    return overlay, out_nii, mesh_path, info


def on_result(followup_ct, mask_file):
    """After a prediction, sync the result slice browser to the lesion."""
    fu, mp = _path(followup_ct), _path(mask_file)
    if fu is None or mp is None:
        return gr.update(), gr.update()
    m = load_volume(mp)
    zs = np.nonzero(m.any(axis=(1, 2)))[0]
    z = int(np.round(zs.mean())) if len(zs) else m.shape[0] // 2
    return render_slice(fu, z, "", mp), gr.update(minimum=0, maximum=m.shape[0] - 1, value=z, step=1)


def browse_result(followup_ct, mask_file, z):
    if _path(followup_ct) is None:
        return gr.update()
    return render_slice(followup_ct, z, "", mask_file)


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1250px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EX = "examples"
EXAMPLES = [
    [f"{EX}/PanTrack_004_bl_ct.nii.gz", f"{EX}/PanTrack_004_fu_ct.nii.gz", f"{EX}/PanTrack_004_bl_mask.nii.gz",
     "233, 204, 60", "245, 198, 60"],   # pancreatic tumor
    [f"{EX}/PanTrack_004_bl_ct.nii.gz", f"{EX}/PanTrack_004_fu_ct.nii.gz", f"{EX}/PanTrack_004_bl_mask.nii.gz",
     "170, 203, 118", "173, 199, 129"],  # liver metastasis
    [f"{EX}/PanTrack_007_bl_ct.nii.gz", f"{EX}/PanTrack_007_fu_ct.nii.gz", f"{EX}/PanTrack_007_bl_mask.nii.gz",
     "309, 159, 60", "305, 193, 60"],   # pancreatic tumor
    [f"{EX}/PanTrack_007_bl_ct.nii.gz", f"{EX}/PanTrack_007_fu_ct.nii.gz", f"{EX}/PanTrack_007_bl_mask.nii.gz",
     "154, 182, 97", "177, 212, 87"],   # liver metastasis
]

with gr.Blocks(title="LongiTrack") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# LongiTrack — longitudinal lesion tracking in CT\n"
            "Point-prompted segmentation of a lesion in a **follow-up** CT, using the lesion's **baseline** appearance "
            "as longitudinal context ([model](https://huggingface.co/ykirchhoff/LongiTrack) · "
            "[paper](https://huggingface.co/papers/2605.23118) · [LongiSeg](https://github.com/MIC-DKFZ/LongiSeg)).  \n"
            "Load two scans, **click the lesion** in each (scroll with the slice slider), then press *Track lesion*. "
            "⚠️ *Research use only — not a medical device.*"
        )
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 1 · Baseline scan")
                bl_ct = gr.File(label="Baseline CT (.nii / .nii.gz)", file_types=[".nii", ".gz", ".mha", ".nrrd"])
                bl_mask = gr.File(label="Baseline lesion mask (optional)", file_types=[".nii", ".gz", ".mha", ".nrrd"])
                bl_view = gr.Image(label="Baseline — click the lesion", interactive=False, type="pil", height=420)
                bl_z = gr.Slider(0, 1, value=0, step=1, label="Baseline slice (z)")
                bl_point = gr.Textbox(label="Baseline point (x, y, z)", placeholder="click on the image")
            with gr.Column():
                gr.Markdown("### 2 · Follow-up scan")
                fu_ct = gr.File(label="Follow-up CT (.nii / .nii.gz)", file_types=[".nii", ".gz", ".mha", ".nrrd"])
                fu_view = gr.Image(label="Follow-up — click the same lesion", interactive=False, type="pil", height=420)
                fu_z = gr.Slider(0, 1, value=0, step=1, label="Follow-up slice (z)")
                fu_point = gr.Textbox(label="Follow-up point (x, y, z)", placeholder="click on the image")
        with gr.Accordion("Advanced settings", open=False):
            version = gr.Dropdown(VERSIONS, value="1.1", label="Model version",
                                  info="1.1: AutoPET IV + PanTrack + HCUCH · 1.0: AutoPET IV only")
        run = gr.Button("Track lesion", variant="primary")
        gr.Markdown("### 3 · Follow-up segmentation")
        with gr.Row():
            with gr.Column():
                out_img = gr.Image(label="Predicted lesion (cyan) — central slice", type="pil", height=420)
                info = gr.Markdown()
            with gr.Column():
                out_mesh = gr.Model3D(label="Lesion 3D surface", height=420, clear_color=[0.08, 0.08, 0.1, 1.0])
                out_mask = gr.File(label="Predicted mask (NIfTI, follow-up geometry)")
        with gr.Accordion("Browse the segmentation slice by slice", open=False):
            res_view = gr.Image(label="Follow-up + prediction", interactive=False, type="pil", height=420)
            res_z = gr.Slider(0, 1, value=0, step=1, label="Slice (z)")

        gr.Examples(
            examples=EXAMPLES,
            inputs=[bl_ct, fu_ct, bl_mask, bl_point, fu_point],
            outputs=[out_img, out_mask, out_mesh, info],
            fn=track_lesion,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples (PanTrack, CC-BY-NC-4.0 — pancreatic tumors & liver metastases)",
        )
        gr.Markdown(
            "<small>Points are voxel indices `(x, y, z)` of the uploaded volumes (x = column, y = row of the shown "
            "slice). Example scans: [PanTrack](https://huggingface.co/datasets/mrokuss/PanTrack) (CC-BY-NC-4.0, "
            "Kirchhoff, Rokuss et al.), z-cropped and resampled to 1.5 mm slices to keep them small. "
            "Model license: CC-BY-NC-SA-4.0.</small>"
        )

    # viewer wiring
    bl_ct.change(on_ct_upload, [bl_ct, bl_point], [bl_view, bl_z], api_visibility="private")
    fu_ct.change(on_ct_upload, [fu_ct, fu_point], [fu_view, fu_z], api_visibility="private")
    bl_z.change(render_slice, [bl_ct, bl_z, bl_point], bl_view, api_visibility="private", trigger_mode="always_last", show_progress="hidden")
    fu_z.change(render_slice, [fu_ct, fu_z, fu_point], fu_view, api_visibility="private", trigger_mode="always_last", show_progress="hidden")
    bl_view.select(on_click, [bl_ct, bl_z], [bl_point, bl_view], api_visibility="private")
    fu_view.select(on_click, [fu_ct, fu_z], [fu_point, fu_view], api_visibility="private")
    bl_point.change(on_point_text, [bl_ct, bl_point], [bl_view, bl_z], api_visibility="private")
    fu_point.change(on_point_text, [fu_ct, fu_point], [fu_view, fu_z], api_visibility="private")

    run.click(track_lesion, [bl_ct, fu_ct, bl_mask, bl_point, fu_point, version],
              [out_img, out_mask, out_mesh, info], api_name="track_lesion")
    out_mask.change(on_result, [fu_ct, out_mask], [res_view, res_z], api_visibility="private")
    res_z.change(browse_result, [fu_ct, out_mask, res_z], res_view, api_visibility="private",
                 trigger_mode="always_last", show_progress="hidden")

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)