The code behind every chart

"""Declarative dataset registry for the Lewis website — the single source of truth
for build_cache.py.

Lewis has no Anu ``series_registry.json``; its data is a tree of CSV/Excel files under
``Outputs/Data/`` (many with a ``[YYYY.MM.DD] `` filename prefix). Each entry below maps a
stable cache **key** to a glob (relative to ``Outputs/``), display metadata, a logical
group, and — where the data has a live public-API source — a ``live`` adapter mapping that
``refresh.py`` uses to snapshot fresh data.

``glob`` is resolved by build_cache.resolve() which picks the most recent matching file, so
the date-prefix churn never breaks the build.

Groups drive the catalog + explorer sections:
  bop          per-country Balance of Payments panels (US/UK/Germany)
  fof          U.S. Flow of Funds / International Investment Position
  fred         FRED/ALFRED trade & exchange-rate series (long)
  worldbank    World Bank WDI panels (wide per-indicator + long all-country)
  itv          International Transfer of Value / unequal-exchange panels
  global       cross-country & regional coverage summaries
  unified      harmonized multi-source long panels
"""
from __future__ import annotations

import pandas as pd

# adapter shorthand
_WB = {"adapter": "worldbank", "cadence": "weekly"}
_FRED = {"adapter": "fred", "cadence": "weekly"}

