"""Stage `sketch`: sample the sketch + user masks onto cells."""
from __future__ import annotations

from pathlib import Path

import numpy as np
from PIL import Image

from . import zones as ZN
from .config import params
from .noise import fbm
from .pipeline import StageError
from .sphere import latlon_to_xyz, xyz_to_latlon

W, H = 2000, 1000
SKETCH = ("land", "mountains", "desert", "rainforest", "trench")
MASKS = ("land_hint", "mountain_hint", "o2_zones", "gravity_zones", "lock")
DEFAULTS = {"warp_km": 1000.0, "warp_freq": 1.2, "detail_warp_km": 200.0, "detail_warp_freq": 5.0, "moves": []}


def load_png01(path: Path) -> np.ndarray:
    return np.array(Image.open(path).convert("L"), dtype=np.float64) / 255.0


def load_mask(path: Path) -> np.ndarray:
    """Grayscale override mask → [−1, 1]; mid-grey (128 / 32768), transparency and missing detail = neutral 0."""
    try:
        im = Image.open(path)
        im.load()
    except (OSError, ValueError) as e:
        raise StageError(f"cannot read mask {path.name}: {e}") from e
    if im.mode in ("I", "I;16", "I;16B", "I;16L"):
        v = (np.array(im).astype(np.float64) - 32768.0) / 32768.0
        dead = 1.0 / 32768.0
    else:
        la = np.array(im.convert("LA")).astype(np.float64)
        v = np.where(la[..., 1] > 0, (la[..., 0] - 127.5) / 127.5, 0.0)
        dead = 1.0 / 255.0
    return np.clip(np.where(np.abs(v) <= dead, 0.0, v), -1.0, 1.0)


def sample_equirect(img: np.ndarray, lat, lon) -> np.ndarray:
    h, w = img.shape
    x = (np.asarray(lon, dtype=np.float64) + 180.0) / 360.0 * w - 0.5
    y = (90.0 - np.asarray(lat, dtype=np.float64)) / 180.0 * h - 0.5
    x0 = np.floor(x).astype(np.int64)
    y0 = np.floor(y).astype(np.int64)
    fx, fy = x - x0, y - y0
    xa, xb = x0 % w, (x0 + 1) % w
    ya, yb = np.clip(y0, 0, h - 1), np.clip(y0 + 1, 0, h - 1)
    top = img[ya, xa] * (1 - fx) + img[ya, xb] * fx
    bot = img[yb, xa] * (1 - fx) + img[yb, xb] * fx
    return top * (1 - fy) + bot * fy


def rotation(a, b):
    """3×3 rotation carrying [lat, lon] a onto b along the great circle (Rodrigues)."""
    pa, pb = latlon_to_xyz(*np.array(a, dtype=np.float64)), latlon_to_xyz(*np.array(b, dtype=np.float64))
    k = np.cross(pa, pb)
    s, c = np.linalg.norm(k), float(np.dot(pa, pb))
    if s < 1e-12:
        return np.eye(3)
    k /= s
    K = np.array([[0, -k[2], k[1]], [k[2], 0, -k[0]], [-k[1], k[0], 0]])
    return np.eye(3) + s * K + (1 - c) * K @ K


def continent_mask(land, lat, lon):
    """Pixels of the sketch landmass containing (lat, lon), 2 px coastal fringe included; lon wraps."""
    from scipy import ndimage
    lab, _ = ndimage.label(land)
    for a, b in zip(lab[:, 0], lab[:, -1]):
        if a and b and a != b:
            lab[lab == b] = a
    h, w = land.shape
    y, x = min(max(int((90.0 - lat) / 180.0 * h), 0), h - 1), int((lon + 180.0) / 360.0 * w) % w
    if not lab[y, x]:
        raise StageError(f"sketch move: no sketch land at [{lat}, {lon}]")
    return ndimage.binary_dilation(lab == lab[y, x], iterations=2)


def warped_latlon(g, P, seed):
    """Sample positions displaced by a two-scale domain warp (RMS ≈ warp_km, detail_warp_km)."""
    p = g.xyz.copy()
    for amp, freq, s in ((P["warp_km"], P["warp_freq"], 201), (P["detail_warp_km"], P["detail_warp_freq"], 211)):
        if amp <= 0:
            continue
        disp = np.stack([fbm(g.xyz, seed + s + k, 4, freq) for k in range(3)], axis=1)
        p = p + (amp / g.radius_km) * disp / max(float(disp.std()), 1e-12)
    return xyz_to_latlon(p)


def run(ctx) -> dict:
    g = ctx.grid
    P = params(ctx.cfg, "sketch", DEFAULTS)
    sk = ctx.root / "sketch"
    missing = [n for n in SKETCH if not (sk / f"{n}.png").exists()]
    if missing:
        raise StageError(f"sketch files missing {missing}: draw them or run `mapgen.py new-world` first")
    lat, lon = warped_latlon(g, P, ctx.seed)     # the drawing is a loose guide: warp it (user masks are not)
    imgs = {n: load_png01(sk / f"{n}.png") for n in SKETCH}
    moved = []                                   # [[sketch.moves]]: rotate whole landmasses across the sphere
    for mv in P["moves"]:
        comp = continent_mask(imgs["land"] > 0.5, *mv["at"])
        moved.append((rotation(mv["at"], mv["to"]), {n: np.where(comp, im, 0.0) for n, im in imgs.items()}))
        imgs = {n: np.where(comp, 0.0, im) for n, im in imgs.items()}
    p = latlon_to_xyz(lat, lon)
    out = {}
    for n in SKETCH:
        v = sample_equirect(imgs[n], lat, lon)
        for rot, parts in moved:
            qlat, qlon = xyz_to_latlon(p @ rot)          # rot.T applied to row vectors
            v = np.maximum(v, sample_equirect(parts[n], qlat, qlon))
        out[f"sk_{n}"] = v
    for m in MASKS:
        p = ctx.root / "masks" / f"{m}.png"
        out[f"m_{m}"] = sample_equirect(load_mask(p), g.lat, g.lon) if p.exists() else np.zeros(g.n)
    zo2, zg = ZN.masks(g.xyz, ctx.tect.get("zone", []), ctx.seed, g.radius_km)   # config zones on top of the masks
    out["m_o2_zones"] = np.clip(out["m_o2_zones"] + zo2, -1.0, 1.0)
    out["m_gravity_zones"] = np.clip(out["m_gravity_zones"] + zg, -1.0, 1.0)
    return out