"""Orbital image coverage of the Moon and Mars.

(a) Mars, HiRISE (0.25-0.5 m/pixel): union area of all RED-channel RDR footprints in
    the PDS cumulative index, by Monte Carlo point-in-polygon with 4 million
    random points; also by year and by map scale.
(b) Moon LROC NAC and Mars CTX / HiRISE from random-point ODE queries
    (s04_fetch_ode_samples.py): fraction of the surface imaged at given pixel scales.

Outputs: data/surface_coverage_summary.csv, data/hirise_coverage_by_year.csv
"""
import os

import numpy as np
import pandas as pd

from common import DATA, RAW, SEED

MARS_AREA_KM2 = 4 * np.pi * 3389.5 ** 2
MOON_AREA_KM2 = 4 * np.pi * 1737.4 ** 2


def load_hirise():
    cols = {1: "file", 5: "product_id", 7: "target", 11: "start_time", 17: "lines", 18: "samples",
            20: "incidence", 35: "min_lat", 36: "max_lat", 37: "min_lon", 38: "max_lon", 39: "map_scale",
            46: "c1lat", 47: "c1lon", 48: "c2lat", 49: "c2lon", 50: "c3lat", 51: "c3lon", 52: "c4lat", 53: "c4lon"}
    df = pd.read_csv(os.path.join(RAW, "RDRCUMINDEX.TAB"), header=None, usecols=list(cols), quotechar='"',
                     skipinitialspace=True)
    df = df.rename(columns=cols)
    for c in ["product_id", "target", "start_time"]:
        df[c] = df[c].astype(str).str.strip()
    n_all = len(df)
    df = df[df.product_id.str.endswith("_RED") & (df.target == "MARS")].copy()
    df["year"] = df.start_time.str[:4].astype(int)
    for c in ["c1lat", "c1lon", "c2lat", "c2lon", "c3lat", "c3lon", "c4lat", "c4lon", "map_scale"]:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    df = df.dropna(subset=["c1lat", "c1lon", "c2lat", "c2lon", "c3lat", "c3lon", "c4lat", "c4lon"])
    return df.reset_index(drop=True), n_all


def in_quad(px, py, qx, qy):
    """Vectorised point-in-convex-or-simple-quadrilateral (ray casting)."""
    inside = np.zeros(px.shape, bool)
    j = 3
    for i in range(4):
        xi, yi, xj, yj = qx[i], qy[i], qx[j], qy[j]
        cond = ((yi > py) != (yj > py)) & (px < (xj - xi) * (py - yi) / (yj - yi + 1e-300) + xi)
        inside ^= cond
        j = i
    return inside


