#!/usr/bin/env python3
"""Compare the Python forecast with ping-collector's on the live database.

Run before setting FORECAST_IN_COLLECTOR=true, and after any change to src/ml
or rust/forecast. On the server, as the service user from the app dir:

    sudo -u water-monitor .venv/bin/python scripts/forecast_parity.py

Both sides get the same moment and the same Open-Meteo rain series (fetched
once here and handed to the collector as a file), so a difference is a porting
bug, not timing. Exit 0 when every row matches: station, horizon, as_of,
source, model and every level exactly; probabilities within a few units in the
last place (numpy's and libm's exp may differ in the last bit).
"""

import argparse
import datetime
import json
import math
import os
import subprocess
import sys
import tempfile
from pathlib import Path

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))

ULPS = 4
EXACT = (
    "station_code",
    "horizon_hours",
    "as_of",
    "source",
    "model_version",
    "trained_at",
    "current_level",
    "predicted_max_level",
    "threshold_warning",
    "threshold_danger",
)


def ulps(a: float, b: float) -> float:
    if a == b:
        return 0
    return abs(a - b) / math.ulp(max(abs(a), abs(b)))


def compare(py: list, rs: list) -> list:
    problems = []
    if len(py) != len(rs):
        problems.append(f"row count: python {len(py)}, collector {len(rs)}")
    for p, r in zip(py, rs):
        where = f"{p.get('station_code')} {p.get('horizon_hours')}h"
        for k in EXACT:
            if p.get(k) != r.get(k):
                problems.append(
                    f"{where} {k}: python {p.get(k)!r}, collector {r.get(k)!r}"
                )
        for k in ("p_warning", "p_danger"):
            d = ulps(float(p[k]), float(r[k]))
            if d > ULPS:
                problems.append(
                    f"{where} {k}: python {p[k]!r}, collector {r[k]!r} ({d:.0f} ulp)"
                )
        ps, rs_ = p.get("stages"), r.get("stages")
        if (ps is None) != (rs_ is None) or (ps and len(ps) != len(rs_)):
            problems.append(f"{where} stages differ")
        elif ps:
            for a, b in zip(ps, rs_):
                if a["stage"] != b["stage"] or a["level"] != b["level"]:
                    problems.append(f"{where} stage {a['stage']} differs")
                elif ulps(a["p_exceed"], b["p_exceed"]) > ULPS:
                    problems.append(
                        f"{where} stage {a['stage']} p_exceed {a['p_exceed']} vs {b['p_exceed']}"
                    )
    return problems


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--collector", default=str(REPO / "bin" / "ping-collector"))
    parser.add_argument("--models-dir", default=str(REPO / "models"))
    parser.add_argument(
        "--db-url", default=None, help="default: the app's configuration"
    )
    parser.add_argument("--now", default=None, help="local time, default the clock")
    args = parser.parse_args()

    from src.ml import predict
    from src.ml import rain as rain_mod
    from src.ml.export import json_path_for

    models = Path(args.models_dir)
    stale = [
        p.name
        for p in models.glob("flood_*.joblib")
        if not json_path_for(p).exists()
        or json_path_for(p).stat().st_mtime < p.stat().st_mtime
    ]
    if stale:
        print(
            f"not exported for the collector: {', '.join(stale)}; run `python -m src.ml.export`"
        )
        return 2

    now = (
        datetime.datetime.fromisoformat(args.now)
        if args.now
        else datetime.datetime.now()
    )
    rain = rain_mod.serving_series()
    if rain is None:
        import pandas as pd

        rain = pd.Series(dtype=float)
    py = predict.get_latest_forecasts(
        db_url=args.db_url, models_dir=models, now=now, rain=rain
    )
    py = json.loads(json.dumps(py, default=float))

    with tempfile.TemporaryDirectory() as tmp:
        rain_file = Path(tmp) / "rain.json"
        rain_file.write_text(
            json.dumps(
                {
                    "timestamps": [t.strftime("%Y-%m-%dT%H:%M:%S") for t in rain.index],
                    "values": [None if math.isnan(v) else float(v) for v in rain],
                }
            )
        )
        env = dict(os.environ, MODELS_DIR=str(models), RUST_LOG="error")
        if args.db_url:
            env["POSTGRES_CONNECTION_STRING"] = args.db_url
        out = subprocess.run(
            [
                args.collector,
                "forecast",
                "--json",
                "--now",
                now.isoformat(),
                "--rain-file",
                str(rain_file),
            ],
            capture_output=True,
            text=True,
            env=env,
            cwd=REPO,
            check=False,
        )
    if out.returncode != 0:
        print(f"ping-collector forecast failed:\n{out.stderr}")
        return 2
    rs = json.loads(out.stdout)

    problems = compare(py, rs)
    models_used = sorted({r["model_version"] for r in py})
    n_model = sum(r["source"] == "model" for r in py)
    print(
        f"as of {now:%Y-%m-%d %H:%M:%S}: {len(py)} rows ({n_model} from models {models_used}), "
        f"rain {int(rain.notna().sum())} h"
    )
    if problems:
        print(f"{len(problems)} DIFFERENCES:")
        for p in problems[:40]:
            print("  " + p)
        return 1
    exact = sum(p[k] == r[k] for p, r in zip(py, rs) for k in ("p_warning", "p_danger"))
    print(
        f"identical: every level and field; probabilities {exact}/{2 * len(py)} bit-identical, "
        f"the rest within {ULPS} ulp"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())