#!/usr/bin/env python3
"""
Full algorithm recalibration — unified Python pipeline.

Recomputes per-band weight overrides for the 10-factor scorer from real
propagation evidence in the prod DB:

  * VHF / UHF (50, 144, 432 MHz): Pearson correlations between PSKR
    `spot_count` density and HRRR features in `pskr_calibration_samples`
    (~520k samples since 2026-05-04 in prod).
  * Microwave (>=902 MHz): Pearson correlations between QSO `distance_km`
    and HRRR features in `contacts` joined to `hrrr_profiles`.

Weight derivation (per band):
  1. Start from the globally-fit default weight vector (`DEFAULT_WEIGHTS`).
  2. Per correlation-backed factor, build a "relative signal strength"
     multiplier = (|band_r| + EPS) / (|ref_r| + EPS), clamped to
     [CLAMP_LO, CLAMP_HI]. The reference is the 10 GHz contacts fit
     because it has the densest evidence in the corpus.
  3. Factors without per-band correlations (rain, season, time_of_day,
     sky, wind) get physics-informed multipliers — same priors used by
     the legacy `derive_band_weights.py`.
  4. Multiply defaults by multipliers, normalize to sum = 1.0.

Outputs (no Elixir source ever modified):
  * priv/algo/band_weights.json — machine-readable per-band overrides
    that BandConfig can load at runtime via Jason.
  * docs/algo-reports/{today}-full-recalibration.md — human report
    with per-band correlations, sample counts, derived weights, and
    a "what moved" diff against the previous JSON file (if present).

Usage:
    direnv allow                              # picks up PROP_PROD_DB_URL
    python3 scripts/recalibrate.py            # writes JSON + report
    python3 scripts/recalibrate.py --dry-run  # prints JSON to stdout

The script is self-contained: psycopg + stdlib only (no pandas, no
numpy). Run it whenever the contact or PSKR corpus grows materially.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import sys
from dataclasses import dataclass, field
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'")


REPO_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_JSON_OUT = REPO_ROOT / "priv" / "algo" / "band_weights.json"
DEFAULT_REPORT_DIR = REPO_ROOT / "docs" / "algo-reports"


# ── Constants ────────────────────────────────────────────────────────────────

# Must match `@weights` in lib/microwaveprop/propagation/band_config.ex.
# These are the globally-fit gradient-descent priors (2026-04-11) that
# every per-band fit modulates.
DEFAULT_WEIGHTS: dict[str, float] = {
    "humidity":      0.1262,
    "time_of_day":   0.0380,
    "td_depression": 0.1010,
    "refractivity":  0.0986,
    "sky":           0.0841,
    "season":        0.1134,
    "wind":          0.0841,
    "rain":          0.1431,
    "pwat":          0.1147,
    "pressure":      0.0967,
}

# Correlation-fed factor → HRRR field. These factors get data-driven
# multipliers per band. The rest get physics priors.
FACTOR_SOURCE: dict[str, str] = {
    "humidity":      "dpc",   # dewpoint carries the moisture signal
    "td_depression": "tc",    # temperature dominates T-Td variance
    "refractivity":  "grad",  # refractivity gradient
    "pressure":      "pr",
    "pwat":          "pwat",
}

# ITU-R P.838-3 rain_k (per BandConfig). Used to scale the `rain` weight
# by a sqrt-dampened ratio against 10 GHz (the score itself already
# scales by rain_k internally).
BAND_RAIN_K: dict[int, float] = {
    50: 0.0, 144: 0.0, 222: 0.0, 432: 0.0,
    902: 0.0, 1_296: 0.0,
    2_304: 0.001, 3_400: 0.002, 5_760: 0.005,
    10_000: 0.010, 24_000: 0.070, 47_000: 0.187,
    68_000: 0.310, 75_000: 0.345,
    122_000: 0.498, 134_000: 0.520, 142_000: 0.530, 145_000: 0.535,
    241_000: 0.550, 288_000: 0.560, 322_000: 0.570,
    403_000: 0.580, 411_000: 0.580,
}

# Season weight multiplier (physics prior — see algo.md and the legacy
# derive_band_weights.py). VHF has large seasonal swings (Es / tropo
# mixing); 24+ GHz peaks in winter due to dry-air absorption.
BAND_SEASON_MULT: dict[int, float] = {
    50: 1.6, 144: 1.4, 222: 1.3, 432: 1.2,
    902: 1.1, 1_296: 1.0,
    2_304: 1.0, 3_400: 1.0, 5_760: 1.0,
    10_000: 1.0, 24_000: 1.1, 47_000: 1.2,
    68_000: 1.3, 75_000: 1.4,
    122_000: 1.5, 134_000: 1.5, 142_000: 1.5, 145_000: 1.5,
    241_000: 1.6, 288_000: 1.6, 322_000: 1.7, 403_000: 1.7, 411_000: 1.7,
}

# Time-of-day multiplier — Finding 8 in algo.md (effect scales with
# frequency). 10 GHz is the baseline at 1.0.
BAND_TOD_MULT: dict[int, float] = {
    50: 0.6, 144: 0.7, 222: 0.8, 432: 0.8,
    902: 0.9, 1_296: 1.0,
    2_304: 1.0, 3_400: 1.0, 5_760: 1.0,
    10_000: 1.0, 24_000: 1.8, 47_000: 2.5,
    68_000: 3.5, 75_000: 4.0,
    122_000: 5.0, 134_000: 5.0, 142_000: 5.0, 145_000: 5.0,
    241_000: 5.0, 288_000: 5.0, 322_000: 5.0, 403_000: 5.0, 411_000: 5.0,
}

# PSKR `band` strings → MHz (PSKR scraper stores ham band labels).
PSKR_BAND_TO_MHZ: dict[str, int] = {
    "6m":   50,
    "2m":   144,
    "70cm": 432,
    "23cm": 1_296,
    "13cm": 2_304,
    "3cm":  10_000,
}

# Per-band data source routing. PSKR for VHF/UHF where spot density is
# rich and contacts are sparse; contacts for microwave where PSKR
# activity is rare and the contact corpus is dense.
PSKR_BANDS = {50, 144, 432}

# Floor on |r| (noise level we can distinguish from zero given typical n).
# Same as legacy derive_band_weights.py for continuity.
EPS = 0.05
CLAMP_LO = 0.5
CLAMP_HI = 2.0
MIN_N_FOR_FIT = 200

# Factor display order in JSON + Markdown. Matches band_config.ex
# comment ordering for diffability.
FACTOR_ORDER = [
    "humidity", "time_of_day", "td_depression", "refractivity",
    "sky", "season", "wind", "rain", "pwat", "pressure",
]


# ── Data structures ─────────────────────────────────────────────────────────


@dataclass
class BandCorr:
    """Per-band correlation snapshot, source-agnostic."""
    band_mhz: int
    source: str  # "pskr" or "contacts"
    n: int
    rho: dict[str, float] = field(default_factory=dict)
    # Optional per-band diagnostics for the report.
    total_label: float | None = None
    avg_label:   float | None = None


# ── SQL ─────────────────────────────────────────────────────────────────────


# Contacts ↔ HRRR (microwave): per-band Pearson(distance_km, feature).
# Same nearest-neighbour join as the legacy recalibrate_algo.py.
CONTACTS_CORR_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
  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 band,
       count(*)::int AS n,
       CORR(dist, tc)   AS rho_tc,
       CORR(dist, dpc)  AS rho_dpc,
       CORR(dist, pr)   AS rho_pr,
       CORR(dist, pwat) AS rho_pwat,
       CORR(dist, grad) AS rho_grad,
       avg(dist)::float AS avg_dist
FROM joined
WHERE band >= 50
GROUP BY band
HAVING count(*) >= 50
ORDER BY band;
"""

