"""Calibration statistics: how the engine's suggestion compares with the
danger forecasters actually published.

These numbers are DESCRIPTIVE. Imported history (e.g. the Sierra archive,
2013-2021) predates the current avalanche-problem guidance and was not produced
with this tool, so disagreement is not necessarily an error on either side.
Imported forecasts carry likelihood and size but no sensitivity/distribution,
so they can calibrate the danger table but not the likelihood matrix.
"""

from __future__ import annotations

import math
from collections import Counter, defaultdict

from psycopg import Connection

from .db import fetchall
from .engine import BANDS, ENGINE_VERSION, bands_of, danger_for, size_class

LEVELS = [1, 2, 3, 4, 5]


def _f(v):
    return None if v is None else float(v)


def suggest_from_recorded(problems: list[dict], config: dict) -> dict:
    """Per-band suggestion for problems that already carry a likelihood (imported forecasts)."""
    out = {"upper": 0, "middle": 0, "lower": 0, "overall": 0}
    for band in BANDS:
        best = 0
        for p in problems:
            if p.get("likelihood") and p.get("size_max") is not None and band in bands_of(p.get("locations")):
                best = max(best, danger_for(int(p["likelihood"]), float(p["size_max"]), config))
        out[band] = best
        out["overall"] = max(out["overall"], best)
    return out


def load_forecasts(conn: Connection, center_id: int, season: str | None = None, era: str | None = None) -> list[dict]:
    return fetchall(
        conn,
        """SELECT f.id, f.valid_date, f.season, f.guidance_era, f.zone_id, f.source,
                  f.danger_upper, f.danger_middle, f.danger_lower,
                  coalesce(json_agg(json_build_object(
                      'type', p.problem_type, 'rank', p.rank, 'locations', p.locations,
                      'likelihood', p.likelihood, 'size_min', p.size_min, 'size_max', p.size_max)
                      ORDER BY p.rank) FILTER (WHERE p.id IS NOT NULL), '[]') AS problems
           FROM forecasts f LEFT JOIN forecast_problems p ON p.forecast_id = f.id
           WHERE f.center_id = %s AND (%s::text IS NULL OR f.season = %s) AND (%s::text IS NULL OR f.guidance_era = %s)
           GROUP BY f.id ORDER BY f.valid_date""",
        (center_id, season, season, era, era),
    )


def load_tool_runs(conn: Connection, center_id: int, season: str | None = None) -> list[dict]:
    rows = fetchall(
        conn,
        """SELECT id, valid_date, final_upper AS danger_upper, final_middle AS danger_middle, final_lower AS danger_lower,
                  suggested_upper, suggested_middle, suggested_lower, final_problems AS problems, config_id
           FROM tool_runs WHERE center_id = %s AND status = 'finalized' ORDER BY valid_date""",
        (center_id,),
    )
    for r in rows:
        d = r["valid_date"]
        start = d.year if d.month >= 8 else d.year - 1
        r["season"] = f"{start}-{str(start + 1)[-2:]}"
    return [r for r in rows if season is None or r["season"] == season]


def band_pairs(items: list[dict], config: dict | None, use_stored_suggestion: bool = False) -> tuple[list[dict], Counter]:
    """(suggested, final) per band. Skips unrated bands and bands no problem touches."""
    pairs, skipped = [], Counter()
    for it in items:
        sugg = (
            {b: it[f"suggested_{b}"] for b in BANDS}
            if use_stored_suggestion
            else suggest_from_recorded(it["problems"] or [], config)
        )
        primary = (it["problems"] or [{}])[0].get("type") if it["problems"] else None
        for b in BANDS:
            final = it[f"danger_{b}"]
            if final is None or final == 0:
                skipped["unrated"] += 1
                continue
            if not sugg[b]:
                skipped["no_problem_in_band"] += 1
                continue
            pairs.append({"id": it["id"], "date": it["valid_date"], "season": it.get("season"), "band": b,
                          "suggested": int(sugg[b]), "final": int(final), "primary": primary})
    return pairs, skipped


