"""
app.py — GazeRefine interactive demo (Hugging Face Space).

Upload flow
-----------
* Standard images (.jpg / .png / etc.) — drag-and-drop or click on the main
  gr.Image widget.  The same widget also accepts clicks to place fixations, so
  there is now only ONE image panel instead of two.
* DICOM files (.dcm) — use the separate "Upload DICOM" file picker.  The file
  is decoded with pydicom and converted to an RGB PIL image before being handed
  to the same fixation / run pipeline.
* Fixation file (.csv / .xlsx / .xls) — use the separate "Upload fixation
  file" picker.  This can contain fixations for one or many images (e.g. an
  eye-tracker export with one row per fixation).  After upload, three
  dropdowns let the user pick which column is the image-name/ID column and
  which columns hold X / Y (and, optionally, duration).  Rows are matched to
  the currently loaded image by filename; X/Y values are auto-detected as
  either normalised [0,1] or raw pixel coordinates.
"""
from __future__ import annotations

import sys
import types
import tempfile
import csv
import os
from pathlib import Path

# ── 1. audioop shim (Python 3.13 removed audioop; pydub needs it) ─────────────
if sys.version_info >= (3, 13):
    for _mod in ("audioop", "pyaudioop"):
        if _mod not in sys.modules:
            sys.modules[_mod] = types.ModuleType(_mod)

# ── 2. Patch starlette Jinja2Templates.TemplateResponse ──────────────────────
import starlette.templating as _st

_orig_TR = _st.Jinja2Templates.TemplateResponse

def _compat_TR(self, *args, **kwargs):
    if args and isinstance(args[0], str) and len(args) >= 2 and isinstance(args[1], dict):
        name       = args[0]
        context    = args[1]
        status_code = args[2] if len(args) > 2 else kwargs.get("status_code", 200)
        headers    = kwargs.get("headers")
        media_type = kwargs.get("media_type")
        background = kwargs.get("background")
        template   = self.get_template(name)
        return _st._TemplateResponse(
            template, context,
            status_code=status_code,
            headers=headers,
            media_type=media_type,
            background=background,
        )
    return _orig_TR(self, *args, **kwargs)

_st.Jinja2Templates.TemplateResponse = _compat_TR  # type: ignore[method-assign]

import gradio as gr

# ── 3. gradio_client schema shim ──────────────────────────────────────────────
try:
    import gradio_client.utils as _gcu
    _orig_inner = _gcu._json_schema_to_python_type

    def _safe_inner(schema, defs=None):
        if not isinstance(schema, dict):
            return "Any"
        if not isinstance(schema.get("additionalProperties"), dict):
            schema = {k: v for k, v in schema.items() if k != "additionalProperties"}
        return _orig_inner(schema, defs)

    _gcu._json_schema_to_python_type = _safe_inner
except Exception:
    pass

# ── 4. huggingface_hub HfFolder shim ─────────────────────────────────────────
try:
    from huggingface_hub import HfFolder  # noqa: F401
except ImportError:
    import huggingface_hub as _hfh
    class _FakeHfFolder:
        @staticmethod
        def get_token(): return None
    _hfh.HfFolder = _FakeHfFolder  # type: ignore[attr-defined]
    sys.modules["huggingface_hub"].HfFolder = _FakeHfFolder  # type: ignore[assignment]

import numpy as np
from PIL import Image, ImageDraw

# ── 5. Path setup ─────────────────────────────────────────────────────────────
_here = Path(__file__).resolve().parent
for _candidate in [_here] + list(_here.parents):
    _s = str(_candidate)
    if _s not in sys.path:
        sys.path.insert(0, _s)

import scripts.predict_single as _predict_module  # noqa: E402
from scripts.predict_single import predict         # noqa: E402

# ── Monkey-patch load_fixation_csv ────────────────────────────────────────────
# Our single-image temp CSV has x,y,duration in raw pixel coordinates with no
# image_name column.  The original loader expects a dataset CSV and returns an
# empty tensor when that column is absent.
# This patch detects the missing column, reads the CSV directly, normalises
# pixel → [0,1], and adds the batch dimension the model requires: (N,3)→(1,N,3).