def hirise_union(df, n_pts=4_000_000, seed=SEED + 7):
    rng = np.random.default_rng(seed)
    z = rng.uniform(-1, 1, n_pts)
    lat = np.degrees(np.arcsin(z))
    lon = rng.uniform(0, 360, n_pts)
    cell = 0.5
    nlon = int(360 / cell)
    ilat = np.clip(((lat + 90) / cell).astype(int), 0, int(180 / cell) - 1)
    ilon = np.clip((lon / cell).astype(int), 0, nlon - 1)
    cid = ilat * nlon + ilon
    order = np.argsort(cid, kind="stable")
    cid_s, lat_s, lon_s = cid[order], lat[order], lon[order]
    starts = np.searchsorted(cid_s, np.arange(int(180 / cell) * nlon + 1))
    hit_year = np.full(n_pts, 9999, np.int32)       # earliest year covered
    hit_scale = np.full(n_pts, 99.0, np.float32)     # finest map scale
    clat = df[["c1lat", "c2lat", "c3lat", "c4lat"]].values
    clon = df[["c1lon", "c2lon", "c3lon", "c4lon"]].values % 360.0
    years = df.year.values
    scales = df.map_scale.values
    for k in range(len(df)):
        la, lo = clat[k], clon[k]
        la0 = la.mean()
        polar = np.abs(la).max() > 84.0
        if polar:
            sgn = 1.0 if la0 > 0 else -1.0
            rr = 90.0 - sgn * la
            qx, qy = rr * np.cos(np.radians(lo)), rr * np.sin(np.radians(lo))
            i0 = int(np.clip(((la.min() + 90) / cell), 0, 359))
            i1 = int(np.clip(((la.max() + 90) / cell), 0, 359))
            cells = []
            if sgn > 0:
                i1 = 359
            else:
                i0 = 0
            sl = slice(starts[i0 * nlon], starts[(i1 + 1) * nlon])
            idx = np.arange(sl.start, sl.stop)
            rp = 90.0 - sgn * lat_s[idx]
            px, py = rp * np.cos(np.radians(lon_s[idx])), rp * np.sin(np.radians(lon_s[idx]))
        else:
            lo0 = lo[0]
            dlo = (lo - lo0 + 180.0) % 360.0 - 180.0      # unwrap around the first corner
            c = np.cos(np.radians(la0))
            qx, qy = dlo * c, la
            j0 = int(np.floor(((lo0 + dlo.min()) % 360.0) / cell))
            nj = int(np.ceil((dlo.max() - dlo.min()) / cell)) + 1
            i0 = int(np.clip((la.min() + 90) / cell, 0, 359))
            i1 = int(np.clip((la.max() + 90) / cell, 0, 359))
            parts = []
            for i in range(i0, i1 + 1):
                for dj in range(nj + 1):
                    cc = i * nlon + (j0 + dj) % nlon
                    if starts[cc + 1] > starts[cc]:
                        parts.append(np.arange(starts[cc], starts[cc + 1]))
            if not parts:
                continue
            idx = np.concatenate(parts)
            px = ((lon_s[idx] - lo0 + 180.0) % 360.0 - 180.0) * c
            py = lat_s[idx]
        if len(idx) == 0:
            continue
        ins = in_quad(px, py, qx, qy)
        if ins.any():
            ii = idx[ins]
            hit_year[ii] = np.minimum(hit_year[ii], years[k])
            hit_scale[ii] = np.minimum(hit_scale[ii], scales[k])
    return hit_year, hit_scale, lat_s


def quad_area_km2(df, R=3389.5):
    """Approximate footprint areas (planar shoelace in local equirectangular coordinates)."""
    la = df[["c1lat", "c2lat", "c3lat", "c4lat"]].values
    lo = df[["c1lon", "c2lon", "c3lon", "c4lon"]].values
    dlo = (lo - lo[:, :1] + 180.0) % 360.0 - 180.0
    x = np.radians(dlo) * np.cos(np.radians(la.mean(axis=1)))[:, None] * R
    y = np.radians(la) * R
    a = 0.5 * np.abs(np.sum(x * np.roll(y, -1, axis=1) - np.roll(x, -1, axis=1) * y, axis=1))
    return a


def wilson(k, n, z=1.96):
    if n == 0:
        return (np.nan, np.nan)
    p = k / n
    den = 1 + z * z / n
    c = (p + z * z / (2 * n)) / den
    h = z * np.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / den
    return (max(0.0, c - h), min(1.0, c + h))