def weighted_kappa(pairs: list[dict]) -> float | None:
    """Quadratic-weighted Cohen's kappa over danger levels 1-5."""
    n = len(pairs)
    if n == 0:
        return None
    k = len(LEVELS)
    obs = [[0] * k for _ in range(k)]
    for p in pairs:
        obs[p["suggested"] - 1][p["final"] - 1] += 1
    rows = [sum(r) for r in obs]
    cols = [sum(obs[i][j] for i in range(k)) for j in range(k)]
    num = den = 0.0
    for i in range(k):
        for j in range(k):
            w = (i - j) ** 2 / (k - 1) ** 2
            num += w * obs[i][j]
            den += w * rows[i] * cols[j] / n
    return None if den == 0 else round(1 - num / den, 4)


def summarize(pairs: list[dict]) -> dict:
    n = len(pairs)
    confusion = [[0] * 6 for _ in range(6)]  # [suggested][final], levels 0..5
    for p in pairs:
        confusion[p["suggested"]][p["final"]] += 1
    if n == 0:
        return {"n": 0, "exact": None, "within_one": None, "bias": None, "kappa": None, "confusion": confusion}
    exact = sum(p["suggested"] == p["final"] for p in pairs)
    within = sum(abs(p["suggested"] - p["final"]) <= 1 for p in pairs)
    return {
        "n": n,
        "exact": round(exact / n, 4),
        "within_one": round(within / n, 4),
        "bias": round(sum(p["final"] - p["suggested"] for p in pairs) / n, 4),  # + = forecasters rated higher
        "kappa": weighted_kappa(pairs),
        "confusion": confusion,
    }


def agreement(conn: Connection, center_id: int, config: dict, source: str = "imported", season: str | None = None, era: str | None = None) -> dict:
    if source == "tool":
        items = load_tool_runs(conn, center_id, season)
        pairs, skipped = band_pairs(items, config, use_stored_suggestion=True)
    else:
        items = load_forecasts(conn, center_id, season, era)
        pairs, skipped = band_pairs(items, config)
    by_band = {b: summarize([p for p in pairs if p["band"] == b]) for b in BANDS}
    groups = defaultdict(list)
    seasons = defaultdict(list)
    for p in pairs:
        groups[p["primary"]].append(p)
        seasons[p["season"]].append(p)
    return {
        "source": source,
        "items": len(items),
        "overall": summarize(pairs),
        "by_band": by_band,
        "by_primary_problem": {str(k): {kk: v for kk, v in summarize(v).items() if kk != "confusion"} for k, v in groups.items() if k},
        "by_season": {k: {kk: v for kk, v in summarize(v).items() if kk != "confusion"} for k, v in sorted(seasons.items()) if k},
        "skipped": dict(skipped),
        "engine_version": ENGINE_VERSION,
    }


def wilson(successes: int, n: int, z: float = 1.96) -> tuple[float, float] | None:
    if n == 0:
        return None
    p = successes / n
    denom = 1 + z * z / n
    centre = (p + z * z / (2 * n)) / denom
    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / denom
    return round(max(0.0, centre - half), 4), round(min(1.0, centre + half), 4)


