#!/usr/bin/env python3
"""
02_annotate.py — add to every USGS event:
  * the sidereal (Lahiri) longitude of all nine grahas at the origin time
    (the Moon's is the one the text uses; the others serve the all-graha view),
  * its nakṣatra in the 27-scheme and the 28-scheme (with Abhijit),
  * the Bṛhat Saṁhitā 32 circle (Vāyu / Agni / Indra / Varuṇa) of that asterism,
  * a Gardner–Knopoff mainshock flag (True = independent event),
  * the country the epicentre lies in, or the nearest country within
    2° of arc (~220 km) for offshore events — subduction-zone quakes off
    Japan, Chile or Sumatra belong to those coasts.

Input : usgs_1980_2025_m45.csv     (from 01_fetch_usgs.py; only M >= 5.0 is kept)
        countries-50m.json          (Natural Earth via world-atlas, TopoJSON)
Output: catalog_annotated.csv.gz

Requirements: numpy pandas scipy pyswisseph geopandas shapely topojson
"""
import argparse
import json
import sys
import time

import geopandas as gpd
import numpy as np
import pandas as pd
from shapely.geometry import shape

import bhukampa_common as bc


def load_countries(path):
    """TopoJSON (world-atlas) -> GeoDataFrame with name + ISO numeric id."""
    import topojson  # pip install topojson
    topo = json.load(open(path))
    tj = topojson.Topology(topo, object_name="countries")
    gj = json.loads(tj.to_geojson())
    rows = []
    for f in gj["features"]:
        rows.append({"name": f["properties"]["name"],
                     "iso_n3": f.get("id"),
                     "geometry": shape(f["geometry"])})
    gdf = gpd.GeoDataFrame(rows, crs="EPSG:4326")
    gdf["geometry"] = gdf.buffer(0)   # heal any self-touching rings
    return gdf


def assign_countries(df, countries, max_deg=2.0):
    pts = gpd.GeoDataFrame(df[["longitude", "latitude"]].copy(),
                           geometry=gpd.points_from_xy(df.longitude, df.latitude),
                           crs="EPSG:4326")
    inside = gpd.sjoin(pts, countries[["name", "geometry"]],
                       how="left", predicate="within")
    inside = inside[~inside.index.duplicated(keep="first")]
    country = inside["name"].copy()
    offshore = country.isna()
    if offshore.any():
        near = gpd.sjoin_nearest(pts[offshore], countries[["name", "geometry"]],
                                 how="left", max_distance=max_deg,
                                 distance_col="deg")
        near = near[~near.index.duplicated(keep="first")]
        country.loc[near.index] = near["name"]
    return country.fillna("(open ocean)"), ~offshore


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--catalog", default="usgs_1980_2025_m45.csv")
    ap.add_argument("--countries", default="countries-50m.json")
    ap.add_argument("--out", default="catalog_annotated.csv.gz")
    ap.add_argument("--max-deg", type=float, default=2.0)
    ap.add_argument("--minmag", type=float, default=5.0)
    a = ap.parse_args()

    t0 = time.time()
    df = pd.read_csv(a.catalog)
    df = df[df["type"].fillna("earthquake").eq("earthquake")].copy()
    df = df[df["mag"] >= a.minmag].copy()      # the service filters on the magnitude
                                                # at query time; a few revised values slip under
    df["time"] = pd.to_datetime(df["time"], utc=True)
    df = df.sort_values("time").reset_index(drop=True)
    print(f"{len(df):,} earthquakes", file=sys.stderr)

    jd = bc.julian_days(df["time"])
    G = bc.graha_sidereal_longitudes(jd)
    df["moon_lon"] = np.round(G["Moon"], 4)
    for g in bc.GRAHAS:
        if g != "Moon":
            df[f"{g}_lon"] = np.round(G[g], 3)
    df["nak27_i"] = bc.nak27_index(df["moon_lon"])
    df["nak28_i"] = bc.nak28_index(df["moon_lon"])
    df["nak27"] = [bc.NAK27[i] for i in df["nak27_i"]]
    df["nak28"] = [bc.NAK28[i] for i in df["nak28_i"]]
    df["mandala"] = df["nak28"].map(bc.MANDALA_OF)
    print(f"moon done  {time.time()-t0:.0f}s", file=sys.stderr)

    df["mainshock"] = bc.decluster_gk(df)
    print(f"decluster done: {df.mainshock.sum():,} mainshocks  {time.time()-t0:.0f}s",
          file=sys.stderr)

    countries = load_countries(a.countries)
    df["country"], df["onland"] = assign_countries(df, countries, a.max_deg)
    print(f"countries done  {time.time()-t0:.0f}s", file=sys.stderr)

    cols = ["time", "latitude", "longitude", "depth", "mag", "magType", "place",
            "id", "moon_lon", "nak27", "nak28", "mandala", "mainshock",
            "country", "onland"] + [f"{g}_lon" for g in bc.GRAHAS if g != "Moon"]
    df[cols].to_csv(a.out, index=False, date_format="%Y-%m-%dT%H:%M:%S.%fZ")
    print(f"wrote {a.out}", file=sys.stderr)


if __name__ == "__main__":
    main()