#!/usr/bin/env python3
"""Reproduce The Economist's GDP-adjusted Big Mac regression, edition by edition.

Mirrors data-generator-v3.R in github.com/TheEconomist/big-mac-data:
  * regression sample = regression_countries (+2021 add-ons from 2021-01-01, +BGR from 2026-07-01)
    -> euro-area members enter individually alongside EUZ, so 2026-07 has 70 observations
  * GDP_bigmac = GDP_local / (local_price / US local_price)      (GDP per person at Big Mac PPP)
  * adj_price  = OLS fit of dollar_price on GDP_bigmac, per edition
  * USD_adjusted = (dollar_price/adj_price) / (US dollar_price/US adj_price) - 1

Usage:
  python3 bigmac_gdp_regression.py path/to/big-mac-source-data-v2.csv [--check path/to/big-mac-adjusted-index.csv]
                                   [--latest path/to/big-mac-full-index.csv] [--out-dir public/data]
Writes big-mac-gdp-regression-by-edition.csv (to --out-dir, default: next to the source file).
With --latest, also writes big-mac-gdp-adjusted-latest.csv: the newest edition, raw vs adjusted,
with both income measures (GDP_dollar at market rates, GDP_bigmac at Big Mac PPP).

Site refresh (run from bigmacindex-app_astro/, against a clone of TheEconomist/big-mac-data):
  python3 scripts/bigmac_gdp_regression.py ../big-mac-data/source-data/big-mac-source-data-v2.csv \
    --check ../big-mac-data/output-data/big-mac-adjusted-index.csv \
    --latest ../big-mac-data/output-data/big-mac-full-index.csv --out-dir public/data
"""
import csv, math, sys, os, collections

REG = set("ARG AUS BRA GBR CAN CHL CHN CZE DNK EGY EUZ HKG HUN IDN ISR JPN MYS MEX NZL NOR PER PHL POL RUS SAU SGP ZAF "
          "KOR SWE CHE TWN THA TUR USA COL PAK IND AUT BEL NLD FIN FRA DEU IRL ITA PRT ESP GRC EST".split())
ADD_2021 = set("ARE CRI LKA UKR URY VNM GTM HND NIC AZE BHR HRV JOR KWT MDA OMN QAT ROU SVK SVN LVA LTU".split())
ADD_2026 = {"BGR"}


def num(x):
    try:
        return float(x)
    except (TypeError, ValueError):
        return None


def ols(xs, ys):
    n = len(xs); mx = sum(xs) / n; my = sum(ys) / n
    sxy = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    sxx = sum((x - mx) ** 2 for x in xs); syy = sum((y - my) ** 2 for y in ys)
    b = sxy / sxx
    return my - b * mx, b, sxy * sxy / (sxx * syy)


def in_sample(iso, date):
    return iso in REG or (iso in ADD_2021 and date >= "2021-01-01") or (iso in ADD_2026 and date >= "2026-07-01")


def write_latest(src_path, full_path, out_dir):
    """Newest edition, one row per published country with an adjusted figure, raw next to adjusted."""
    rows = [r for r in csv.DictReader(open(full_path, newline=""))]
    latest = max(r["date"] for r in rows)
    gdp_dollar = {r["iso_a3"]: r["GDP_dollar"] for r in csv.DictReader(open(src_path, newline="")) if r["date"] == latest}
    keep = [r for r in rows if r["date"] == latest and r["USD_adjusted"]]
    keep.sort(key=lambda r: float(r["USD_adjusted"]))
    fields = ["date", "iso_a3", "name", "currency_code", "local_price", "dollar_ex", "dollar_price",
              "GDP_dollar", "GDP_bigmac", "adj_price", "USD_raw", "USD_adjusted"]
    out_path = os.path.join(out_dir, "big-mac-gdp-adjusted-latest.csv")
    with open(out_path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=fields); w.writeheader()
        for r in keep:
            w.writerow({**{f: r.get(f, "") for f in fields}, "GDP_dollar": gdp_dollar.get(r["iso_a3"], "")})
    print(f"wrote {out_path} ({len(keep)} countries, {latest})")


def main(src_path, check_path=None, latest_path=None, out_dir=None):
    by_date = collections.defaultdict(list)
    for r in csv.DictReader(open(src_path, newline="")):
        g, lp, de = num(r["GDP_local"]), num(r["local_price"]), num(r["dollar_ex"])
        if not g or g <= 0 or not lp or not de:
            continue
        if in_sample(r["iso_a3"], r["date"]):
            by_date[r["date"]].append((r["iso_a3"], g, lp, de))

    out_rows, fits = [], {}
    for d in sorted(by_date):
        rows = by_date[d]
        us_price = next(lp for iso, g, lp, de in rows if iso == "USA")
        xs = [g / (lp / us_price) for iso, g, lp, de in rows]          # GDP per person at Big Mac PPP
        ys = [lp / de for iso, g, lp, de in rows]                       # dollar price
        a, b, r2 = ols(xs, ys)
        _, _, r2_log = ols([math.log(x) for x in xs], [math.log(y) for y in ys])
        fits[d] = (a, b)
        out_rows.append({"date": d, "n_obs": len(rows), "intercept_usd": round(a, 4),
                         "slope_usd_per_1000_gdp": round(b * 1000, 5), "r2_linear": round(r2, 4),
                         "r2_loglog": round(r2_log, 4)})

    out_dir = out_dir or os.path.dirname(os.path.abspath(src_path))
    out_path = os.path.join(out_dir, "big-mac-gdp-regression-by-edition.csv")
    with open(out_path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(out_rows[0].keys())); w.writeheader(); w.writerows(out_rows)
    print(f"wrote {out_path} ({len(out_rows)} editions)")

    if check_path:  # confirm the fit reproduces The Economist's published adj_price
        worst = 0.0
        for r in csv.DictReader(open(check_path, newline="")):
            a, b = fits[r["date"]]
            worst = max(worst, abs(a + b * float(r["GDP_bigmac"]) - float(r["adj_price"])))
        print(f"max |fitted - published adj_price| = {worst:.2e}")

    if latest_path:
        write_latest(src_path, latest_path, out_dir)


if __name__ == "__main__":
    args = sys.argv[1:]
    opt = lambda name: args[args.index(name) + 1] if name in args else None
    main(args[0], opt("--check"), opt("--latest"), opt("--out-dir"))