# PSKR samples already carry HRRR fields (the PSKR sampler joins them at
# ingestion). Label is `spot_count` per (hour, midpoint, band) cell.
# Filter to HRRR-complete rows so all five correlations are computed
# from the same denominator.
PSKR_CORR_SQL = """
SELECT band,
       count(*)::int AS n,
       CORR(spot_count, surface_temp_c)            AS rho_tc,
       CORR(spot_count, surface_dewpoint_c)        AS rho_dpc,
       CORR(spot_count, surface_pressure_mb)       AS rho_pr,
       CORR(spot_count, pwat_mm)                   AS rho_pwat,
       CORR(spot_count, min_refractivity_gradient) AS rho_grad,
       sum(spot_count)::bigint AS total_spots,
       avg(spot_count)::float  AS avg_spots
FROM pskr_calibration_samples
WHERE surface_pressure_mb IS NOT NULL
  AND min_refractivity_gradient IS NOT NULL
GROUP BY band
ORDER BY band;
"""


# ── Loaders ─────────────────────────────────────────────────────────────────


def _row_to_rho(row: dict) -> dict[str, float]:
    return {
        "tc":   row["rho_tc"]   or 0.0,
        "dpc":  row["rho_dpc"]  or 0.0,
        "pr":   row["rho_pr"]   or 0.0,
        "pwat": row["rho_pwat"] or 0.0,
        "grad": row["rho_grad"] or 0.0,
    }


