#!/usr/bin/env python3
"""
Map Georectification App

Workflow:
  1. Select projection type (with guidance on which points to pick)
  2. Click known lat/lon positions on the map → add Ground Control Points (GCPs)
  3. Compute the affine pixel→geographic transform
  4. Switch to Query tab → click anywhere to get lat/lon, label and save points

Run with:
  uv run app.py
"""

import csv
import io
import tempfile

import gradio as gr
from PIL import Image

from georect import (
    IMAGE_PATH,
    PROJECTIONS,
    gcps_df,
    get_proj4,
    points_df,
    px_to_lonlat,
    render_calib,
    render_query,
    solve_affine,
)


ZOOM_JS = """
const MIN_SCALE = 0.5, MAX_SCALE = 20, STEP = 0.15;

function initZoom(container) {
  if (container.dataset.zoomInit) return;
  container.dataset.zoomInit = '1';

  let scale = 1, panX = 0, panY = 0;
  let isPanning = false, startX = 0, startY = 0;

  function getImg() { return container.querySelector('img'); }

  function apply() {
    const img = getImg();
    if (!img) return;
    img.style.transformOrigin = '0 0';
    img.style.transform = `translate(${panX}px, ${panY}px) scale(${scale})`;
  }

  container.addEventListener('wheel', (e) => {
    if (!e.shiftKey) return;
    e.preventDefault();
    e.stopPropagation();
    const img = getImg();
    if (!img) return;

    const rect = img.getBoundingClientRect();
    const cx = (e.clientX - rect.left) / scale;
    const cy = (e.clientY - rect.top) / scale;

    const prev = scale;
    const delta = e.deltaY > 0 ? -STEP : STEP;
    scale = Math.min(MAX_SCALE, Math.max(MIN_SCALE, scale * (1 + delta)));

    panX += cx * (prev - scale);
    panY += cy * (prev - scale);
    apply();
  }, {passive: false});

  // middle-click drag to pan (avoids conflict with left-click point selection)
  container.addEventListener('mousedown', (e) => {
    if (e.button === 1) {
      isPanning = true;
      startX = e.clientX - panX;
      startY = e.clientY - panY;
      e.preventDefault();
    }
  });
  window.addEventListener('mousemove', (e) => {
    if (!isPanning) return;
    panX = e.clientX - startX;
    panY = e.clientY - startY;
    apply();
  });
  window.addEventListener('mouseup', () => { isPanning = false; });

  // double-click to reset zoom
  container.addEventListener('dblclick', () => {
    scale = 1; panX = 0; panY = 0;
    apply();
  });

  container.style.overflow = 'hidden';
}

function scan() {
  document.querySelectorAll('.image-frame').forEach(initZoom);
}
scan();
new MutationObserver(scan).observe(document.body, {childList: true, subtree: true});
"""