import pandas as _pd
import torch as _torch

try:
    from gazerefine.gaze import load_fixation_csv as _orig_load_fixation_csv
except Exception:
    _orig_load_fixation_csv = None

def _patched_load_fixation_csv(csv_path, image_width=1, image_height=1, image_name=None):
    df = _pd.read_csv(csv_path)
    print(f"[PATCH] load_fixation_csv — columns: {list(df.columns)}, rows: {len(df)}")

    if "image_name" in df.columns and _orig_load_fixation_csv is not None:
        print("[PATCH] image_name column present — using original loader")
        return _orig_load_fixation_csv(csv_path, image_width=image_width,
                                       image_height=image_height, image_name=image_name)

    x   = df["x"].values.astype(float)
    y   = df["y"].values.astype(float)
    dur = df["duration"].values.astype(float)

    x_n   = x   / max(float(image_width),  1.0)
    y_n   = y   / max(float(image_height), 1.0)
    dur_n = dur / (dur.max() + 1e-8)

    # model expects (B, N, 3) — add batch dim
    fixations = _torch.tensor(
        list(zip(x_n, y_n, dur_n)), dtype=_torch.float32
    ).unsqueeze(0)  # (N, 3) → (1, N, 3)

    print(f"[PATCH] tensor shape: {tuple(fixations.shape)}")
    print(f"[PATCH] fixations (x_norm, y_norm, dur_norm):\n{fixations[0]}")
    return fixations

_predict_module.load_fixation_csv = _patched_load_fixation_csv


# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

PRESETS = {
    "Colonoscopy / polyp (Kvasir-SEG settings)": "colonoscopy",
    "Grayscale MRI / CT (prostate-MRI settings)": "mri",
}
POINT_COLORS = ["#ff3b30", "#ff9500", "#ffcc00", "#34c759", "#5ac8fa", "#007aff", "#af52de"]

_NO_COL = "— none —"


def dcm_to_pil(dcm_path: str) -> Image.Image:
    """Load a DICOM file and return an RGB PIL image."""
    import pydicom
    dcm = pydicom.dcmread(dcm_path)
    arr = dcm.pixel_array.astype(np.float32)
    arr = arr - arr.min()
    arr = arr / (arr.max() + 1e-8)
    arr = (arr * 255).astype(np.uint8)
    # Handle multi-frame / greyscale / RGB DICOM
    if arr.ndim == 2:
        return Image.fromarray(arr, mode="L").convert("RGB")
    if arr.ndim == 3 and arr.shape[0] in (1, 3, 4):
        # (C, H, W) → (H, W, C)
        arr = arr.transpose(1, 2, 0)
    return Image.fromarray(arr).convert("RGB")