def load_contacts_correlations(conn) -> dict[int, BandCorr]:
    print("• loading contacts ↔ HRRR correlations", file=sys.stderr)
    with conn.cursor(row_factory=dict_row) as cur:
        cur.execute(CONTACTS_CORR_SQL)
        rows = cur.fetchall()
    out: dict[int, BandCorr] = {}
    for r in rows:
        band = int(r["band"])
        out[band] = BandCorr(
            band_mhz=band,
            source="contacts",
            n=int(r["n"]),
            rho=_row_to_rho(r),
            avg_label=r.get("avg_dist"),
        )
    return out


def load_pskr_correlations(conn) -> dict[int, BandCorr]:
    print("• loading PSKR ↔ HRRR correlations", file=sys.stderr)
    with conn.cursor(row_factory=dict_row) as cur:
        cur.execute(PSKR_CORR_SQL)
        rows = cur.fetchall()
    out: dict[int, BandCorr] = {}
    for r in rows:
        band_label = r["band"]
        mhz = PSKR_BAND_TO_MHZ.get(band_label)
        if mhz is None:
            # Unknown PSKR band string — log and skip rather than guess.
            print(f"  (skipping unknown PSKR band {band_label!r})", file=sys.stderr)
            continue
        out[mhz] = BandCorr(
            band_mhz=mhz,
            source="pskr",
            n=int(r["n"]),
            rho=_row_to_rho(r),
            total_label=float(r["total_spots"]) if r["total_spots"] is not None else None,
            avg_label=r.get("avg_spots"),
        )
    return out


def merge_correlations(
    pskr: dict[int, BandCorr],
    contacts: dict[int, BandCorr],
) -> dict[int, BandCorr]:
    """Per-band routing: PSKR for VHF/UHF, contacts for everything else."""
    out: dict[int, BandCorr] = {}
    for band in set(pskr) | set(contacts):
        if band in PSKR_BANDS and band in pskr:
            out[band] = pskr[band]
        elif band in contacts:
            out[band] = contacts[band]
        elif band in pskr:
            # No contacts but PSKR has it (e.g. 23cm/13cm if we expand routing).
            out[band] = pskr[band]
    return out


# ── Weight derivation ──────────────────────────────────────────────────────


def derive_weights(
    band_mhz: int,
    corrs: dict[int, BandCorr],
) -> tuple[dict[str, float] | None, dict[str, float]]:
    """Returns (weights, multipliers). weights=None means inherit defaults.

    `multipliers` is always populated so the report can show why a band
    moved (or why it didn't qualify for a fit).
    """
    band_corr = corrs.get(band_mhz)
    # Reference is always the 10 GHz contacts fit — densest, most trusted.
    ref_corr = corrs.get(10_000)

    multipliers: dict[str, float] = {}

    # Correlation-fed factors first (skipped if either side missing).
    if band_corr is not None and ref_corr is not None:
        for factor, field_name in FACTOR_SOURCE.items():
            band_r = abs(band_corr.rho[field_name]) + EPS
            ref_r = abs(ref_corr.rho[field_name]) + EPS
            ratio = band_r / ref_r
            multipliers[factor] = max(CLAMP_LO, min(CLAMP_HI, ratio))
    else:
        for factor in FACTOR_SOURCE:
            multipliers[factor] = 1.0

    # Physics priors for the rest.
    rain_k = BAND_RAIN_K.get(band_mhz, 0.010)
    rain_k_ref = BAND_RAIN_K[10_000]
    if rain_k > 0:
        multipliers["rain"] = max(0.2, min(3.0, (rain_k / rain_k_ref) ** 0.5))
    else:
        multipliers["rain"] = 0.1

    multipliers["season"] = BAND_SEASON_MULT.get(band_mhz, 1.0)
    multipliers["time_of_day"] = BAND_TOD_MULT.get(band_mhz, 1.0)
    multipliers["sky"] = 1.0
    multipliers["wind"] = 1.0

    # No fit if we don't have enough data on this band.
    if band_corr is None or band_corr.n < MIN_N_FOR_FIT:
        return None, multipliers

    raw = {k: DEFAULT_WEIGHTS[k] * multipliers[k] for k in DEFAULT_WEIGHTS}
    total = sum(raw.values())
    weights = {k: round(v / total, 4) for k, v in raw.items()}
    return weights, multipliers


# ── Output ─────────────────────────────────────────────────────────────────


