#!/usr/bin/env python3
"""
Algo recalibration analysis runner.

Connects to the prod (or any) Postgres, joins contacts to the nearest HRRR
grid profile within ±0.07° / ±1h (matching `Weather.find_nearest_hrrr/3`),
and produces a Markdown report covering:

  * row counts and date ranges for the key tables
  * per-band contact distribution including bands that BandConfig hasn't seen
  * monthly sounding ducting refresh
  * native HRRR profile duct stats
  * per-band Pearson correlations of contact distance vs HRRR fields
  * binned distance distributions for HPBL and pressure at 10 GHz

Output is saved to docs/algo-reports/YYYY-MM-DD-recalibration.md so it can
be diffed against algo.md or committed alongside it.

Connection string is read from PROP_PROD_DB_URL (set in .envrc) or the
positional --dsn argument. Dependency-light: psycopg + pandas only.

Usage:
    direnv allow            # picks up PROP_PROD_DB_URL from .envrc
    python3 scripts/recalibrate_algo.py
    # or override:
    python3 scripts/recalibrate_algo.py --dsn 'postgres://user:pw@host/db' --out path.md
"""

from __future__ import annotations

import argparse
import datetime as dt
import os
import sys
import textwrap
from pathlib import Path

try:
    import psycopg
    from psycopg.rows import dict_row
except ImportError:
    sys.exit(
        "psycopg is required: pip install 'psycopg[binary]>=3.1' pandas"
    )

try:
    import pandas as pd
except ImportError:
    sys.exit("pandas is required: pip install pandas")


REPORT_HEADER = """\
# Recalibration Analysis Report — {today}

> Auto-generated by `scripts/recalibrate_algo.py`. This report is the
> input for `algo.md` updates — diff against the prior dated section to
> see what's actually moved.

Connection: `{dsn_redacted}`
Statement timeout: {stmt_timeout}
"""


# ─── SQL fragments ──────────────────────────────────────────────────────────

ROW_COUNTS_SQL = """
SELECT 'contacts' tbl, count(*)::bigint n,
       min(qso_timestamp)::date lo, max(qso_timestamp)::date hi FROM contacts
UNION ALL SELECT 'surface_observations', count(*),
       min(observed_at)::date, max(observed_at)::date FROM surface_observations
UNION ALL SELECT 'soundings', count(*),
       min(observed_at)::date, max(observed_at)::date FROM soundings
UNION ALL SELECT 'hrrr_profiles', count(*),
       min(valid_time)::date, max(valid_time)::date FROM hrrr_profiles
UNION ALL SELECT 'hrrr_native_profiles', count(*),
       min(valid_time)::date, max(valid_time)::date FROM hrrr_native_profiles
UNION ALL SELECT 'era5_profiles', count(*),
       min(valid_time)::date, max(valid_time)::date FROM era5_profiles
UNION ALL SELECT 'iemre_observations', count(*), NULL::date, NULL::date
       FROM iemre_observations
UNION ALL SELECT 'nexrad_observations', count(*),
       min(observed_at)::date, max(observed_at)::date FROM nexrad_observations
UNION ALL SELECT 'rtma_observations', count(*),
       min(valid_time)::date, max(valid_time)::date FROM rtma_observations
UNION ALL SELECT 'terrain_profiles', count(*), NULL::date, NULL::date
       FROM terrain_profiles
UNION ALL SELECT 'propagation_scores', count(*),
       min(valid_time)::date, max(valid_time)::date FROM propagation_scores
ORDER BY tbl;
"""

CONTACTS_BY_BAND_SQL = """
SELECT band::int AS band_mhz, count(*) AS contacts,
       ROUND(AVG(distance_km::numeric)) AS avg_km,
       MAX(distance_km::numeric) AS max_km
FROM contacts
WHERE pos1 IS NOT NULL AND distance_km IS NOT NULL
  AND distance_km < 3000 AND flagged_invalid = false AND band >= 902
GROUP BY 1 ORDER BY 1;
"""

SOUNDING_MONTHLY_SQL = """
SELECT EXTRACT(MONTH FROM observed_at)::int AS month,
       count(*) AS soundings,
       ROUND(100.0 * count(*) FILTER (WHERE ducting_detected) / count(*), 1) AS ducting_pct,
       ROUND(AVG(min_refractivity_gradient)::numeric, 1) AS avg_min_grad,
       ROUND(AVG(precipitable_water_mm)::numeric, 1) AS avg_pwat
FROM soundings
WHERE min_refractivity_gradient IS NOT NULL
GROUP BY 1 ORDER BY 1;
"""

