"""Map projections of the equirectangular relief: Equal Earth, Mollweide (equal-area) and orthographic globes."""
from __future__ import annotations

import numpy as np
from PIL import Image, ImageDraw

from .sketch import sample_equirect
from .sphere import east_north, latlon_to_xyz, xyz_to_latlon

BACKGROUND = (16, 18, 24)
_A1, _A2, _A3, _A4 = 1.340264, -0.081106, 0.000893, 0.003796
_M = np.sqrt(3.0) / 2.0
EE_XMAX = 2.0 * np.sqrt(3.0) * np.pi / (3.0 * _A1)
EE_YMAX = _A1 * np.pi / 3 + _A2 * (np.pi / 3) ** 3 + _A3 * (np.pi / 3) ** 7 + _A4 * (np.pi / 3) ** 9
MW_XMAX, MW_YMAX = 2.0 * np.sqrt(2.0), np.sqrt(2.0)


def equal_earth_forward(lat, lon):
    t = np.arcsin(_M * np.sin(np.radians(lat)))
    d = _A1 + 3 * _A2 * t**2 + 7 * _A3 * t**6 + 9 * _A4 * t**8
    x = 2 * np.sqrt(3.0) * np.radians(lon) * np.cos(t) / (3 * d)
    y = _A1 * t + _A2 * t**3 + _A3 * t**7 + _A4 * t**9
    return x, y


def equal_earth_inverse(x, y):
    t = np.asarray(y, dtype=np.float64) / _A1
    for _ in range(12):
        f = _A1 * t + _A2 * t**3 + _A3 * t**7 + _A4 * t**9 - y
        t = t - f / (_A1 + 3 * _A2 * t**2 + 7 * _A3 * t**6 + 9 * _A4 * t**8)
    d = _A1 + 3 * _A2 * t**2 + 7 * _A3 * t**6 + 9 * _A4 * t**8
    lat = np.degrees(np.arcsin(np.clip(np.sin(t) / _M, -1, 1)))
    lon = np.degrees(3 * np.asarray(x) * d / (2 * np.sqrt(3.0) * np.cos(t)))
    return lat, lon


def mollweide_forward(lat, lon):
    phi = np.radians(lat)
    t = phi.copy() if isinstance(phi, np.ndarray) else np.array(phi, dtype=np.float64)
    for _ in range(30):
        f = 2 * t + np.sin(2 * t) - np.pi * np.sin(phi)
        t = t - f / np.maximum(2 + 2 * np.cos(2 * t), 1e-12)
    return MW_XMAX / np.pi * np.radians(lon) * np.cos(t), MW_YMAX * np.sin(t)


def mollweide_inverse(x, y):
    t = np.arcsin(np.clip(np.asarray(y) / MW_YMAX, -1, 1))
    lat = np.degrees(np.arcsin(np.clip((2 * t + np.sin(2 * t)) / np.pi, -1, 1)))
    lon = np.degrees(np.pi * np.asarray(x) / (MW_XMAX * np.maximum(np.cos(t), 1e-12)))
    return lat, lon


def orthographic_inverse(x, y, lat0, lon0):
    """Unit-disc view coords → lat/lon on the visible hemisphere centred on (lat0, lon0); ok = inside the disc."""
    x, y = np.asarray(x, dtype=np.float64), np.asarray(y, dtype=np.float64)
    rho2 = x**2 + y**2
    ok = rho2 <= 1.0
    z = np.sqrt(np.clip(1.0 - rho2, 0.0, 1.0))
    c = latlon_to_xyz(np.array([lat0]), np.array([lon0]))
    e, n = east_north(c)
    p = x[..., None] * e[0] + y[..., None] * n[0] + z[..., None] * c[0]
    lat, lon = xyz_to_latlon(p)
    return lat, lon, ok


def _sample_rgb(img, lat, lon):
    """Bilinear RGB samples, read straight from the uint8 channels (each sampled value promotes to float64 exactly,
    so no full-size float copy of a big image is needed)."""
    return np.stack([sample_equirect(img[..., k], lat, lon) for k in range(3)], axis=-1)