def build_json(
    corrs: dict[int, BandCorr],
    pskr_only: dict[int, BandCorr],
    contacts_only: dict[int, BandCorr],
    generated_at: dt.datetime,
) -> dict:
    """Stable JSON shape — keys sorted, numbers rounded for diffability."""
    all_bands = sorted(set(BAND_RAIN_K.keys()) | set(corrs.keys()))
    overrides: dict[str, dict] = {}
    for band in all_bands:
        weights, _mults = derive_weights(band, corrs)
        if weights is None:
            continue
        bc = corrs.get(band)
        overrides[str(band)] = {
            "weights": {k: weights[k] for k in FACTOR_ORDER},
            "n_samples": bc.n if bc else 0,
            "source": bc.source if bc else "default",
        }

    return {
        "generated_at": generated_at.isoformat(timespec="seconds"),
        "generated_by": "scripts/recalibrate.py",
        "schema_version": 1,
        "global_weights": {k: DEFAULT_WEIGHTS[k] for k in FACTOR_ORDER},
        "data_sources": {
            "pskr": {
                "bands":         sorted(pskr_only.keys()),
                "total_samples": sum(bc.n for bc in pskr_only.values()),
            },
            "contacts": {
                "bands":         sorted(contacts_only.keys()),
                "total_samples": sum(bc.n for bc in contacts_only.values()),
            },
        },
        "band_overrides": overrides,
    }


def load_prior_json(path: Path) -> dict | None:
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text())
    except (OSError, json.JSONDecodeError) as exc:
        print(f"  (could not read prior JSON at {path}: {exc})", file=sys.stderr)
        return None


def diff_weights(
    prior: dict | None, current: dict,
) -> list[tuple[str, str, float, float, float]]:
    """Returns rows (band, factor, prior_w, new_w, delta) sorted by |delta|."""
    if not prior:
        return []
    prior_overrides = prior.get("band_overrides", {})
    rows: list[tuple[str, str, float, float, float]] = []
    for band, entry in current.get("band_overrides", {}).items():
        prior_entry = prior_overrides.get(band, {}).get("weights", {})
        new_entry = entry["weights"]
        for factor in FACTOR_ORDER:
            old = float(prior_entry.get(factor, DEFAULT_WEIGHTS[factor]))
            new = float(new_entry[factor])
            delta = new - old
            if abs(delta) >= 0.005:  # 0.5pp threshold to keep noise out
                rows.append((band, factor, old, new, delta))
    rows.sort(key=lambda r: abs(r[4]), reverse=True)
    return rows


def render_report(
    payload: dict,
    corrs: dict[int, BandCorr],
    diffs: list[tuple[str, str, float, float, float]],
    dsn_redacted: str,
) -> str:
    today = payload["generated_at"][:10]
    lines: list[str] = []
    lines.append(f"# Full Recalibration Report — {today}\n")
    lines.append(
        "> Auto-generated by `scripts/recalibrate.py`. The companion "
        "JSON file at `priv/algo/band_weights.json` is the machine-"
        "readable output that Elixir loads at runtime. This report is "
        "the human-readable audit trail.\n"
    )
    lines.append(f"Connection: `{dsn_redacted}`")
    lines.append(f"Generated: `{payload['generated_at']}`\n")

    # ── Data sources
    lines.append("## Data sources\n")
    src = payload["data_sources"]
    lines.append(
        f"- **PSKR** (VHF/UHF, label = `spot_count`): bands "
        f"{src['pskr']['bands']}, {src['pskr']['total_samples']:,} samples"
    )
    lines.append(
        f"- **Contacts ↔ HRRR** (microwave, label = `distance_km`): "
        f"{len(src['contacts']['bands'])} bands, "
        f"{src['contacts']['total_samples']:,} matched contacts"
    )
    lines.append("")

    # ── Per-band evidence summary
    lines.append("## Per-band evidence\n")
    lines.append("| Band | Source | n | label avg | ρ_tc | ρ_dpc | ρ_pr | ρ_pwat | ρ_grad |")
    lines.append("|------|--------|---:|----------:|-----:|------:|-----:|-------:|-------:|")
    for band in sorted(corrs):
        bc = corrs[band]
        avg = f"{bc.avg_label:.2f}" if bc.avg_label is not None else "—"
        lines.append(
            f"| {band} MHz | {bc.source} | {bc.n:,} | {avg} | "
            f"{bc.rho['tc']:+.3f} | {bc.rho['dpc']:+.3f} | {bc.rho['pr']:+.3f} | "
            f"{bc.rho['pwat']:+.3f} | {bc.rho['grad']:+.3f} |"
        )
    lines.append("")

    # ── Derived overrides
    lines.append("## Derived per-band weights\n")
    overrides = payload["band_overrides"]
    if not overrides:
        lines.append("_No band had ≥{} samples — nothing to override._\n".format(MIN_N_FOR_FIT))
    else:
        header = "| Band | n | source | " + " | ".join(FACTOR_ORDER) + " |"
        sep    = "|------|---:|--------|" + "|".join(["----:"] * len(FACTOR_ORDER)) + "|"
        lines.append(header)
        lines.append(sep)
        for band_str in sorted(overrides, key=int):
            entry = overrides[band_str]
            ws = entry["weights"]
            cells = " | ".join(f"{ws[f]:.4f}" for f in FACTOR_ORDER)
            lines.append(
                f"| {band_str} MHz | {entry['n_samples']:,} | "
                f"{entry['source']} | {cells} |"
            )
    lines.append("")

    # ── Diff against prior
    lines.append("## Changes vs prior `band_weights.json`\n")
    if not diffs:
        lines.append(
            "_No prior file, or no factor moved by ≥0.005. First-run "
            "values can be compared to the in-tree `@band_configs` "
            "weights in `lib/microwaveprop/propagation/band_config.ex`._\n"
        )
    else:
        lines.append("Largest weight moves (≥0.5 pp), sorted by magnitude:\n")
        lines.append("| Band | Factor | Prior | New | Δ |")
        lines.append("|------|--------|------:|----:|--:|")
        for band, factor, old, new, delta in diffs[:30]:
            lines.append(
                f"| {band} MHz | {factor} | {old:.4f} | {new:.4f} | {delta:+.4f} |"
            )
        if len(diffs) > 30:
            lines.append(f"\n_(+{len(diffs) - 30} smaller moves omitted)_")
    lines.append("")

    # ── How to apply
    lines.append("## Applying these weights\n")
    lines.append(
        "This script writes `priv/algo/band_weights.json` and stops "
        "there — it never edits `band_config.ex`. To make the weights "
        "live, `BandConfig.weights(band_mhz)` needs a one-time edit to "
        "read this JSON at app start (cache in `:persistent_term`) and "
        "merge `band_overrides[band].weights` over the in-source "
        "defaults. Until then, the JSON is documentation only."
    )
    return "\n".join(lines)


