#!/usr/bin/env python3
"""
05_ancient_regions.py — the regions Varāhamihira names, tested; and, for
the Nepal region, the nakṣatra of every graha, not only the Moon.

Bṛhat Saṁhitā 32.22 puts the Vaidehakas (Videha = Mithilā: north Bihar and
the Nepal Tarai) and the Kirātas (the Himalayan hill peoples) under the
Varuṇa circle; 14.29 places Kāśmīra in the north-eastern ninth of the
Kūrma-vibhāga but assigns it to no circle.  Of the peoples he names, only
these two areas have enough M >= 5.0 earthquakes since 1980 to test.

Regions (WGS84):
  Kāśmīra  — the Kashmir Valley and its rim, 73.5–76.5 E, 33.0–35.5 N
  Videha   — Nepal (country polygon) ∪ north Bihar (83.5–88.3 E, 25.4–27.6 N)
             — "the Nepal region"
  Nepal    — the country polygon alone, for comparison

Two nulls.  For the Moon (and any fast graha) the natural null is the
fraction of TIME the graha spends in each nakṣatra over 1980–2025.  For the
slow grahas — Jupiter, Saturn, the nodes — that null is confounded with the
growth of the catalogue (far more events are recorded after 2004 than
before), so the honest reference is the WORLD catalogue itself, same
magnitude floor, same raw/mainshock choice: does the sky over Nepal's
earthquakes differ from the sky over everyone's?  Both are computed and
reported for every graha; the essay uses the time null for the Moon and
the world null for the rest.

For each region × graha × {raw, mainshocks} × {27, 28 asterisms}: counts,
chi-square against both nulls, power at w = 0.3, and the one-sided binomial
test of the Varuṇa-circle share against the world share (and, for the Moon,
against the time share as well).

Input : catalog_annotated.csv.gz, countries-50m.json
Output: regions.json, regions.csv
"""
import argparse
import json

import geopandas as gpd
import numpy as np
import pandas as pd
from scipy import stats
from shapely.geometry import box
from shapely.ops import unary_union

import bhukampa_common as bc

GRAHAS = bc.GRAHAS


def regions(countries_path):
    import topojson
    from shapely.geometry import shape
    topo = json.load(open(countries_path))
    gj = json.loads(topojson.Topology(topo, object_name="countries").to_geojson())
    nepal = [shape(f["geometry"]) for f in gj["features"] if f["properties"]["name"] == "Nepal"][0]
    return {
        "Kāśmīra": box(73.5, 33.0, 76.5, 35.5),
        "Videha": unary_union([nepal, box(83.5, 25.4, 88.3, 27.6)]),
        "Nepal": nepal,
    }


def lon_col(g):
    return "moon_lon" if g == "Moon" else f"{g}_lon"


def counts(d, g):
    lon = d[lon_col(g)].to_numpy(float)
    return (np.bincount(bc.nak27_index(lon), minlength=27)[:27],
            np.bincount(bc.nak28_index(lon), minlength=28)[:28])


def varuna_share_test(c28, p0):
    members = [bc.NAK28.index(n) for n in bc.MANDALA["Varuṇa"]]
    k = int(c28[members].sum()); n = int(c28.sum())
    if n == 0:
        return {"k": 0, "n": 0, "p0": p0, "ratio": float("nan"), "p_greater": float("nan"), "p_two_sided": float("nan")}
    return {"k": k, "n": n, "p0": float(p0), "observed": k / n, "ratio": (k / n) / p0,
            "p_greater": float(stats.binomtest(k, n, p0, alternative="greater").pvalue),
            "p_two_sided": float(stats.binomtest(k, n, p0).pvalue)}


def circle_shares(c28, ref28):
    """observed share, reference share and two-sided binomial p for each circle"""
    out = {}
    n = int(c28.sum())
    for circ, names in bc.MANDALA.items():
        idx = [bc.NAK28.index(x) for x in names]
        k = int(c28[idx].sum()); p0 = float(ref28[idx].sum())
        out[circ] = {"k": k, "n": n, "p0": p0, "observed": k / n if n else float("nan"),
                     "ratio": (k / n) / p0 if n and p0 else float("nan"),
                     "p_greater": float(stats.binomtest(k, n, p0, alternative="greater").pvalue) if n else float("nan"),
                     "p_two_sided": float(stats.binomtest(k, n, p0).pvalue) if n else float("nan")}
    return out