def empirical_table(conn: Connection, center_id: int, config: dict, source: str = "imported", season: str | None = None,
                    mode: str = "single", era: str | None = None, min_n: int = 10) -> dict:
    """P(final danger | likelihood, size class) from bands with one located problem
    (mode='single') or using each band's dominant problem (mode='dominant')."""
    items = load_tool_runs(conn, center_id, season) if source == "tool" else load_forecasts(conn, center_id, season, era)
    cells: dict[str, Counter] = defaultdict(Counter)
    for it in items:
        for b in BANDS:
            final = it[f"danger_{b}"]
            if not final:
                continue
            located = [
                p for p in (it["problems"] or [])
                if p.get("likelihood") and p.get("size_max") is not None and b in bands_of(p.get("locations"))
            ]
            if not located or (mode == "single" and len(located) != 1):
                continue
            chosen = max(located, key=lambda p: (danger_for(int(p["likelihood"]), float(p["size_max"]), config), -p["rank"]))
            key = f"{int(chosen['likelihood'])}-{size_class(float(chosen['size_max']), config)}"
            cells[key][int(final)] += 1
    out = {}
    for key, counts in cells.items():
        n = sum(counts.values())
        mode_level, mode_count = counts.most_common(1)[0]
        lk, sc = key.split("-")
        out[key] = {
            "likelihood": int(lk),
            "size_class": int(sc),
            "n": n,
            "counts": [counts.get(level, 0) for level in range(6)],
            "mode": mode_level,
            "mode_share": round(mode_count / n, 4),
            "mode_ci": wilson(mode_count, n),
            "current": config["dangerTable"][lk][int(sc) - 1],
            "sparse": n < min_n,
        }
    return {"source": source, "mode": mode, "items": len(items), "min_n": min_n, "cells": out}


def replay_config(conn: Connection, center_id: int, candidate: dict, active: dict, holdout_season: str | None = None) -> dict:
    """Backtest a candidate config against imported history, optionally on one held-out season."""
    items = load_forecasts(conn, center_id, holdout_season)
    pairs_c, _ = band_pairs(items, candidate)
    pairs_a, skipped = band_pairs(items, active)
    by_key = {(p["id"], p["band"]): p for p in pairs_a}
    changed = []
    for p in pairs_c:
        a = by_key.get((p["id"], p["band"]))
        if a and a["suggested"] != p["suggested"]:
            changed.append({"forecast_id": p["id"], "date": p["date"], "band": p["band"], "active": a["suggested"],
                            "candidate": p["suggested"], "final": p["final"]})
    return {
        "holdout_season": holdout_season,
        "active": summarize(pairs_a),
        "candidate": summarize(pairs_c),
        "n_changed": len(changed),
        "changed": changed[:200],
        "skipped": dict(skipped),
    }


def backfill(conn: Connection, center_id: int, config_row: dict) -> int:
    """Store the engine's suggestion for every imported forecast under a config version."""
    n = 0
    for it in load_forecasts(conn, center_id):
        s = suggest_from_recorded(it["problems"] or [], config_row["config"])
        conn.execute(
            """INSERT INTO forecast_suggestions (forecast_id, config_id, engine_version, suggested_upper, suggested_middle,
                   suggested_lower, suggested_overall)
               VALUES (%s, %s, %s, %s, %s, %s, %s)
               ON CONFLICT (forecast_id, config_id, engine_version) DO UPDATE SET
                   suggested_upper = EXCLUDED.suggested_upper, suggested_middle = EXCLUDED.suggested_middle,
                   suggested_lower = EXCLUDED.suggested_lower, suggested_overall = EXCLUDED.suggested_overall""",
            (it["id"], config_row["id"], ENGINE_VERSION, s["upper"], s["middle"], s["lower"], s["overall"]),
        )
        n += 1
    return n


def problem_usage(conn: Connection, center_id: int) -> list[dict]:
    """How often each problem type was listed per season — shows how problem use changed over the years."""
    return fetchall(
        conn,
        """SELECT f.season, p.problem_type, count(*)::int AS n, count(DISTINCT f.id)::int AS forecasts
           FROM forecast_problems p JOIN forecasts f ON f.id = p.forecast_id
           WHERE f.center_id = %s GROUP BY f.season, p.problem_type ORDER BY f.season, p.problem_type""",
        (center_id,),
    )


__all__ = ["agreement", "backfill", "empirical_table", "problem_usage", "replay_config", "suggest_from_recorded", "_f"]