NATIVE_DUCT_SQL = """
SELECT
  CASE WHEN best_duct_band_ghz IS NULL THEN 'none'
       WHEN best_duct_band_ghz < 5 THEN '<5 GHz'
       WHEN best_duct_band_ghz < 15 THEN '5-15 GHz'
       WHEN best_duct_band_ghz < 30 THEN '15-30 GHz'
       WHEN best_duct_band_ghz < 75 THEN '30-75 GHz'
       ELSE '75+ GHz' END AS band_bin,
  count(*) AS profiles,
  ROUND(AVG(inversion_top_m)::numeric) AS avg_inv_top_m,
  ROUND(AVG(theta_e_jump_k)::numeric, 1) AS avg_theta_e_jump,
  ROUND(AVG(bulk_richardson)::numeric, 2) AS avg_richardson
FROM hrrr_native_profiles GROUP BY 1 ORDER BY 1;
"""

# Per-band contact ↔ HRRR Pearson correlations.
PER_BAND_JOIN_SQL = """
WITH joined AS (
  SELECT DISTINCT ON (c.id)
    c.id, c.band::int AS band, c.distance_km::float AS dist,
    h.surface_temp_c AS tc, h.surface_dewpoint_c AS dpc,
    h.surface_pressure_mb AS pr, h.pwat_mm AS pwat,
    h.min_refractivity_gradient AS grad,
    h.surface_refractivity AS sref, h.hpbl_m AS hpbl
  FROM contacts c
  JOIN hrrr_profiles h
    ON h.lat BETWEEN (c.pos1->>'lat')::float - 0.07
                 AND (c.pos1->>'lat')::float + 0.07
   AND h.lon BETWEEN (c.pos1->>'lon')::float - 0.07
                 AND (c.pos1->>'lon')::float + 0.07
   AND h.valid_time BETWEEN c.qso_timestamp - INTERVAL '1 hour'
                       AND c.qso_timestamp + INTERVAL '1 hour'
  WHERE c.pos1 IS NOT NULL AND c.distance_km < 3000 AND c.flagged_invalid = false
  ORDER BY c.id, ABS(EXTRACT(EPOCH FROM h.valid_time - c.qso_timestamp))
)
SELECT * FROM joined;
"""

# Empty-table / data-quality smoke checks.
DATA_GAP_SQL = """
SELECT
  (SELECT count(*) FROM era5_profiles) AS era5_rows,
  (SELECT count(*) FROM hrrr_climatology) AS hrrr_climatology_rows,
  (SELECT count(*) FROM rtma_observations) AS rtma_rows,
  (SELECT count(*) FROM metar_5min_observations) AS metar5_rows,
  (SELECT count(*) FROM contacts WHERE hrrr_status = 'complete') AS hrrr_complete_contacts,
  (SELECT count(*) FROM contacts WHERE hrrr_status != 'complete') AS hrrr_pending_contacts;
"""


# ─── Reporting helpers ──────────────────────────────────────────────────────


def fetch_df(conn, sql: str) -> pd.DataFrame:
    with conn.cursor(row_factory=dict_row) as cur:
        cur.execute(sql)
        rows = cur.fetchall()
    return pd.DataFrame(rows)


def md_table(df: pd.DataFrame) -> str:
    if df.empty:
        return "_(no rows)_\n"
    return df.to_markdown(index=False, floatfmt=".3f") + "\n"


def correlations_per_band(df: pd.DataFrame, min_samples: int = 30) -> pd.DataFrame:
    fields = ["pr", "dpc", "pwat", "sref", "grad", "tc", "hpbl"]
    rows = []
    for band, group in df.groupby("band"):
        if len(group) < min_samples:
            continue
        row = {"band_mhz": int(band), "n": len(group)}
        for f in fields:
            try:
                row[f"rho_{f}"] = round(group["dist"].corr(group[f]), 3)
            except Exception:
                row[f"rho_{f}"] = None
        rows.append(row)
    return pd.DataFrame(rows).sort_values("band_mhz")


def hpbl_bins(df: pd.DataFrame, band_mhz: int = 10000) -> pd.DataFrame:
    sub = df[(df["band"] == band_mhz) & df["hpbl"].notna()].copy()
    if sub.empty:
        return pd.DataFrame()
    bins = [0, 200, 500, 1000, 1500, 2000, 1e9]
    labels = ["<200", "200-500", "500-1000", "1000-1500", "1500-2000", ">=2000"]
    sub["bin"] = pd.cut(sub["hpbl"], bins=bins, labels=labels, right=False)
    out = (
        sub.groupby("bin", observed=True)["dist"]
        .agg(["count", "mean", "median"])
        .round(1)
        .rename(columns={"count": "n", "mean": "avg_km", "median": "p50_km"})
        .reset_index()
    )
    return out


