"""Rapid-intensification probabilities from the exported models (numpy + pandas only, no pickle).

    import pandas as pd
    from predict import load, predict
    models = load("models.json")
    fixes = pd.read_csv("points.csv.gz", keep_default_na=False, na_values=[""])   # from the HF dataset
    p = predict(models["part2_M1c_EP"], fixes[fixes.basin == "EP"])

Each input row is one best-track fix; the columns needed are listed in each model's "features" entry and are
the columns of the dataset's points.csv.gz. The output is the calibrated probability that the 24-h intensity
change starting at that fix is at least 30 kt (RI30).
"""
import json
import numpy as np
import pandas as pd

MISSING_FLAGS = ["dV6", "dV12", "speed_kmh", "u_kmh", "v_kmh", "poleward_kmh"]
BASINS = ["NA", "EP", "WP", "NI", "SI", "SP"]


def load(path="models.json"):
    return json.load(open(path))["models"]


def design(df, features, basin_dummies=False):
    X = df[features].astype(float).copy()
    for c in MISSING_FLAGS:
        if c in X:
            X[c + "_na"] = X[c].isna().astype(float)
            X[c] = X[c].fillna(0.0)
    X["V2"] = X.V ** 2
    X["dist_land_km"] = np.log1p(X.dist_land_km)
    X["age_h"] = np.log1p(X.age_h)
    if basin_dummies:
        for b in BASINS[1:]:
            X["b_" + b] = (df.basin == b).astype(float)
    return X


def predict(model, df, calibrated=True):
    X = design(df, model["features"], model["basin_dummies"])
    assert list(X.columns) == model["columns"], "column order differs from the fitted model"
    Z = (X.values - np.array(model["scaler_mean"])) / np.array(model["scaler_scale"])
    p = 1.0 / (1.0 + np.exp(-(Z @ np.array(model["coef"]) + model["intercept"])))
    if not calibrated:
        return p
    iso = model["isotonic"]
    if "per_basin" in iso:  # pooled model: one calibration curve per basin
        out = np.empty(len(df))
        for b, curve in iso["per_basin"].items():
            k = (df.basin == b).values
            out[k] = np.interp(p[k], curve["x"], curve["y"])
        return out
    return np.interp(p, iso["x"], iso["y"])