def main():
    rows = []
    # ---------------- HiRISE union from the PDS index ----------------
    df, n_all = load_hirise()
    hy, hs, lat_s = hirise_union(df)
    n = len(hy)
    cov = (hy < 9999).mean()
    err = np.sqrt(cov * (1 - cov) / n)
    areas = quad_area_km2(df[np.abs(df[["c1lat", "c2lat", "c3lat", "c4lat"]]).max(axis=1) < 84])
    print(f"[surf] HiRISE RDR index: {n_all} rows, {len(df)} RED Mars images, years {df.year.min()}-{df.year.max()}")
    print(f"[surf] HiRISE union coverage of Mars: {100 * cov:.3f}% +- {100 * err:.3f}% (Monte Carlo, {n} points); "
          f"sum of footprint areas (|lat|<84) = {100 * areas.sum() / MARS_AREA_KM2:.2f}% of Mars")
    rows.append(dict(body="Mars", dataset="HiRISE RED (PDS RDR index)", metric="union fraction of surface, any scale",
                     value=cov, lo=cov - 1.96 * err, hi=cov + 1.96 * err, n=n))
    f25 = (hs <= 0.3).mean()
    rows.append(dict(body="Mars", dataset="HiRISE RED (PDS RDR index)", metric="union fraction at <=0.3 m/pixel map scale",
                     value=f25, lo=f25 - 1.96 * np.sqrt(f25 / n), hi=f25 + 1.96 * np.sqrt(f25 / n), n=n))
    rows.append(dict(body="Mars", dataset="HiRISE RED (PDS RDR index)", metric="sum of footprint areas / surface (repeat coverage included)",
                     value=areas.sum() / MARS_AREA_KM2, lo=np.nan, hi=np.nan, n=len(df)))
    by = []
    for y in range(int(df.year.min()), int(df.year.max()) + 1):
        by.append(dict(year=y, n_images_cum=int((df.year <= y).sum()), union_fraction=float((hy <= y).mean())))
    pd.DataFrame(by).to_csv(os.path.join(DATA, "hirise_coverage_by_year.csv"), index=False)
    for lo_, hi_ in [(0, 30), (30, 60), (60, 90)]:
        m = (np.abs(lat_s) >= lo_) & (np.abs(lat_s) < hi_)
        v = (hy[m] < 9999).mean()
        rows.append(dict(body="Mars", dataset="HiRISE RED (PDS RDR index)", metric=f"union fraction, |lat| {lo_}-{hi_} deg",
                         value=v, lo=np.nan, hi=np.nan, n=int(m.sum())))
    # ---------------- ODE random-point samples ----------------
    for tag, body, ds, thr_list in [("moon_nac", "Moon", "LROC NAC EDR (ODE point sample)", [0.6, 0.8, 1.0, 1.5, 2.0]),
                                    ("mars_ctx", "Mars", "CTX EDR (ODE point sample)", [6.0, 7.0]),
                                    ("mars_hirise", "Mars", "HiRISE RDR (ODE point sample; cross-check of the index union)", [])]:
        fn = os.path.join(DATA, f"{tag}_samples.csv")
        if not os.path.exists(fn):
            print("[surf] missing", fn)
            continue
        d = pd.read_csv(fn)
        n = len(d)
        k = int((d.n_images > 0).sum())
        lo_, hi_ = wilson(k, n)
        rows.append(dict(body=body, dataset=ds, metric="fraction of random points with >=1 image", value=k / n, lo=lo_, hi=hi_, n=n))
        k2 = int((d.n_images >= 2).sum())
        lo_, hi_ = wilson(k2, n)
        rows.append(dict(body=body, dataset=ds, metric="fraction with >=2 images", value=k2 / n, lo=lo_, hi=hi_, n=n))
        if not thr_list:          # HiRISE: coverage only (the ODE resolution field is not in m/pixel for this product type)
            continue
        rows.append(dict(body=body, dataset=ds, metric="median number of images per point", value=float(d.n_images.median()), lo=np.nan, hi=np.nan, n=n))
        rows.append(dict(body=body, dataset=ds, metric="median best pixel scale (m)", value=float(d.best_res.median()), lo=float(d.best_res.quantile(0.05)), hi=float(d.best_res.quantile(0.95)), n=n))
        for thr in thr_list:
            k = int((d.best_res <= thr).sum())
            lo_, hi_ = wilson(k, n)
            rows.append(dict(body=body, dataset=ds, metric=f"fraction with best pixel scale <= {thr} m", value=k / n, lo=lo_, hi=hi_, n=n))
            k = int((d.best_res_goodlight <= thr).sum())
            lo_, hi_ = wilson(k, n)
            rows.append(dict(body=body, dataset=ds, metric=f"fraction with <= {thr} m and incidence 20-85 deg", value=k / n, lo=lo_, hi=hi_, n=n))
    out = pd.DataFrame(rows)
    out.to_csv(os.path.join(DATA, "surface_coverage_summary.csv"), index=False)
    pd.set_option("display.width", 220)
    pd.set_option("display.max_colwidth", 70)
    print(out.round(4).to_string(index=False))
    return out


if __name__ == "__main__":
    main()