def pressure_bins(df: pd.DataFrame, band_mhz: int = 10000) -> pd.DataFrame:
    sub = df[(df["band"] == band_mhz) & df["pr"].notna()].copy()
    if sub.empty:
        return pd.DataFrame()
    bins = [0, 990, 1000, 1010, 1020, 9999]
    labels = ["<990", "990-1000", "1000-1010", "1010-1020", ">=1020"]
    sub["bin"] = pd.cut(sub["pr"], bins=bins, labels=labels, right=False)
    out = (
        sub.groupby("bin", observed=True)["dist"]
        .agg(["count", "mean", "median"])
        .round(1)
        .rename(columns={"count": "n", "mean": "avg_km", "median": "p50_km"})
        .reset_index()
    )
    return out


def redact(dsn: str) -> str:
    if "@" not in dsn:
        return dsn
    head, tail = dsn.split("@", 1)
    if "://" in head and ":" in head.split("://", 1)[1]:
        scheme, rest = head.split("://", 1)
        user = rest.split(":", 1)[0]
        return f"{scheme}://{user}:***@{tail}"
    return dsn


# ─── Main ───────────────────────────────────────────────────────────────────


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dsn",
        default=os.environ.get("PROP_PROD_DB_URL"),
        help="Postgres connection string (defaults to $PROP_PROD_DB_URL)",
    )
    parser.add_argument(
        "--out",
        default=None,
        help="Output Markdown path (defaults to docs/algo-reports/YYYY-MM-DD-recalibration.md)",
    )
    parser.add_argument(
        "--statement-timeout",
        default="20min",
        help="Postgres statement_timeout (default: 20min)",
    )
    args = parser.parse_args()

    if not args.dsn:
        print(
            "error: PROP_PROD_DB_URL not set and --dsn not given; populate "
            "via .envrc",
            file=sys.stderr,
        )
        return 2

    today = dt.date.today().isoformat()
    out_path = Path(
        args.out
        or Path(__file__).resolve().parent.parent
        / "docs"
        / "algo-reports"
        / f"{today}-recalibration.md"
    )
    out_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"connecting to {redact(args.dsn)}", file=sys.stderr)

    sections: list[str] = []
    sections.append(
        REPORT_HEADER.format(
            today=today,
            dsn_redacted=redact(args.dsn),
            stmt_timeout=args.statement_timeout,
        )
    )

    with psycopg.connect(args.dsn, autocommit=True) as conn:
        with conn.cursor() as cur:
            cur.execute(f"SET statement_timeout = '{args.statement_timeout}';")

        print("• row counts", file=sys.stderr)
        sections.append("\n## Row counts and date ranges\n")
        sections.append(md_table(fetch_df(conn, ROW_COUNTS_SQL)))

        print("• data gaps", file=sys.stderr)
        gaps = fetch_df(conn, DATA_GAP_SQL).iloc[0].to_dict()
        sections.append("\n## Data gaps (these should NOT be zero in a healthy prod)\n")
        sections.append(
            "\n".join(f"- **{k}**: {v}" for k, v in gaps.items()) + "\n"
        )

        print("• contacts by band", file=sys.stderr)
        sections.append("\n## Contacts by band (≥902 MHz)\n")
        sections.append(md_table(fetch_df(conn, CONTACTS_BY_BAND_SQL)))

        print("• sounding monthly", file=sys.stderr)
        sections.append("\n## Monthly sounding ducting\n")
        sections.append(md_table(fetch_df(conn, SOUNDING_MONTHLY_SQL)))

        print("• native HRRR ducts", file=sys.stderr)
        sections.append(
            "\n## HRRR native-profile duct distribution (best supportable band)\n"
        )
        sections.append(md_table(fetch_df(conn, NATIVE_DUCT_SQL)))

        print("• per-band correlations (this is the slow one)", file=sys.stderr)
        joined = fetch_df(conn, PER_BAND_JOIN_SQL)

    if joined.empty:
        sections.append("\n_No contacts ↔ HRRR matches were found._\n")
    else:
        sections.append(
            f"\n## Contact ↔ HRRR per-band Pearson correlations "
            f"(matched n={len(joined)})\n"
        )
        sections.append(md_table(correlations_per_band(joined)))

        sections.append("\n## 10 GHz: HPBL bin distance distribution\n")
        sections.append(md_table(hpbl_bins(joined)))

        sections.append("\n## 10 GHz: pressure bin distance distribution\n")
        sections.append(md_table(pressure_bins(joined)))

    out_path.write_text("\n".join(sections))
    print(f"wrote {out_path}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())