def draw_points(image: Image.Image, points: list) -> Image.Image:
    """Overlay fixation circles on a copy of `image`.

    `points`: list of (x_px, y_px, duration) in original-image pixel coords.
    """
    if image is None:
        return None
    vis = image.convert("RGB").copy()
    draw = ImageDraw.Draw(vis)
    w, h = vis.size
    r = max(6, min(w, h) // 80)
    for i, (x_px, y_px, dur) in enumerate(points):
        color = POINT_COLORS[i % len(POINT_COLORS)]
        rad = r * (0.6 + 0.8 * dur)
        draw.ellipse(
            [x_px - rad, y_px - rad, x_px + rad, y_px + rad],
            outline=color, width=3,
        )
        draw.text((x_px + rad + 2, y_px - rad), str(i + 1), fill=color)
    return vis


def read_table(path: str) -> "_pd.DataFrame":
    """Load a .csv / .xlsx / .xls fixation file into a DataFrame."""
    ext = Path(path).suffix.lower()
    if ext in (".xlsx", ".xls"):
        return _pd.read_excel(path)
    # Sniff delimiter for csv/tsv/txt — eye-tracker exports are sometimes
    # tab-separated even with a .csv extension.
    return _pd.read_csv(path, sep=None, engine="python")


def normalize_xy(x_vals: np.ndarray, y_vals: np.ndarray, img_w: int, img_h: int):
    """Convert X/Y column values to pixel coords for the given image size.

    Values already in [0, 1] (inclusive, with a little slack for rounding)
    are treated as normalised; anything else is assumed to already be raw
    pixel coordinates and is left as-is (but clamped to the image bounds).
    """
    looks_normalized = (
        np.nanmax(x_vals) <= 1.05 and np.nanmax(y_vals) <= 1.05
        and np.nanmin(x_vals) >= -0.05 and np.nanmin(y_vals) >= -0.05
    )
    if looks_normalized:
        x_px = np.clip(x_vals, 0, 1) * img_w
        y_px = np.clip(y_vals, 0, 1) * img_h
    else:
        x_px = np.clip(x_vals, 0, img_w)
        y_px = np.clip(y_vals, 0, img_h)
    return x_px, y_px


# ─────────────────────────────────────────────────────────────────────────────
# Event handlers
# ─────────────────────────────────────────────────────────────────────────────


_UPLOAD_LABEL    = "Drop / click to load  .jpg  .png  .bmp  .tif  .tiff  .webp  .dcm"
_FIXATION_LABEL  = "Click to place fixations"
_FIXFILE_LABEL   = "Upload fixation file (.csv / .xlsx / .xls) — optional"


def _resolve_path(file_obj):
    """Extract a filesystem path from whatever gr.File passes."""
    if isinstance(file_obj, str):
        return file_obj
    if isinstance(file_obj, dict):
        return file_obj.get("name") or file_obj.get("path") or file_obj.get("tmp_path") or ""
    if hasattr(file_obj, "name"):
        return file_obj.name
    return ""


def on_file_upload(file_obj):
    """Load any image or DICOM and switch the panel to fixation-click mode."""
    _no_change = (None, [], gr.update(), gr.update(), gr.update(), gr.update())

    if file_obj is None:
        return _no_change

    image_name = ""

    # gr.Image gives PIL/numpy; gr.File gives a path
    if isinstance(file_obj, Image.Image):
        pil = file_obj.convert("RGB")
    elif isinstance(file_obj, np.ndarray):
        pil = Image.fromarray(file_obj).convert("RGB")
    else:
        path = _resolve_path(file_obj)
        if not path:
            gr.Warning("Could not resolve file path.")
            return _no_change
        image_name = Path(path).name
        ext = Path(path).suffix.lower()
        try:
            pil = dcm_to_pil(path) if ext == ".dcm" else Image.open(path).convert("RGB")
        except Exception as e:
            gr.Warning(f"Could not load file: {e}")
            return _no_change

    print(f"[DEBUG] on_file_upload — size={pil.size}  name={image_name!r}")
    # Switch: hide upload zone, show image panel + delete button
    return (
        pil,                                    # orig_image_state
        [],                                     # points_state
        image_name,                              # image_name_state
        gr.update(visible=False),               # upload_zone  → hide
        gr.update(value=pil, visible=True,
                  label=_FIXATION_LABEL),       # image_panel  → show with image
        gr.update(visible=True),                # delete_btn   → show
    )


def on_select(orig_image: Image.Image, points: list, duration: float, evt: gr.SelectData):
    """Record a fixation click in original-image pixel coords."""
    if orig_image is None:
        gr.Warning("Upload an image first.")
        return points, gr.update()
    x_px, y_px = float(evt.index[0]), float(evt.index[1])
    new_points = points + [(x_px, y_px, duration)]
    print(f"[DEBUG] fixation #{len(new_points)}: x={x_px:.1f}  y={y_px:.1f}  dur={duration}")
    return new_points, draw_points(orig_image, new_points)


def on_clear(orig_image):
    """Remove all fixations but keep the current image."""
    if orig_image is None:
        return [], gr.update()
    return [], gr.update(value=orig_image)


def on_delete():
    """Delete the current image and return to upload mode."""
    return (
        None,                                   # orig_image_state
        [],                                     # points_state
        "",                                      # image_name_state
        gr.update(value=None, visible=True),    # upload_zone  → show (reset)
        gr.update(value=None, visible=False),   # image_panel  → hide
        gr.update(visible=False),               # delete_btn   → hide
    )


# ── Fixation-file upload → column mapping ────────────────────────────────────

def on_fixfile_upload(file_obj):
    """Load the fixation table and populate the column-mapping dropdowns."""
    _hide = (
        None, gr.update(visible=False),
        gr.update(choices=[], value=None), gr.update(choices=[], value=None),
        gr.update(choices=[], value=None), gr.update(choices=[], value=None),
        gr.update(visible=False),
    )
    if file_obj is None:
        return _hide

    path = _resolve_path(file_obj)
    if not path:
        gr.Warning("Could not resolve fixation file path.")
        return _hide

    try:
        df = read_table(path)
    except Exception as e:
        gr.Warning(f"Could not read fixation file: {e}")
        return _hide

    if df.empty or len(df.columns) == 0:
        gr.Warning("Fixation file appears to be empty.")
        return _hide

    cols = [str(c) for c in df.columns]
    print(f"[DEBUG] fixation file loaded — columns: {cols}, rows: {len(df)}")

    def _guess(*keywords, fallback=None):
        for c in cols:
            cl = c.lower()
            if any(k in cl for k in keywords):
                return c
        return fallback if fallback is not None else cols[0]

    guess_id  = _guess("image", "id", "name", "file", fallback=cols[0])
    # exact / boundary-aware matches first (avoids "fix_index" matching "x"),
    # then fall back to a bare trailing "x" / "y".
    guess_x = _guess("fix_x", "pos_x", "gaze_x", fallback=None)
    if guess_x is None:
        guess_x = next((c for c in cols if c.lower().rstrip("_") .endswith("x")
                         and "index" not in c.lower()), cols[0])
    guess_y = _guess("fix_y", "pos_y", "gaze_y", fallback=None)
    if guess_y is None:
        guess_y = next((c for c in cols if c.lower().rstrip("_").endswith("y")
                         and "index" not in c.lower()), cols[0])

    dur_choices = [_NO_COL] + cols
    guess_dur = _guess("duration", "dur", fallback=_NO_COL)

    return (
        df.to_json(),                                          # fixfile_df_state (serialized)
        gr.update(visible=True),                                # mapping_row → show
        gr.update(choices=cols, value=guess_id),                # id_col_dd
        gr.update(choices=cols, value=guess_x),                 # x_col_dd
        gr.update(choices=cols, value=guess_y),                 # y_col_dd
        gr.update(choices=dur_choices, value=guess_dur),         # dur_col_dd
        gr.update(visible=True),                                # apply_fix_btn → show
    )


def on_apply_fixfile(fixfile_json, id_col, x_col, y_col, dur_col,
                      orig_image, image_name):
    """Match rows to the currently loaded image (by filename) and load
    them as fixation points, replacing whatever points are currently set.

    If no rows match the loaded image's filename, nothing is loaded — the
    existing points (if any) are left untouched, and the user is warned so
    they can check the ID column / image filename instead of silently
    getting fixations for the wrong image."""
    if orig_image is None:
        gr.Warning("Load an image first, then apply the fixation file.")
        return gr.update(), gr.update()
    if not fixfile_json:
        gr.Warning("Upload a fixation file first.")
        return gr.update(), gr.update()
    if not id_col or not x_col or not y_col:
        gr.Warning("Pick the ID, X and Y columns first.")
        return gr.update(), gr.update()
    if not image_name:
        gr.Warning(
            "Couldn't determine the loaded image's filename (this can "
            "happen if the image was pasted/dropped without a filename). "
            "Re-upload the image as a file and try again."
        )
        return gr.update(), gr.update()

    df = _pd.read_json(fixfile_json)

    # Match by exact filename first, then by stem-without-extension, so the
    # mapping still works if the fixation file's IMAGE column omits the
    # extension or uses a different one than the uploaded image.
    mask = df[id_col].astype(str) == image_name
    if not mask.any():
        stem_no_ext = Path(image_name).stem
        mask = df[id_col].astype(str).apply(lambda v: Path(str(v)).stem) == stem_no_ext

    sub = df[mask]

    if sub.empty:
        gr.Warning(
            f"No rows in the fixation file match the loaded image "
            f"('{image_name}'). Nothing was loaded — check that the ID "
            f"column values match the image filename."
        )
        return gr.update(), gr.update()

    w, h = orig_image.size
    x_vals = sub[x_col].astype(float).to_numpy()
    y_vals = sub[y_col].astype(float).to_numpy()
    x_px, y_px = normalize_xy(x_vals, y_vals, w, h)

    if dur_col and dur_col != _NO_COL and dur_col in sub.columns:
        dur_raw = sub[dur_col].astype(float).to_numpy()
        dmax = float(np.nanmax(dur_raw)) if len(dur_raw) else 1.0
        dur_n = dur_raw / (dmax + 1e-8)
    else:
        dur_n = np.full(len(sub), 1.0)

    new_points = [
        (float(xp), float(yp), float(d))
        for xp, yp, d in zip(x_px, y_px, dur_n)
    ]
    print(f"[DEBUG] loaded {len(new_points)} fixations from file for image '{image_name}'")

    return new_points, draw_points(orig_image, new_points)


def run(orig_image: Image.Image, points: list, preset_name: str, threshold: float):
    import traceback, uuid
    print(f"[DEBUG] run — points={len(points)}  preset={preset_name}")

    if orig_image is None:
        gr.Warning("Upload an image first.")
        return None, None, None
    if not points:
        gr.Warning("Click on the image at least once to place a fixation (or load a fixation file).")
        return None, None, None

    preset_key = PRESETS[preset_name]
    w, h = orig_image.size

    shared_stem       = f"gazerefine_{uuid.uuid4().hex}"
    tmp_img_path      = os.path.join(tempfile.gettempdir(), f"{shared_stem}.png")
    fixation_csv_path = os.path.join(tempfile.gettempdir(), f"{shared_stem}.csv")

    orig_image.convert("RGB").save(tmp_img_path)

    with open(fixation_csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["x", "y", "duration"])
        for x_px, y_px, dur in points:
            writer.writerow([x_px, y_px, dur])

    print(f"[DEBUG] image {w}x{h}  |  {len(points)} fixations  |  preset={preset_key}  thr={threshold}")

    with open(fixation_csv_path) as f:
        print(f"[DEBUG] CSV:\n{f.read()}")

    try:
        out = predict(
            image_path=tmp_img_path,
            fixation_csv=fixation_csv_path,
            preset=preset_key,
            threshold=threshold,
            return_all=True,
        )
    except Exception as e:
        print(f"[ERROR] predict() raised: {e}")
        traceback.print_exc()
        gr.Warning(f"Prediction failed: {e}")
        return None, None, None
    finally:
        for p in (tmp_img_path, fixation_csv_path):
            try:
                os.unlink(p)
            except OSError:
                pass

    mask_arr = np.array(out["mask"])
    print(f"[DEBUG] mask non-zero: {(mask_arr > 0).sum()} / {mask_arr.size}")
    return out["gaze_overlay"], out["mask_overlay"], out["mask"]


# ─────────────────────────────────────────────────────────────────────────────
# UI
# ─────────────────────────────────────────────────────────────────────────────

with gr.Blocks(title="GazeRefine — gaze-guided zero-shot segmentation") as demo:
    gr.Markdown(
        """
        # 👁️ GazeRefine — Expert Gaze as a Test-Time Prompt
        Training-free, zero-shot medical image segmentation.
        Upload an image or DICOM, click to place fixations (or load a fixation
        file), then hit **Run**.
        """
    )

    orig_image_state = gr.State(None)
    points_state      = gr.State([])
    image_name_state  = gr.State("")     # filename of the currently loaded image
    fixfile_df_state  = gr.State(None)   # serialized DataFrame (to_json) of the uploaded fixation file

    with gr.Row():
        # ── Left column ───────────────────────────────────────────────────────
        with gr.Column(scale=1):

            # ── Upload zone (visible when no image loaded) ────────────────────
            upload_zone = gr.File(
                label=_UPLOAD_LABEL,
                file_types=[".jpg", ".jpeg", ".png", ".bmp",
                            ".tif", ".tiff", ".webp", ".gif", ".dcm"],
                file_count="single",
                visible=True,
                elem_id="upload_zone",
            )

            # ── Image panel (hidden until image loaded; click to fixate) ──────
            image_panel = gr.Image(
                type="pil",
                label=_FIXATION_LABEL,
                height=430,
                interactive=False,          # no toolbar → .select fires on click
                show_download_button=False,
                visible=False,
                elem_id="image_panel",
            )

            # ── Delete button (hidden until image loaded) ─────────────────────
            delete_btn = gr.Button("🗑 Delete image — load another", visible=False, variant="secondary")

            # ── Fixation file upload (optional alternative to manual clicks) ──
            with gr.Accordion("📄 Load fixations from file", open=False):
                fixfile_upload = gr.File(
                    label=_FIXFILE_LABEL,
                    file_types=[".csv", ".xlsx", ".xls", ".tsv", ".txt"],
                    file_count="single",
                    elem_id="fixfile_upload",
                )
                with gr.Row(visible=False) as mapping_row:
                    id_col_dd  = gr.Dropdown(label="Image / ID column", choices=[])
                    x_col_dd   = gr.Dropdown(label="X column", choices=[])
                    y_col_dd   = gr.Dropdown(label="Y column", choices=[])
                    dur_col_dd = gr.Dropdown(label="Duration column (optional)", choices=[])
                apply_fix_btn = gr.Button(
                    "📥 Load fixations for current image", visible=False,
                )

            # ── Controls ──────────────────────────────────────────────────────
            with gr.Row():
                duration_slider = gr.Slider(
                    0.1, 1.0, value=1.0, step=0.1,
                    label="Fixation duration weight",
                )
                clear_btn = gr.Button("✖ Clear fixations")

            preset = gr.Radio(
                list(PRESETS.keys()), value=list(PRESETS.keys())[0],
                label="Preset",
            )
            threshold = gr.Slider(
                0.1, 0.9, value=0.5, step=0.05,
                label="Mask threshold",
            )
            run_btn = gr.Button("▶ Run GazeRefine", variant="primary")

        # ── Right column: outputs ─────────────────────────────────────────────
        with gr.Column(scale=1):
            gaze_out = gr.Image(label="Gaze prior", height=260)
            with gr.Row():
                mask_overlay_out = gr.Image(label="Mask overlay", height=260)
                mask_only_out    = gr.Image(label="Binary mask",  height=260)

    # ── Event wiring ──────────────────────────────────────────────────────────

    _upload_outputs = [orig_image_state, points_state, image_name_state,
                        upload_zone, image_panel, delete_btn]

    upload_zone.upload(on_file_upload, inputs=[upload_zone], outputs=_upload_outputs)
    upload_zone.change(on_file_upload, inputs=[upload_zone], outputs=_upload_outputs)

    image_panel.select(
        on_select,
        inputs=[orig_image_state, points_state, duration_slider],
        outputs=[points_state, image_panel],
    )

    clear_btn.click(
        on_clear,
        inputs=[orig_image_state],
        outputs=[points_state, image_panel],
    )

    delete_btn.click(
        on_delete,
        outputs=[orig_image_state, points_state, image_name_state,
                 upload_zone, image_panel, delete_btn],
    )

    _fixfile_outputs = [fixfile_df_state, mapping_row, id_col_dd, x_col_dd, y_col_dd, dur_col_dd, apply_fix_btn]

    fixfile_upload.upload(on_fixfile_upload, inputs=[fixfile_upload], outputs=_fixfile_outputs)
    fixfile_upload.change(on_fixfile_upload, inputs=[fixfile_upload], outputs=_fixfile_outputs)

    apply_fix_btn.click(
        on_apply_fixfile,
        inputs=[fixfile_df_state, id_col_dd, x_col_dd, y_col_dd, dur_col_dd,
                orig_image_state, image_name_state],
        outputs=[points_state, image_panel],
    )

    run_btn.click(
        run,
        inputs=[orig_image_state, points_state, preset, threshold],
        outputs=[gaze_out, mask_overlay_out, mask_only_out],
    )

    gr.Markdown(
        "Method: GazeRefine — frozen DINOv3 + gaze-anchored prototypes + recurrent "
        "foreground/background refinement, entirely training-free."
    )

demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)