DATASETS: dict[str, dict] = {
    # ---- Balance of Payments (constructed panels; not live) ----
    "us_bop": {
        "glob": "Data/Results/*USdata_annual_pct.csv",
        "name": "United States — Balance of Payments (% GDP)",
        "group": "bop", "source": "BEA International Transactions (BPM6)",
        "units": "% of GDP", "country_scope": "US", "year_range": [1960, 2024],
    },
    "uk_bop": {
        "glob": "Data/Results/*UKdata_annual_pct.csv",
        "name": "United Kingdom — Balance of Payments (% GDP)",
        "group": "bop", "source": "ONS Balance of Payments (Pink Book)",
        "units": "% of GDP", "country_scope": "UK", "year_range": [1948, 2023],
    },
    "ger_bop": {
        "glob": "Data/Results/*GERdata_annual_pct.csv",
        "name": "Germany — Balance of Payments (% GDP)",
        "group": "bop", "source": "Deutsche Bundesbank Balance of Payments",
        "units": "% of GDP", "country_scope": "DE", "year_range": [1991, 2024],
    },
    "cross_country_bop": {
        "glob": "Data/Results/*cross_country_comparison.xlsx",
        "name": "Cross-country BoP comparison",
        "group": "bop", "source": "BEA / ONS / Bundesbank (harmonized)",
        "units": "% of GDP", "country_scope": "US, UK, DE", "year_range": [1960, 2024],
    },

    # ---- Flow of Funds / IIP (U.S.) ----
    "bea_iip": {
        "glob": "Data/Robin/FLOW_OF_FUNDS/BEA_IIP/*bea_iip_data.csv",
        "name": "U.S. International Investment Position (BEA, long)",
        "group": "fof", "source": "BEA International Investment Position",
        "units": "USD millions", "country_scope": "US", "year_range": [1976, 2024],
    },
    "row_us_holdings": {
        "glob": "Data/Robin/FLOW_OF_FUNDS/ROW_HOLDINGS/*row_us_holdings.csv",
        "name": "Foreign (Rest-of-World) holdings of U.S. assets, by instrument",
        "group": "fof",
        "source": "Federal Reserve Z.1 Financial Accounts — Rest of the World (FRED BOGZ1LM2630*)",
        "units": "USD millions (market value)", "country_scope": "US (vs Rest of World)",
        "year_range": [1945, 2025],
    },
    "us_gov_borrowing": {
        "glob": "Data/Robin/FLOW_OF_FUNDS/GOV_BORROWING/*us_gov_borrowing.csv",
        "name": "U.S. federal government balance (surplus/deficit)",
        "group": "fof",
        "source": "U.S. Treasury / OMB via FRED (FYFSGDA188S, FYFSD)",
        "units": "% of GDP and USD millions", "country_scope": "US",
        "year_range": [1901, 2025],
    },

"""Build the canonical Parquet cache + site manifest for the Lewis website.

Reads the declarative registry in ``datasets.py`` (single source of truth), resolves each
glob against ``Outputs/`` (newest match wins, so the ``[YYYY.MM.DD] `` filename prefix never
breaks the build), normalizes to a canonical schema, writes one Parquet per dataset under
``site_data/cache/``, and emits ``site_manifest.json`` which drives the Explorer, Catalog,
downloads, freshness badges, and the country index.

Run:  PYTHONIOENCODING=utf-8 python webapp/data_pipeline/build_cache.py
"""
from __future__ import annotations

import json
import sys
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from app import config as C  # noqa: E402
from app.services import country_service  # noqa: E402
from data_pipeline.datasets import DATASETS, composite_coverage  # noqa: E402
from data_pipeline import units as U  # noqa: E402

ID_COLS = {"country_code", "country_name", "year", "date", "indicator",
           "indicator_code", "indicator_name", "units", "source", "frequency"}


def _now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def resolve(glob_rel: str) -> Path | None:
    """Newest file under Outputs/ matching the glob (handles date-prefixed names)."""
    matches = sorted(C.OUTPUTS.glob(glob_rel), key=lambda p: p.stat().st_mtime, reverse=True)
    return matches[0] if matches else None


def read_panel(path: Path) -> pd.DataFrame:
    try:
        if path.suffix.lower() in (".xlsx", ".xls"):
            return pd.read_excel(path)
        return pd.read_csv(path, low_memory=False)
    except Exception as e:  # noqa: BLE001
        print(f"  ! read error {path.name}: {e}")
        return pd.DataFrame()


def normalize(df: pd.DataFrame) -> pd.DataFrame:
    if df.empty:
        return df
    df = df.loc[:, ~df.columns.astype(str).str.match(r"Unnamed: ?\d+")]
    rename: dict = {}
    for c in df.columns:
        cl = str(c).strip()
        if cl == "Year":
            rename[c] = "year"
        elif cl == "Date":
            rename[c] = "date"
        elif cl == "country_iso3":
            rename[c] = "country_code"
    df = df.rename(columns=rename)
    # promote a bare 'country' code column when there is no explicit country_code
    if "country_code" not in df.columns and "country" in df.columns and "country_name" in df.columns:
        df = df.rename(columns={"country": "country_code"})
    if "year" in df.columns:
        yr = pd.to_numeric(df["year"], errors="coerce")
        try:
            df["year"] = yr.astype("Int64")
        except TypeError:
            # fractional / non-integer values in a 'year' column — keep numeric float
            df["year"] = yr
    return df


def metric_cols(df: pd.DataFrame) -> list[str]:
    return [str(c) for c in df.columns if c not in ID_COLS and c not in ("year", "date")]

"""World Bank WDI adapter (no key required).

Docs: https://api.worldbank.org/v2/
Maps a Lewis dataset key to a WDI indicator and returns a tidy DataFrame
(country_code, country_name, year, <indicator>).
"""
from __future__ import annotations

import requests
import pandas as pd

BASE = "https://api.worldbank.org/v2"

# Lewis dataset key -> WDI indicator
INDICATORS = {
    "wb_current_account_pct_gdp": "BN.CAB.XOKA.GD.ZS",
    "wb_current_account_usd": "BN.CAB.XOKA.CD",
    "wb_gdp_current_usd": "NY.GDP.MKTP.CD",
    "wb_exports_usd": "BX.GSR.GNFS.CD",
    "wb_imports_usd": "BM.GSR.GNFS.CD",
}


def fetch_indicator(indicator: str) -> pd.DataFrame:
    rows: list[dict] = []
    page = 1
    while True:
        r = requests.get(f"{BASE}/country/all/indicator/{indicator}",
                         params={"format": "json", "per_page": 20000, "page": page}, timeout=60)
        r.raise_for_status()
        js = r.json()
        if not isinstance(js, list) or len(js) < 2 or js[1] is None:
            break
        meta, data = js[0], js[1]
        for d in data:
            rows.append({"country_code": (d.get("countryiso3code") or "").strip(),
                         "country_name": (d.get("country") or {}).get("value"),
                         "year": int(d["date"]) if d.get("date") else None,
                         "value": d.get("value")})
        if page >= meta.get("pages", 1):
            break
        page += 1
    df = pd.DataFrame(rows).dropna(subset=["year"])
    if not df.empty:
        df["year"] = df["year"].astype(int)
    return df


def fetch(dataset_key: str) -> pd.DataFrame:
    ind = INDICATORS.get(dataset_key)
    if not ind:
        raise ValueError(f"worldbank adapter has no indicator for {dataset_key}")
    return fetch_indicator(ind).rename(columns={"value": ind})

"""FRED adapter (key from FRED_API_KEY env var via config).

Docs: https://fred.stlouisfed.org/docs/api/fred/
Fetches the U.S. external-sector trade series behind the ``fred_trade`` panel,
returned long (date, year, value, series_id) to match the cached schema.
"""
from __future__ import annotations

import requests
import pandas as pd

from app import config as C

BASE = "https://api.stlouisfed.org/fred/series/observations"

# Lewis dataset key -> FRED series ids
SERIES_MAP = {
    "fred_trade": ["BOPGSTB", "BOPGTB", "BOPTEXP", "BOPTIMP", "BOPBCA", "NETFI"],
}


def _fetch_one(fred_id: str, key: str) -> pd.DataFrame:
    r = requests.get(BASE, params={"series_id": fred_id, "api_key": key, "file_type": "json"}, timeout=60)
    r.raise_for_status()
    obs = r.json().get("observations", [])
    df = pd.DataFrame(obs)
    if df.empty:
        return df
    df["value"] = pd.to_numeric(df["value"].replace(".", None), errors="coerce")
    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    df["year"] = df["date"].dt.year
    df["series_id"] = fred_id
    return df[["date", "year", "value", "series_id"]]


def fetch(dataset_key: str) -> pd.DataFrame:
    key = C.get_api_key("fred")
    if not key:
        raise RuntimeError("FRED API key not available; set FRED_API_KEY "
                           "(free key: https://fred.stlouisfed.org/docs/api/api_key.html)")
    ids = SERIES_MAP.get(dataset_key)
    if not ids:
        raise ValueError(f"fred adapter has no mapping for {dataset_key}")
    frames = [_fetch_one(fid, key) for fid in ids]
    frames = [f for f in frames if not f.empty]
    return pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()

"""Live refresh orchestrator.

For each dataset with a live adapter (see datasets.py): fetch via the adapter ->
validate -> snapshot to site_data/raw_feeds/ + bump freshness log, OR keep last-good
and log the failure. A failed fetch never fabricates or overwrites the curated cache.

The display cache holds curated Lewis panels; overwriting them with a raw API dump would
change the schema the charts depend on. Refresh therefore keeps a validated raw snapshot +
freshness stamp; promoting raw feeds into display panels is a deliberate rebuild step.

Run:  python webapp/data_pipeline/refresh.py [--only worldbank|fred] [--key wb_current_account_pct_gdp] [--dry-run]
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from app import config as C  # noqa: E402
from data_pipeline import validators  # noqa: E402
from data_pipeline.datasets import live_targets  # noqa: E402
from data_pipeline.adapters import worldbank, fred  # noqa: E402

ADAPTERS = {"worldbank": worldbank, "fred": fred}
RAW_FEEDS = C.SITE_DATA / "raw_feeds"


def _now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def _load_log() -> dict:
    return json.loads(C.REFRESH_LOG.read_text(encoding="utf-8")) if C.REFRESH_LOG.exists() else {}


def _save_log(log: dict) -> None:
    C.REFRESH_LOG.write_text(json.dumps(log, indent=2, default=str), encoding="utf-8")


def _atomic_write_parquet(df: pd.DataFrame, path: Path) -> None:
    tmp = path.with_suffix(".parquet.tmp")
    df.to_parquet(tmp, index=False)
    os.replace(tmp, path)


def refresh_key(key: str, adapter_name: str, log: dict) -> tuple[bool, str]:
    adapter = ADAPTERS.get(adapter_name)
    if adapter is None:
        return False, f"unknown adapter {adapter_name}"
    RAW_FEEDS.mkdir(parents=True, exist_ok=True)
    raw_path = RAW_FEEDS / f"{key}__{adapter_name}.parquet"
    old = None
    if raw_path.exists():
        try:
            old = pd.read_parquet(raw_path)
        except Exception:
            old = None
    try:
        new = adapter.fetch(key)
    except Exception as e:  # noqa: BLE001
        return False, f"fetch error: {e}"
    ok, reason = validators.validate_fetch(new, old)
    if not ok:
        return False, f"validation failed: {reason} (kept last-good)"
    _atomic_write_parquet(new, raw_path)
    log.setdefault(adapter_name, {})[key] = {
        "last_refresh": _now(), "rows": int(len(new)),
        "raw_feed": f"raw_feeds/{key}__{adapter_name}.parquet"}
    return True, f"snapshot {len(new)} rows -> raw_feeds/"


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--only", choices=list(ADAPTERS))

"""Refresh validators — guard against bad/short/empty fetches.

A fetch that fails validation must NOT overwrite good data (no-freeze, no-synthetic).
Returns (ok, reason).
"""
from __future__ import annotations

import pandas as pd


def validate_fetch(new: pd.DataFrame, old: pd.DataFrame | None,
                   min_rows: int = 1, shrink_tolerance: float = 0.9) -> tuple[bool, str]:
    if new is None or new.empty:
        return False, "fetch returned empty frame"
    if len(new) < min_rows:
        return False, f"too few rows ({len(new)} < {min_rows})"
    id_cols = {"country_code", "country_name", "year", "date", "record_date"}
    data_cols = [c for c in new.columns if c not in id_cols]
    if data_cols and new[data_cols].notna().sum().sum() == 0:
        return False, "all data columns are NaN"
    if old is not None and not old.empty and len(new) < len(old) * shrink_tolerance:
        return False, f"row count shrank ({len(new)} < {len(old)}*{shrink_tolerance})"
    return True, "ok"

"""Chart service — themed Plotly figure builders returned as JSON.

Column names verified against the real cached panels (2026-06-01):
  us_bop / uk_bop / ger_bop : 'year' + '<Component> Balance_pct' columns, e.g.
      'Current Account Balance_pct', 'Merchandise Trade Balance_pct',
      'Service Trade Balance_pct', 'Goods and Services Balance_pct',
      'Primary Income Balance_pct', 'Secondary Income Balance_pct',
      'Capital Account Balance_pct'; US financial account proxy =
      'Asset/Liability Balance_pct'; UK/DE = 'Financial Account Balance_pct'.
  bea_iip  : long  (date, series_id, value, title, ...)
  itv_master / itv_country_means : country_code, year, ca_pct_gdp, cp_class,
      gdp_pc, labor_share, wage_proxy, gdp_usd, ...
  wb_*_pct_gdp / wb_*_usd : wide  (year + ISO3 columns)
  fred_trade : long  (date, value, series_id, category)
  global_country_summary : country/country_name/latest_gdp_usd_billions/total_observations
  global_regional : region/latest_aggregate_gdp_trillions
"""
from __future__ import annotations

import json
from typing import Callable

import pandas as pd
import plotly.graph_objects as go

from app.services import data_service as D
from app.services import country_service as _CS
from app.services.labels import var_label as _var_label


def _cn(code):
    """ISO3 -> full country name for chart legends (falls back to the code)."""
    try:
        return _CS.full_name(str(code))
    except Exception:
        return code

# v1.1: figures ship THEME-NEUTRAL. The Arcanum Site Kit (ark-plotly.js, merged
# client-side in static/js/app.js) supplies the font color, grid colors, the
# accent-led colorway (crimson first for lewis), transparent backgrounds, and
# live light/dark re-theming on the toggle. No server-imposed plotly_white
# template, white backgrounds, or navy PALETTE. The only colors this module
# still owns are semantic DATA encodings (historical-event marker = crimson),
# which read on both light and dark backgrounds.
EVENT_COLOR = "#dc2626"  # crimson accent — historical-event vline/annotation

BOP_KEY = {"us": "us_bop", "uk": "uk_bop", "de": "ger_bop"}
BOP_LABEL = {"us": "United States", "uk": "United Kingdom", "de": "Germany"}
FIN_ACCT = ["Financial Account Balance_pct", "Asset/Liability Balance_pct"]


# ============================ Units guard (UNITS_VALIDATION_STANDARD) =========
# A chart that overlays traces on a shared axis MUST keep them on one unit, or
# declare an explicit normalization (pct_gdp / index / log). This is the guard
# that prevents the Lewis Germany≈0 incident (mixed-unit overlay) from recurring.
import logging

_LOG = logging.getLogger("lewis.chart_units")
# Acceptable cross-trace magnitude spread on a shared axis: per the rubric a
# series >100x off the others is a units/scale break, not a real difference.
_MAGNITUDE_RATIO_LIMIT = 100.0


def _trace_units(key: str, column: str) -> str | None:
    """Declared units for (dataset, column) from the manifest's column_units."""
    entry = D.dataset_entry(key) or {}
    cu = entry.get("column_units") or {}
    return cu.get(str(column)) or entry.get("units")


def assert_same_unit(traces: list[tuple], *, chart: str) -> tuple[bool, str]:
    """Validate that overlaid traces share one unit and one magnitude band.

    ``traces`` is a list of ``(key, column, series)`` tuples. Returns
    ``(ok, reason)``; logs and returns False if units differ or any trace is
    >100x off the cross-trace median magnitude (the magnitude sanity scan).
    """
    declared = {(_trace_units(k, c) or "?") for k, c, _ in traces if c is not None}
    if len([u for u in declared if u != "?"]) > 1:
        reason = f"mixed units {sorted(declared)} on a shared axis"

"""Download service — stream any cached dataset as csv / xlsx / parquet,
plus a bulk ZIP of every cached dataset (CSV + a provenance manifest).

Per the estate DOWNLOAD_AND_FORMATS standard (CSV / XLSX / Parquet, NO JSON)
JSON is not an offered download format at any layer — the API route rejects a
``.json`` request with 400 (mirrors gerhard's format gate)."""
from __future__ import annotations

import io
import json
import zipfile
from datetime import datetime, timezone

import pandas as pd
from fastapi import HTTPException
from fastapi.responses import StreamingResponse, FileResponse

from app import config as C
from app.services import data_service as D
from app.services.labels import column_label as _column_label

# DNA-permitted download media only (JSON is banned — see module docstring).
MEDIA = {"csv": "text/csv",
         "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
         "parquet": "application/octet-stream"}


def _prov(key: str) -> str:
    entry = D.dataset_entry(key) or {}
    prov = entry.get("provenance", {}) or {}
    src = prov.get("source") or entry.get("source") or "Lewis"
    return f"# Lewis — {entry.get('name', key)} | source: {src} | exported from the Lewis website\n"


def stream(key: str, fmt: str):
    fmt = fmt.lower()
    if fmt not in MEDIA:
        raise HTTPException(400, f"Unsupported format '{fmt}'")
    df = D.load_dataset(key)
    if df.empty:
        raise HTTPException(404, f"No data for '{key}'")
    filename = f"lewis_{key}.{fmt}"

    if fmt == "parquet":
        path = C.CACHE_DIR / f"{key}.parquet"
        if path.exists():
            return FileResponse(path, media_type=MEDIA[fmt], filename=filename)
        buf = io.BytesIO(); df.to_parquet(buf, index=False); buf.seek(0)
    elif fmt == "csv":
        buf = io.BytesIO((_prov(key) + df.to_csv(index=False)).encode("utf-8"))
    else:  # xlsx
        buf = io.BytesIO()
        with pd.ExcelWriter(buf, engine="openpyxl") as xl:
            df.to_excel(xl, index=False, sheet_name="data")
        buf.seek(0)

    buf.seek(0)
    return StreamingResponse(buf, media_type=MEDIA[fmt],
                             headers={"Content-Disposition": f'attachment; filename="{filename}"'})


def country_itv(iso: str):
    """Stream ONE country's ITV / external-sector panel (itv_master filtered by
    country_code) as CSV — the per-country download for the country-profile ITV
    chart, the only chart class that previously lacked a download."""
    iso = iso.upper()
    df = D.load_dataset("itv_master")
    if df.empty or "country_code" not in df.columns:
        raise HTTPException(404, "ITV master dataset unavailable.")
    d = df[df["country_code"] == iso]
    if d.empty:
        raise HTTPException(404, f"No ITV panel for '{iso}'.")
    header = (f"# Lewis — ITV / value-transfer indicators for {iso} | "
              f"source: itv_master | exported from the Lewis website\n")
    buf = io.BytesIO((header + d.to_csv(index=False)).encode("utf-8"))
    buf.seek(0)
    return StreamingResponse(buf, media_type=MEDIA["csv"],
                             headers={"Content-Disposition": f'attachment; filename="lewis_{iso}_itv.csv"'})