# ── Plumbing ───────────────────────────────────────────────────────────────


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


def main() -> int:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    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-json",
        type=Path,
        default=DEFAULT_JSON_OUT,
        help=f"JSON output path (default: {DEFAULT_JSON_OUT.relative_to(REPO_ROOT)})",
    )
    parser.add_argument(
        "--out-report",
        type=Path,
        default=None,
        help="Markdown report path (default: docs/algo-reports/{today}-full-recalibration.md)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print JSON to stdout, do not write files",
    )
    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",
            file=sys.stderr,
        )
        return 2

    today = dt.date.today().isoformat()
    out_report = args.out_report or (DEFAULT_REPORT_DIR / f"{today}-full-recalibration.md")

    print(f"connecting to {redact(args.dsn)}", file=sys.stderr)
    with psycopg.connect(args.dsn, autocommit=True) as conn:
        with conn.cursor() as cur:
            cur.execute(f"SET statement_timeout = '{args.statement_timeout}'")

        pskr = load_pskr_correlations(conn)
        contacts = load_contacts_correlations(conn)

    corrs = merge_correlations(pskr, contacts)
    print(
        f"• merged: {len(corrs)} bands "
        f"({sum(1 for b in corrs if corrs[b].source == 'pskr')} PSKR, "
        f"{sum(1 for b in corrs if corrs[b].source == 'contacts')} contacts)",
        file=sys.stderr,
    )

    generated_at = dt.datetime.now(dt.timezone.utc).replace(microsecond=0)
    payload = build_json(corrs, pskr, contacts, generated_at)

    prior = load_prior_json(args.out_json)
    diffs = diff_weights(prior, payload)

    if args.dry_run:
        print(json.dumps(payload, indent=2, sort_keys=False))
        print(f"\n(dry-run: would write {len(payload['band_overrides'])} overrides to {args.out_json})", file=sys.stderr)
        return 0

    args.out_json.parent.mkdir(parents=True, exist_ok=True)
    args.out_json.write_text(json.dumps(payload, indent=2, sort_keys=False) + "\n")
    print(f"• wrote {args.out_json} ({len(payload['band_overrides'])} overrides)", file=sys.stderr)

    out_report.parent.mkdir(parents=True, exist_ok=True)
    out_report.write_text(render_report(payload, corrs, diffs, redact(args.dsn)) + "\n")
    print(f"• wrote {out_report}", file=sys.stderr)

    return 0


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