#!/usr/bin/env python3
"""
Per-band weight derivation from the full-corpus correlation analysis.

Reads Pearson correlations that `recalibrate_algo.py` wrote into the DB
and emits a `@band_weight_overrides` Elixir map ready to drop into
`lib/microwaveprop/propagation/band_config.ex`.

The derivation rule is deliberately simple and legible, because the
correlations are noisy (|r| typically 0.05–0.25) and we'd rather preserve
the globally-fit default weights as a prior than over-fit per-band.

Rule:
  1. Start from the default weight vector (gradient-descent fit, global).
  2. Compute a per-factor "relative signal strength" ratio
     s_band,factor = |r_band,factor| / |r_10GHz,factor|
     with a small ε floor so zero-correlation factors don't collapse.
  3. Clamp the ratio to [0.4, 2.5] — noisy per-band correlations shouldn't
     move a factor more than 2.5× or less than 0.4× the prior.
  4. Factors we don't have direct correlations for (sky, wind, rain,
     season, time_of_day) carry a physics-informed multiplier per band:
       * rain scales with `rain_k` (ITU-R P.838 specific attenuation
         coefficient) normalized to 10 GHz's rain_k.
       * season scales up at VHF/low-UHF where Es-like physics amplifies
         monthly variation (we don't model Es directly).
       * time_of_day scales with frequency (Finding 8 in algo.md).
       * sky/wind stay flat — no per-band evidence either direction.
  5. Multiply default weights by the multipliers, then normalize to 1.0.

The output goes to stdout as Elixir source ready to paste into BandConfig.
"""
from __future__ import annotations

import argparse
import os
import sys
from dataclasses import dataclass

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


# Default weights (must match `@weights` in band_config.ex). These are the
# globally-fit gradient-descent priors — the per-band deltas modulate them.
DEFAULT_WEIGHTS = {
    "rain": 0.1362,
    "humidity": 0.1243,
    "pwat": 0.1128,
    "season": 0.1112,
    "refractivity": 0.1049,
    "pressure": 0.1032,
    "td_depression": 0.0978,
    "sky": 0.08,
    "wind": 0.08,
    "time_of_day": 0.0496,
}

# Correlation-fed factor → source-field mapping. These factors get data-
# driven per-band multipliers; the rest get physics-informed multipliers.
FACTOR_SOURCE = {
    "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 values (per BandConfig). Used for `rain` weight
# scaling — rain attenuation grows by ~4 orders of magnitude from VHF
# to sub-mm, so `rain` weight should track that.
BAND_RAIN_K = {
    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 by band. VHF has large seasonal swings driven
# by tropo-mixing physics (summer peak). Microwave seasonal swings are
# contest-schedule artefacts but the physics (summer convective mixing
# hurting 24+ GHz) still applies. No per-band evidence in the correlation
# report so we use physics priors.
BAND_SEASON_MULT = {
    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 by band (Finding 8 in algo.md — effect scales
# with frequency). 10 GHz is the baseline.
BAND_TOD_MULT = {
    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,
}

# EPS sets an absolute floor on |r| so two near-zero correlations don't
# produce a giant ratio. Set at the noise floor we can reliably distinguish
# from zero given typical n. Tightened along with the clamp so small
# corpora don't drag the weights around.
EPS = 0.05
CLAMP_LO = 0.5
CLAMP_HI = 2.0
MIN_N_FOR_FIT = 200  # bands with fewer matched contacts use physics only


@dataclass
class BandCorr:
    band_mhz: int
    n: int
    rho: dict[str, float]


PER_BAND_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(*) 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
FROM joined
WHERE band >= 50
GROUP BY band
HAVING count(*) >= 50
ORDER BY band;
"""


def load_correlations(dsn: str) -> dict[int, BandCorr]:
    with psycopg.connect(dsn) as conn:
        with conn.cursor(row_factory=dict_row) as cur:
            cur.execute("SET statement_timeout = '30min'")
            cur.execute(PER_BAND_SQL)
            rows = cur.fetchall()
    out = {}
    for r in rows:
        out[int(r["band"])] = BandCorr(
            band_mhz=int(r["band"]),
            n=int(r["n"]),
            rho={
                "tc": r["rho_tc"] or 0.0,
                "dpc": r["rho_dpc"] or 0.0,
                "pr": r["rho_pr"] or 0.0,
                "pwat": r["rho_pwat"] or 0.0,
                "grad": r["rho_grad"] or 0.0,
            },
        )
    return out


def derive_weights(
    band_mhz: int, corrs: dict[int, BandCorr]
) -> dict[str, float] | None:
    """Returns a normalized weight map or None if the band should use defaults."""
    band_corr = corrs.get(band_mhz)
    ref_corr = corrs.get(10_000)

    if band_corr is None or band_corr.n < MIN_N_FOR_FIT or ref_corr is None:
        return None

    multipliers = {}

    # Data-driven multipliers for correlation-backed factors.
    for factor, field in FACTOR_SOURCE.items():
        band_r = abs(band_corr.rho[field]) + EPS
        ref_r = abs(ref_corr.rho[field]) + EPS
        ratio = band_r / ref_r
        multipliers[factor] = max(CLAMP_LO, min(CLAMP_HI, ratio))

    # Physics-driven multipliers for the rest.
    # Rain: weight tracks rain_k but with sqrt-dampening. A 7× rain_k jump
    # (10→24 GHz) should not become a 7× weight jump — the score itself
    # already scales with rain_k via ITU-R P.838 in `score_rain`, so the
    # weight only needs to capture how often rain drives the outcome.
    rain_k = BAND_RAIN_K.get(band_mhz, 0.010)
    rain_k_10g = BAND_RAIN_K[10_000]
    if rain_k > 0:
        ratio = rain_k / rain_k_10g
        rain_mult = max(0.2, min(3.0, ratio**0.5))
    else:
        rain_mult = 0.1
    multipliers["rain"] = rain_mult

    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

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


def format_elixir(band_mhz: int, weights: dict[str, float]) -> str:
    # Preserve the factor ordering used in band_config.ex comments.
    order = [
        "humidity", "time_of_day", "td_depression", "refractivity",
        "sky", "season", "wind", "rain", "pwat", "pressure",
    ]
    lines = [f"      {k}: {weights[k]:.4f}" for k in order]
    inner = ",\n".join(lines)
    return f"    weights: %{{\n{inner}\n    }}"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dsn", default=os.environ.get("PROP_DB_URL"))
    args = parser.parse_args()
    if not args.dsn:
        print("error: PROP_DB_URL not set and --dsn not given", file=sys.stderr)
        return 2

    print(f"loading correlations from {args.dsn}", file=sys.stderr)
    corrs = load_correlations(args.dsn)

    print(
        f"correlations available for {len(corrs)} bands: "
        f"{sorted(corrs.keys())}",
        file=sys.stderr,
    )

    # Every band in BandConfig — we emit an entry for each so the reader
    # sees at a glance which bands got fit and which inherited defaults.
    all_bands = sorted(set(BAND_RAIN_K.keys()) | set(corrs.keys()))

    sections = []
    for band in all_bands:
        weights = derive_weights(band, corrs)
        band_corr = corrs.get(band)
        n = band_corr.n if band_corr else 0
        if weights is None:
            sections.append(
                f"# {band} MHz — n={n}, uses default @weights (insufficient signal)"
            )
        else:
            sections.append(
                f"# {band} MHz — n={n}, derived from corpus correlations\n"
                f"{format_elixir(band, weights)}"
            )
    print("\n\n".join(sections))
    return 0


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