georect.py

"""
Core logic for map georectification:
  - Projection definitions and helpers
  - Affine pixel→projected transform (least-squares fit)
  - Image rendering (crosshair markers)
  - DataFrame helpers for GCPs and picked points
"""

from pathlib import Path

import numpy as np
import pandas as pd
import pyproj
from PIL import Image, ImageDraw

IMAGE_PATH = Path(__file__).parent / "img.jpg"

# ── Projection definitions ────────────────────────────────────────────────────

PROJECTIONS = {
    "Albers Equal Area (CONUS)": {
        "proj4": "+proj=aea +lat_1=29.5 +lat_2=45.5 +lat_0=37.5 +lon_0=-96 +x_0=0 +y_0=0 +datum=NAD83 +units=m",
        "hint": (
            "**Albers Equal Area** – standard for US geological/geographic maps.  \n"
            "Click on **lat/lon grid intersections** visible on the map "
            "(e.g. where the 40°N line meets the 100°W line).  \n"
            "Need **≥ 3 points**, ideally spread to the corners."
        ),
    },
    "Lambert Conformal Conic (CONUS)": {
        "proj4": "+proj=lcc +lat_1=33 +lat_2=45 +lat_0=39 +lon_0=-96 +x_0=0 +y_0=0 +datum=NAD83 +units=m",
        "hint": (
            "**Lambert Conformal Conic** – common for North American maps.  \n"
            "Pick **≥ 4** lat/lon grid intersections spread across the map."
        ),
    },
    "Mercator": {
        "proj4": "+proj=merc +datum=WGS84 +units=m",
        "hint": (
            "**Mercator** – horizontal grid lines are evenly spaced; "
            "vertical spacing increases toward the poles.  \n"
            "Pick **≥ 3** widely-spaced lat/lon intersections."
        ),
    },
    "Equirectangular (Plate Carrée)": {
        "proj4": None,
        "hint": (
            "**Equirectangular** – lon/lat vary linearly with pixel x/y.  \n"
            "Two diagonally-opposite corner points are enough."
        ),
    },
    "Custom PROJ4 String": {
        "proj4": "custom",
        "hint": "Enter a valid PROJ4 (or `EPSG:XXXX`) string below, then add control points.",
    },
}


# ── Geo math ──────────────────────────────────────────────────────────────────

def get_proj4(proj_name: str, custom: str) -> str | None:
    if proj_name == "Custom PROJ4 String":
        return custom.strip() or None
    return PROJECTIONS[proj_name]["proj4"]


def _lonlat_to_proj(lon: float, lat: float, proj4: str | None):
    if proj4 is None:
        return lon, lat
    return pyproj.Proj(proj4)(lon, lat)


def _proj_to_lonlat(x: float, y: float, proj4: str | None):
    if proj4 is None:
        return x, y  # lon, lat
    return pyproj.Proj(proj4)(x, y, inverse=True)  # lon, lat


def solve_affine(gcps: list, proj4: str | None) -> tuple:
    """
    Fit an affine pixel→projected transform from a list of GCP dicts
    (each with keys px, py, lon, lat).

    Returns (cx, cy, message) where cx/cy are the 3-element coefficient
    vectors [a, b, c] such that x_proj = a*px + b*py + c.
    """
    n = len(gcps)
    if n < 3:
        return None, None, f"Need ≥ 3 control points (have {n})."

    A, Bx, By = [], [], []
    for g in gcps:
        px, py = float(g["px"]), float(g["py"])
        try:
            xp, yp = _lonlat_to_proj(float(g["lon"]), float(g["lat"]), proj4)
        except Exception as e:
            return None, None, f"Projection error at ({g['lon']}, {g['lat']}): {e}"
        A.append([px, py, 1.0])
        Bx.append(xp)
        By.append(yp)

    A = np.array(A)
    Bx = np.array(Bx)
    By = np.array(By)
    cx, *_ = np.linalg.lstsq(A, Bx, rcond=None)
    cy, *_ = np.linalg.lstsq(A, By, rcond=None)

    residuals = np.sqrt((A @ cx - Bx) ** 2 + (A @ cy - By) ** 2)
    rmse = float(np.sqrt(np.mean(residuals ** 2)))
    unit = "m" if proj4 else "°"
    approx = f" (~{rmse / 111_000 * 60:.2f}′ arc)" if proj4 else ""

    if n == 3:
        summary = (
            f"✅ Transform fitted from {n} points. "
            f"RMSE: {rmse:.2f} {unit} *(always 0 with exactly 3 pts — add more for a meaningful residual)*"
        )
    else:
        summary = f"✅ Transform fitted from {n} points. RMSE: {rmse:.2f} {unit}{approx}"

    per_point = "\n".join(
        f"  - GCP {i+1} ({g['lon']}°E, {g['lat']}°N): {r:.2f} {unit}"
        for i, (g, r) in enumerate(zip(gcps, residuals))
    )
    msg = f"{summary}\n\n**Per-point residuals:**\n{per_point}"
    return cx.tolist(), cy.tolist(), msg


def px_to_lonlat(px: float, py: float, cx, cy, proj4: str | None) -> tuple:
    """Convert pixel coordinates to (lon, lat) using the fitted transform."""
    if cx is None:
        return None, None
    v = np.array([float(px), float(py), 1.0])
    return _proj_to_lonlat(float(np.dot(cx, v)), float(np.dot(cy, v)), proj4)


# ── Image rendering ───────────────────────────────────────────────────────────

def _crosshair(draw: ImageDraw.ImageDraw, px, py, color, r=10, w=2, label=None):
    draw.ellipse([px - r, py - r, px + r, py + r], outline=color, width=w)
    draw.line([px - r - 4, py, px + r + 4, py], fill=color, width=w)
    draw.line([px, py - r - 4, px, py + r + 4], fill=color, width=w)
    if label:
        draw.text((px + r + 4, py - r), label, fill=color)


def render_calib(base: Image.Image, gcps: list) -> Image.Image:
    """Return a copy of base with GCP crosshairs drawn in red."""
    img = base.copy()
    draw = ImageDraw.Draw(img)
    for i, g in enumerate(gcps):
        _crosshair(draw, int(g["px"]), int(g["py"]), "#FF4444", label=str(i + 1))
    return img


def render_query(base: Image.Image, gcps: list, picked: list) -> Image.Image:
    """Return a copy of base with grey GCP markers and cyan picked-point markers."""
    img = base.copy()
    draw = ImageDraw.Draw(img)
    for g in gcps:
        _crosshair(draw, int(g["px"]), int(g["py"]), "#888888", r=6, w=1)
    for pt in picked:
        _crosshair(draw, int(pt["px"]), int(pt["py"]), "#44CCFF", label=pt.get("label", ""))
    return img


# ── DataFrame helpers ─────────────────────────────────────────────────────────

def gcps_df(gcps: list) -> pd.DataFrame:
    cols = ["#", "Pixel X", "Pixel Y", "Lon °E", "Lat °N"]
    if not gcps:
        return pd.DataFrame(columns=cols)
    return pd.DataFrame(
        [[i + 1, g["px"], g["py"], g["lon"], g["lat"]] for i, g in enumerate(gcps)],
        columns=cols,
    )


def points_df(pts: list) -> pd.DataFrame:
    cols = ["Label", "Lat °N", "Lon °E", "Pixel X", "Pixel Y"]
    if not pts:
        return pd.DataFrame(columns=cols)
    return pd.DataFrame(
        [[p.get("label", ""), round(p["lat"], 6), round(p["lon"], 6), p["px"], p["py"]]
         for p in pts],
        columns=cols,
    )