def reproject(img, proj, width):
    """Equirectangular RGB → equal-area world map ('equal_earth' | 'mollweide')."""
    xmax, ymax, inv = {"equal_earth": (EE_XMAX, EE_YMAX, equal_earth_inverse),
                       "mollweide": (MW_XMAX, MW_YMAX, mollweide_inverse)}[proj]
    height = int(round(width * ymax / xmax))
    xs = ((np.arange(width) + 0.5) / width * 2 - 1) * xmax
    ys = (1 - (np.arange(height) + 0.5) / height * 2) * ymax
    X, Y = np.meshgrid(xs, ys)
    lat, lon = inv(X, Y)
    ok = np.isfinite(lat) & np.isfinite(lon) & (np.abs(lon) <= 180.0)
    rgb = _sample_rgb(img, np.where(ok, lat, 0.0), np.where(ok, lon, 0.0))
    out = np.where(ok[..., None], rgb, np.array(BACKGROUND, dtype=np.float64))
    return np.clip(np.round(out), 0, 255).astype(np.uint8)


def globe(img, lat0, lon0, size):
    """Orthographic view centred on (lat0, lon0), gentle limb darkening."""
    v = (np.arange(size) + 0.5) / size * 2 - 1
    X, Y = np.meshgrid(v, -v)
    lat, lon, ok = orthographic_inverse(X, Y, lat0, lon0)
    rgb = _sample_rgb(img, lat, lon) * (0.72 + 0.28 * np.sqrt(np.clip(1 - X**2 - Y**2, 0, 1)))[..., None]
    out = np.where(ok[..., None], rgb, np.array(BACKGROUND, dtype=np.float64))
    return np.clip(np.round(out), 0, 255).astype(np.uint8)


def graticule(arr, proj, step=30):
    """Draw a light lat/lon grid (every `step`°) onto an equal-area world map produced by `reproject`."""
    fwd, xmax, ymax = {"equal_earth": (equal_earth_forward, EE_XMAX, EE_YMAX),
                       "mollweide": (mollweide_forward, MW_XMAX, MW_YMAX)}[proj]
    h, w = arr.shape[:2]
    im = Image.fromarray(arr)
    draw = ImageDraw.Draw(im)
    to_px = lambda x, y: list(zip(((x / xmax + 1) / 2 * w).tolist(), ((1 - y / ymax) / 2 * h).tolist()))
    t = np.linspace(-90, 90, 181)
    for lon in range(-180, 181, step):
        draw.line(to_px(*fwd(t, np.full_like(t, float(lon)))), fill=(200, 210, 225), width=1)
    s = np.linspace(-180, 180, 361)
    for lat in range(-90 + step, 90, step):
        draw.line(to_px(*fwd(np.full_like(s, float(lat)), s)), fill=(200, 210, 225), width=1)
    return np.array(im)


def continent_centres(g, ocean, min_share=0.03):
    """Centroids (lat, lon) of land bodies holding ≥ min_share of all land, largest first."""
    from .graph import components
    land = ~np.asarray(ocean)
    lab = components(g, land)
    area = np.bincount(lab[land], weights=g.area_km2[land])
    out = []
    for k in np.argsort(-area):
        if area[k] < min_share * area.sum():
            break
        m = lab == k
        c = (g.xyz[m] * g.area_km2[m, None]).sum(axis=0)
        la, lo = xyz_to_latlon(c[None])
        out.append((float(la[0]), float(lo[0])))
    return out


def write_all(relief, pdir, g, ocean, width, globe_size=1024):
    """Write proj_equal_earth.png, proj_mollweide.png, globe_*.png and globes_sheet.png into pdir."""
    for proj in ("equal_earth", "mollweide"):
        Image.fromarray(graticule(reproject(relief, proj, width), proj)).save(pdir / f"proj_{proj}.png")
    views = [("north_pole", 90.0, 0.0), ("south_pole", -90.0, 0.0)]
    views = [(f"continent_{i + 1}", la, lo) for i, (la, lo) in enumerate(continent_centres(g, ocean))] + views
    tiles = []
    for name, la, lo in views:
        im = Image.fromarray(globe(relief, la, lo, globe_size))
        im.save(pdir / f"globe_{name}.png")
        tiles.append((f"{name} ({la:.0f}°, {lo:.0f}°)", im))
    cols, t = 4, globe_size // 2
    rows = (len(tiles) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * t, rows * (t + 16)), BACKGROUND)
    draw = ImageDraw.Draw(sheet)
    for k, (label, im) in enumerate(tiles):
        x, y = (k % cols) * t, (k // cols) * (t + 16)
        sheet.paste(im.resize((t, t)), (x, y + 16))
        draw.text((x + 4, y + 2), label, fill=(230, 230, 230))
    sheet.save(pdir / "globes_sheet.png")