def analyse_graha(d, g, world_ref, time_ref):
    c27, c28 = counts(d, g)
    w27, w28 = world_ref[g]
    t27, t28 = time_ref[g]
    rw27, rw28 = bc.chisq_gof(c27, w27), bc.chisq_gof(c28, w28)
    rt27, rt28 = bc.chisq_gof(c27, t27), bc.chisq_gof(c28, t28)
    members = [bc.NAK28.index(n) for n in bc.MANDALA["Varuṇa"]]
    return {"counts27": c27.tolist(), "counts28": c28.tolist(),
            "world": {"chi2_27": rw27["chi2"], "p27": rw27["p"], "chi2_28": rw28["chi2"], "p28": rw28["p"], "w27": rw27["w"]},
            "time": {"chi2_27": rt27["chi2"], "p27": rt27["p"], "chi2_28": rt28["chi2"], "p28": rt28["p"], "w27": rt27["w"]},
            "power27": bc.chisq_power(int(c27.sum()), 26), "power28": bc.chisq_power(int(c28.sum()), 27),
            "circles_world": circle_shares(c28, w28),
            "circles_time": circle_shares(c28, t28),
            "varuna_world": varuna_share_test(c28, float(w28[members].sum())),
            "varuna_time": varuna_share_test(c28, float(t28[members].sum()))}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--catalog", default="catalog_annotated.csv.gz")
    ap.add_argument("--countries", default="countries-50m.json")
    ap.add_argument("--json", default="regions.json")
    ap.add_argument("--csv", default="regions.csv")
    a = ap.parse_args()

    df = pd.read_csv(a.catalog)
    pts = gpd.GeoSeries(gpd.points_from_xy(df.longitude, df.latitude), crs="EPSG:4326")

    # time-weighted null for every graha (6-hourly grid over the window)
    t = pd.date_range("1980-01-01", "2026-01-01", freq="6h", tz="UTC", inclusive="left")
    G = bc.graha_sidereal_longitudes(bc.julian_days(pd.Series(t)))
    time_ref = {}
    for g in GRAHAS:
        c27 = np.bincount(bc.nak27_index(G[g]), minlength=27)[:27]
        c28 = np.bincount(bc.nak28_index(G[g]), minlength=28)[:28]
        time_ref[g] = (c27 / c27.sum(), c28 / c28.sum())

    out = {"grahas": GRAHAS,
           "expected_time": {g: {"p27": time_ref[g][0].tolist(), "p28": time_ref[g][1].tolist()} for g in GRAHAS},
           "world": {}, "regions": {}}
    world_ref = {}
    for tag, d in (("raw", df), ("main", df[df.mainshock])):
        world_ref[tag] = {}
        out["world"][tag] = {"n": int(len(d)), "grahas": {}}
        for g in GRAHAS:
            c27, c28 = counts(d, g)
            world_ref[tag][g] = (c27 / c27.sum(), c28 / c28.sum())
            out["world"][tag]["grahas"][g] = {"counts27": c27.tolist(), "counts28": c28.tolist(),
                                              "p27_time": bc.chisq_gof(c27, time_ref[g][0])["p"],
                                              "varuna_share": float(c28[[bc.NAK28.index(n) for n in bc.MANDALA["Varuṇa"]]].sum() / c28.sum())}

    rows = []
    for name, geom in regions(a.countries).items():
        sub = df[pts.within(geom).to_numpy()]
        res = {"bounds": list(geom.bounds)}
        for tag, d in (("raw", sub), ("main", sub[sub.mainshock])):
            res[tag] = {"n": int(len(d)),
                        "by_mandala": {c: int(d.mandala.eq(c).sum()) for c in bc.MANDALA},
                        "top": [{"time": str(r.time)[:16], "mag": float(r.mag), "place": str(r.place),
                                 "nak28": r.nak28, "mandala": r.mandala}
                                for r in d.sort_values("mag", ascending=False).head(8).itertuples()],
                        "grahas": {g: analyse_graha(d, g, world_ref[tag], time_ref) for g in GRAHAS}}
            for g in GRAHAS:
                r = res[tag]["grahas"][g]
                rows.append({"region": name, "catalog": tag, "graha": g, "n": res[tag]["n"],
                             "p27_time": r["time"]["p27"], "p27_world": r["world"]["p27"],
                             "p28_world": r["world"]["p28"], "power27": r["power27"],
                             "varuna_k": r["varuna_world"]["k"], "varuna_ratio_world": r["varuna_world"]["ratio"],
                             "varuna_p_world": r["varuna_world"]["p_greater"],
                             "varuna_ratio_time": r["varuna_time"]["ratio"], "varuna_p_time": r["varuna_time"]["p_greater"]})
        out["regions"][name] = res
        m = res["main"]["grahas"]["Moon"]
        print(f"{name:8s} n raw={res['raw']['n']} main={res['main']['n']}  Moon(main): p27 time={m['time']['p27']:.3g} world={m['world']['p27']:.3g} "
              f"power={m['power27']:.2f}  Varuṇa {m['varuna_time']['k']}/{m['varuna_time']['n']} = {m['varuna_time']['ratio']:.2f}× (p={m['varuna_time']['p_greater']:.3g})")
        for g in GRAHAS:
            r = res["main"]["grahas"][g]
            print(f"    {g:8s} p27(world)={r['world']['p27']:.3g}  p27(time)={r['time']['p27']:.3g}  Varuṇa(world) {r['varuna_world']['ratio']:.2f}× p={r['varuna_world']['p_greater']:.3g}")
    pd.DataFrame(rows).to_csv(a.csv, index=False, float_format="%.5g")
    json.dump(out, open(a.json, "w", encoding="utf-8"), ensure_ascii=False, separators=(",", ":"))


if __name__ == "__main__":
    main()
