"""Model/config hand-over to the Rust forecaster (src/ml/export.py)."""

import gzip
import json

import numpy as np
import pytest
from scipy.special import expit
from sklearn.ensemble import (
    HistGradientBoostingClassifier,
    HistGradientBoostingRegressor,
)

from src.ml import export


def test_rust_flood_config_is_current():
    """rust/forecast/flood_config.json is compiled into ping-collector.

    If this fails, a threshold/stage/lead/notify constant changed: run
    `uv run python -m src.ml.export --config` and rebuild the collector.
    """
    assert export.CONFIG_PATH.read_text(encoding="utf-8") == export.flood_config_text()


def _data(n=600, seed=7):
    rng = np.random.default_rng(seed)
    X = rng.normal(size=(n, 4))
    X[rng.random(size=X.shape) < 0.1] = np.nan  # NaN routing is exported too
    y = np.nan_to_num(X[:, 0]) * 2 + np.nan_to_num(X[:, 1]) + rng.normal(size=n) * 0.1
    return X, y


def _walk(head, X):
    out = np.full(len(X), head["baseline"])
    for tree in head["trees"]:
        thr = [np.inf if t is None else t for t in tree["threshold"]]
        for i, row in enumerate(X):
            node = 0
            while not tree["leaf"][node]:
                v = row[tree["feature"][node]]
                left = tree["missing_left"][node] if np.isnan(v) else v <= thr[node]
                node = tree["left"][node] if left else tree["right"][node]
            out[i] += tree["value"][node]
    return out


def test_exported_trees_reproduce_sklearn_exactly():
    X, y = _data()
    reg = HistGradientBoostingRegressor(max_iter=30, random_state=0).fit(X, y)
    clf = HistGradientBoostingClassifier(max_iter=30, random_state=0).fit(X, y > 0.5)
    # through JSON, as the Rust side reads it
    reg_head = json.loads(json.dumps(export.export_head(reg)))
    clf_head = json.loads(json.dumps(export.export_head(clf)))
    assert reg_head["kind"] == "regressor" and clf_head["kind"] == "classifier"
    assert np.array_equal(_walk(reg_head, X), reg.predict(X))
    assert np.array_equal(expit(_walk(clf_head, X)), clf.predict_proba(X)[:, 1])


def test_bundle_json_roundtrip(tmp_path):
    X, y = _data()
    reg = HistGradientBoostingRegressor(max_iter=5, random_state=0).fit(X, y)
    bundle = {
        "station_code": "P.1",
        "model_version": "hgb-v3+test",
        "trained_at": "2026-10-01T03:30:00",
        "regression_target": "rise",
        "feature_names": ["a", "b", "c", "d"],
        "horizons": [6, 12, 24],
        "sigma": {6: 0.15, 12: 0.15, 24: 0.2},
        "thresholds": {"warning": 3.7, "danger": 4.2},
        "heads": {"max_24": reg, "warn_24": None},
        "skipped_heads": {"warn_24": "too few positives"},
    }
    path = tmp_path / "flood_P.1.json.gz"
    export.write_bundle_json(bundle, path)
    with gzip.open(path, "rt", encoding="utf-8") as fh:
        doc = json.load(fh)
    assert doc["format_version"] == export.FORMAT_VERSION
    assert doc["sigma"] == {"6": 0.15, "12": 0.15, "24": 0.2}
    assert doc["heads"]["warn_24"] is None
    assert doc["heads"]["max_24"]["kind"] == "regressor"
    assert not (tmp_path / "flood_P.1.json.gz.tmp").exists()
    first = path.read_bytes()
    export.write_bundle_json(bundle, path)
    assert path.read_bytes() == first  # reproducible bytes (gzip mtime 0)


def test_export_models_dir_skips_up_to_date(tmp_path):
    joblib = pytest.importorskip("joblib")
    X, y = _data()
    reg = HistGradientBoostingRegressor(max_iter=3, random_state=0).fit(X, y)
    bundle = {
        "station_code": "P.20",
        "model_version": "hgb-v3+test",
        "feature_names": ["a", "b", "c", "d"],
        "horizons": [6],
        "sigma": {6: 0.15},
        "thresholds": {"warning": 2.35, "danger": 2.8},
        "heads": {"max_6": reg},
    }
    joblib.dump(bundle, tmp_path / "flood_P.20.joblib")
    assert export.export_models_dir(tmp_path) == 1
    assert (tmp_path / "flood_P.20.json.gz").exists()
    assert export.export_models_dir(tmp_path) == 0
    assert export.export_models_dir(tmp_path, force=True) == 1