"""Single-point Taiwan intensity prediction from the trained P-Alert GMM.

Input : earthquake (M, depth, epicenter lat/lon) + ONE target lat/lon
Output: PGA (gal), PGV (cm/s), CWA intensity

Self-contained: needs only numpy + onnxruntime. No dependency on the rest of
this repo, so you can drop the file (plus the .onnx) into your own program.

CLI:
  python predict_point.py --M 6.5 --depth 10 --ev-lat 23.9 --ev-lon 121.6 \
                          --t-lat 24.0 --t-lon 121.7

Library:
  from predict_point import predict_point
  pga, pgv, level, label = predict_point(6.5, 10.0, 23.9, 121.6, 24.0, 121.7)
"""
import argparse
import math
import os
import sys

import numpy as np
import onnxruntime as ort

_HERE = os.path.dirname(os.path.abspath(__file__))
_MODEL_NAME = "TDN-IntensityGMM-TW-524k-v1.onnx"
# the artifact sits at the repo root when published, under data/ during development
DEFAULT_MODEL = next(
    (p for p in (os.path.join(_HERE, _MODEL_NAME), os.path.join(_HERE, "data", _MODEL_NAME))
     if os.path.exists(p)),
    os.path.join(_HERE, _MODEL_NAME),
)

PGA_BOUNDS = [0.8, 2.5, 8, 25, 80, 140, 250, 440, 800]   # gal
PGV_BOUNDS = [0.2, 0.7, 1.9, 5.7, 15, 30, 50, 80, 140]   # cm/s
LEVELS = ["0", "1", "2", "3", "4", "5弱", "5強", "6弱", "6強", "7"]

R_EARTH_KM = 6371.0


def make_features(M, depth, ev_lat, ev_lon, t_lat, t_lon):
    """The model's 10 input features, in the exact trained order."""
    p1, p2 = math.radians(ev_lat), math.radians(t_lat)
    dl = math.radians(t_lon - ev_lon)
    a = math.sin((p2 - p1) / 2) ** 2 + math.cos(p1) * math.cos(p2) * math.sin(dl / 2) ** 2
    dist = 2 * R_EARTH_KM * math.asin(math.sqrt(min(1.0, a)))
    lnR = math.log(max(math.hypot(dist, depth), 3.0))
    az = math.atan2(
        math.sin(dl) * math.cos(p2),
        math.cos(p1) * math.sin(p2) - math.sin(p1) * math.cos(p2) * math.cos(dl),
    )
    return [M, depth, dist, lnR, ev_lat, ev_lon, t_lat, t_lon, math.sin(az), math.cos(az)]


def level_from_pga_pgv(pga, pgv):
    """CWA intensity index (0..9) = max of the PGA bin and the PGV bin."""
    lv = sum(1 for b in PGA_BOUNDS if pga >= b)
    lv = max(lv, sum(1 for b in PGV_BOUNDS if pgv >= b))
    return lv


class IntensityGMM:
    def __init__(self, onnx_path=DEFAULT_MODEL):
        self.sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
        self.input_name = self.sess.get_inputs()[0].name

    def predict_points(self, M, depth, ev_lat, ev_lon, points):
        """points = [(t_lat, t_lon), ...] -> (pga[], pgv[], level[])."""
        X = np.array(
            [make_features(M, depth, ev_lat, ev_lon, la, lo) for la, lo in points],
            dtype=np.float32,
        )
        pga, pgv = self.sess.run(None, {self.input_name: X})
        pga = pga.ravel().astype(float)
        pgv = pgv.ravel().astype(float)
        levels = [level_from_pga_pgv(float(p), float(v)) for p, v in zip(pga, pgv)]
        return pga, pgv, levels

    def predict_one(self, M, depth, ev_lat, ev_lon, t_lat, t_lon):
        """One target point -> (PGA gal, PGV cm/s, intensity index, intensity label)."""
        pga, pgv, levels = self.predict_points(M, depth, ev_lat, ev_lon, [(t_lat, t_lon)])
        lvl = levels[0]
        return float(pga[0]), float(pgv[0]), lvl, LEVELS[lvl]


def predict_point(M, depth, ev_lat, ev_lon, t_lat, t_lon, model=None):
    """Convenience one-liner for a single point. Returns (pga, pgv, level, label)."""
    return (model or IntensityGMM()).predict_one(M, depth, ev_lat, ev_lon, t_lat, t_lon)


def main():
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass

    ap = argparse.ArgumentParser(description="Single-point PGA/PGV/intensity prediction")
    ap.add_argument("--M", type=float, required=True, help="magnitude (ML)")
    ap.add_argument("--depth", type=float, required=True, help="focal depth (km)")
    ap.add_argument("--ev-lat", type=float, required=True, dest="ev_lat", help="epicenter latitude")
    ap.add_argument("--ev-lon", type=float, required=True, dest="ev_lon", help="epicenter longitude")
    ap.add_argument("--t-lat", type=float, required=True, dest="t_lat", help="target latitude")
    ap.add_argument("--t-lon", type=float, required=True, dest="t_lon", help="target longitude")
    ap.add_argument("--model", default=DEFAULT_MODEL, help="path to the ONNX model")
    args = ap.parse_args()

    pga, pgv, level, label = predict_point(
        args.M, args.depth, args.ev_lat, args.ev_lon, args.t_lat, args.t_lon,
        IntensityGMM(args.model),
    )
    print(f"target  ({args.t_lat}, {args.t_lon})")
    print(f"PGA     {pga:.3f} gal")
    print(f"PGV     {pgv:.3f} cm/s")
    print(f"震度    {label} ({level})")


if __name__ == "__main__":
    main()