def build_app():
    base = Image.open(IMAGE_PATH).convert("RGB")

    with gr.Blocks(title="Map Georectification") as demo:

        # ── Shared state ──────────────────────────────────────────────────────
        base_st    = gr.State(base)
        gcps_st    = gr.State([])
        tfm_st     = gr.State({"cx": None, "cy": None, "proj4": None})
        picked_st  = gr.State([])
        last_px_st = gr.State(None)
        last_py_st = gr.State(None)

        gr.Markdown(
            "## 🗺 Map Georectification\n"
            "**① Calibrate** – choose projection and add Ground Control Points (GCPs).  "
            "**② Compute** – fit pixel→coordinate transform.  "
            "**③ Query** – click anywhere to get lat/lon coordinates."
        )

        with gr.Tabs():

            # ── Tab 1 · Calibrate ─────────────────────────────────────────────
            with gr.Tab("① Calibrate"):
                with gr.Row():
                    with gr.Column(scale=3):
                        calib_img = gr.Image(
                            value=base,
                            label="Click to place a control point",
                            type="pil",
                            interactive=True,
                            buttons=[],
                        )
                    with gr.Column(scale=2):
                        proj_sel = gr.Dropdown(
                            list(PROJECTIONS.keys()),
                            value="Albers Equal Area (CONUS)",
                            label="Map Projection",
                        )
                        proj_hint = gr.Markdown(PROJECTIONS["Albers Equal Area (CONUS)"]["hint"])
                        custom_box = gr.Textbox(
                            label="Custom PROJ4 / EPSG string",
                            placeholder="e.g.  +proj=aea +lat_1=29.5 ...  or  EPSG:5070",
                            visible=False,
                        )

                        gr.Markdown("---\n**Add Control Point**  \n"
                                    "*Click a known location on the map, then enter its coordinates:*")
                        with gr.Row():
                            px_box = gr.Number(label="Pixel X", precision=0, interactive=False)
                            py_box = gr.Number(label="Pixel Y", precision=0, interactive=False)
                        with gr.Row():
                            lon_box = gr.Number(label="Longitude °E  (negative = West)")
                            lat_box = gr.Number(label="Latitude °N")
                        with gr.Row():
                            add_btn = gr.Button("Add GCP", variant="primary")
                            del_btn = gr.Button("Remove Last")
                            clr_btn = gr.Button("Clear All", variant="stop")

                        gcp_tbl = gr.Dataframe(
                            value=gcps_df([]), interactive=False, wrap=True,
                            label="Control Points",
                        )
                        compute_btn = gr.Button(
                            "⚙ Compute Georectification", variant="primary", size="lg"
                        )
                        status_md = gr.Markdown("*Add ≥ 3 control points, then click Compute.*")

            # ── Tab 2 · Query ─────────────────────────────────────────────────
            with gr.Tab("② Query"):
                warn_md = gr.Markdown(
                    "> ⚠️ No transform yet – go to the **Calibrate** tab first.",
                    visible=True,
                )
                with gr.Row():
                    with gr.Column(scale=3):
                        query_img = gr.Image(
                            value=base,
                            label="Click to query geographic coordinates",
                            type="pil",
                            interactive=True,
                            buttons=[],
                        )
                    with gr.Column(scale=2):
                        gr.Markdown("**Last click**")
                        with gr.Row():
                            q_lat = gr.Number(label="Latitude °N",  precision=6, interactive=False)
                            q_lon = gr.Number(label="Longitude °E", precision=6, interactive=False)
                        q_info = gr.Markdown("*Click on the map.*")

                        q_label = gr.Textbox(
                            label="Label for this point",
                            placeholder="e.g. Sample location A",
                        )
                        save_btn = gr.Button("Save Point", variant="primary")

                        gr.Markdown("---\n**Saved Points**")
                        pts_tbl = gr.Dataframe(
                            value=points_df([]), interactive=False, wrap=True,
                        )
                        with gr.Row():
                            exp_btn     = gr.Button("Export CSV")
                            clr_pts_btn = gr.Button("Clear Points")
                        csv_out = gr.File(label="Download CSV", visible=False)

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

        def on_proj(name):
            return PROJECTIONS[name]["hint"], gr.update(visible=(name == "Custom PROJ4 String"))

        proj_sel.change(on_proj, proj_sel, [proj_hint, custom_box])

        def on_calib_select(evt: gr.SelectData):
            return int(evt.index[0]), int(evt.index[1])

        calib_img.select(on_calib_select, outputs=[px_box, py_box])

        def on_upload(img):
            if img is None:
                img = base
            new_base = img.convert("RGB")
            tfm = {"cx": None, "cy": None, "proj4": None}
            return (
                new_base,                       # base_st
                [], gcps_df([]),                # gcps_st, gcp_tbl
                tfm,                            # tfm_st
                [], points_df([]),              # picked_st, pts_tbl
                new_base, new_base,             # calib_img, query_img
                gr.update(visible=True),        # warn_md
                "*Add ≥ 3 control points, then click Compute.*",
                None, None,                     # px_box, py_box
                None, None,                     # last_px_st, last_py_st
            )

        upload_outputs = [
            base_st, gcps_st, gcp_tbl, tfm_st, picked_st, pts_tbl,
            calib_img, query_img, warn_md, status_md,
            px_box, py_box, last_px_st, last_py_st,
        ]
        calib_img.upload(on_upload, calib_img, upload_outputs)
        query_img.upload(on_upload, query_img, upload_outputs)

        def on_add(base_img, gcps, px, py, lon, lat):
            if None in (px, py, lon, lat):
                return gcps, render_calib(base_img, gcps), gcps_df(gcps)
            new = gcps + [{"px": float(px), "py": float(py),
                           "lon": float(lon), "lat": float(lat)}]
            return new, render_calib(base_img, new), gcps_df(new)

        add_btn.click(on_add, [base_st, gcps_st, px_box, py_box, lon_box, lat_box],
                      [gcps_st, calib_img, gcp_tbl])

        def on_del(base_img, gcps):
            new = gcps[:-1]
            return new, render_calib(base_img, new), gcps_df(new)

        del_btn.click(on_del, [base_st, gcps_st], [gcps_st, calib_img, gcp_tbl])

        def on_clr_gcps(base_img):
            return [], base_img, gcps_df([])

        clr_btn.click(on_clr_gcps, base_st, [gcps_st, calib_img, gcp_tbl])

        def on_compute(base_img, gcps, proj_name, custom):
            proj4 = get_proj4(proj_name, custom)
            cx, cy, msg = solve_affine(gcps, proj4)
            tfm = {"cx": cx, "cy": cy, "proj4": proj4}
            return tfm, msg, gr.update(visible=(cx is None)), render_query(base_img, gcps, [])

        compute_btn.click(
            on_compute, [base_st, gcps_st, proj_sel, custom_box],
            [tfm_st, status_md, warn_md, query_img],
        )

        def on_query_select(tfm, evt: gr.SelectData):
            px, py = int(evt.index[0]), int(evt.index[1])
            cx, cy, proj4 = tfm["cx"], tfm["cy"], tfm["proj4"]
            if cx is None:
                return px, py, None, None, "*No transform – calibrate first.*"
            lon, lat = px_to_lonlat(px, py, cx, cy, proj4)
            msg = f"Pixel **({px}, {py})** → **{lat:.5f}° N,  {lon:.5f}° E**"
            return px, py, float(lat), float(lon), msg

        query_img.select(
            on_query_select, [tfm_st],
            [last_px_st, last_py_st, q_lat, q_lon, q_info],
        )

        def on_save(base_img, last_px, last_py, tfm, gcps, picked, label):
            if last_px is None or tfm["cx"] is None:
                return picked, points_df(picked), render_query(base_img, gcps, picked)
            lon, lat = px_to_lonlat(last_px, last_py, tfm["cx"], tfm["cy"], tfm["proj4"])
            pt = {
                "px": last_px, "py": last_py, "lon": lon, "lat": lat,
                "label": label.strip() or f"Point {len(picked) + 1}",
            }
            new = picked + [pt]
            return new, points_df(new), render_query(base_img, gcps, new)

        save_btn.click(
            on_save, [base_st, last_px_st, last_py_st, tfm_st, gcps_st, picked_st, q_label],
            [picked_st, pts_tbl, query_img],
        )

        def on_clr_pts(base_img, gcps):
            return [], points_df([]), render_query(base_img, gcps, [])

        clr_pts_btn.click(on_clr_pts, [base_st, gcps_st], [picked_st, pts_tbl, query_img])

        def on_export(picked):
            if not picked:
                return gr.update(visible=False)
            buf = io.StringIO()
            w = csv.writer(buf)
            w.writerow(["Label", "Latitude_N", "Longitude_E", "Pixel_X", "Pixel_Y"])
            for p in picked:
                w.writerow([p.get("label", ""), p["lat"], p["lon"], p["px"], p["py"]])
            with tempfile.NamedTemporaryFile(
                mode="w", suffix=".csv", prefix="georect_points_", delete=False
            ) as f:
                f.write(buf.getvalue())
            return gr.update(value=f.name, visible=True)

        exp_btn.click(on_export, picked_st, csv_out)

    return demo


if __name__ == "__main__":
    build_app().launch(theme=gr.themes.Soft(), js=ZOOM_JS)