economy_countries

"""
economy_countries.py

Build a country-level toy economy dataset.

Sources:
- World Bank WDI bulk CSV for country metadata and economic indicators.
- REST Countries for capital city and country centroid coordinates.

Output:
- assets/toy/economy/economy_countries.csv
"""

from pathlib import Path
import csv
import tempfile
import zipfile

import pandas as pd
import requests


YEAR = 2023
OUT = Path(__file__).with_name("economy_countries.csv")

WDI_BULK_URL = "https://databank.worldbank.org/data/download/WDI_CSV.zip"
REST_COUNTRIES_URL = "https://restcountries.com/v3.1/all"
REQUEST_TIMEOUT = 120

INDICATORS = {
    "SP.POP.TOTL": "population",
    "NY.GDP.MKTP.CD": "gdp_usd",
    "NY.GDP.PCAP.CD": "gdp_per_capita_usd",
    "FP.CPI.TOTL.ZG": "inflation_pct",
    "SL.UEM.TOTL.ZS": "unemployment_pct",
    "NE.EXP.GNFS.ZS": "exports_pct_gdp",
    "NE.IMP.GNFS.ZS": "imports_pct_gdp",
}


def download_wdi_bulk(tmpdir: str) -> Path:
    print("Downloading World Bank WDI bulk CSV")
    zip_path = Path(tmpdir) / "WDI_CSV.zip"

    with requests.get(WDI_BULK_URL, stream=True, timeout=REQUEST_TIMEOUT) as response:
        response.raise_for_status()
        with zip_path.open("wb") as handle:
            for chunk in response.iter_content(chunk_size=1024 * 1024):
                if chunk:
                    handle.write(chunk)

    return zip_path


def read_wdi_countries(archive: zipfile.ZipFile) -> pd.DataFrame:
    country_file = next(
        name for name in archive.namelist()
        if name.endswith("WDICountry.csv")
    )
    records = []

    with archive.open(country_file) as raw:
        reader = csv.DictReader((line.decode("utf-8-sig") for line in raw))
        for row in reader:
            region = row.get("Region")
            if not region or region == "Aggregates":
                continue

            records.append(
                {
                    "country": row.get("Short Name") or row.get("Table Name"),
                    "iso3": row.get("Country Code"),
                    "iso2": row.get("2-alpha code"),
                    "region": region,
                    "income_level": row.get("Income Group"),
                }
            )

    return (
        pd.DataFrame.from_records(records)
        .sort_values("country")
        .reset_index(drop=True)
    )


def read_wdi_indicators(archive: zipfile.ZipFile) -> dict[str, pd.DataFrame]:
    data_file = next(
        name for name in archive.namelist()
        if name.endswith("WDICSV.csv") or name.endswith("WDIData.csv")
    )
    wanted_codes = set(INDICATORS)
    year_column = str(YEAR)
    records = []

    with archive.open(data_file) as raw:
        reader = csv.DictReader((line.decode("utf-8-sig") for line in raw))
        for row in reader:
            code = row.get("Indicator Code")
            if code not in wanted_codes:
                continue

            records.append(
                {
                    "iso3": row.get("Country Code"),
                    "column": INDICATORS[code],
                    "value": row.get(year_column),
                }
            )

    raw_data = pd.DataFrame.from_records(records)
    indicators = {}

    for column in INDICATORS.values():
        data = raw_data.loc[raw_data["column"] == column, ["iso3", "value"]].copy()
        data.rename(columns={"value": column}, inplace=True)
        data[column] = pd.to_numeric(data[column], errors="coerce")
        indicators[column] = data.drop_duplicates("iso3")

    return indicators


def fetch_wdi_bulk() -> tuple[pd.DataFrame, dict[str, pd.DataFrame]]:
    with tempfile.TemporaryDirectory() as tmpdir:
        zip_path = download_wdi_bulk(tmpdir)
        with zipfile.ZipFile(zip_path) as archive:
            countries = read_wdi_countries(archive)
            indicators = read_wdi_indicators(archive)

    return countries, indicators


def fetch_country_geography() -> pd.DataFrame:
    print("Fetching country capitals and coordinates from REST Countries")
    response = requests.get(
        REST_COUNTRIES_URL,
        params={"fields": "cca3,capital,latlng"},
        timeout=REQUEST_TIMEOUT,
    )
    response.raise_for_status()

    records = []
    for row in response.json():
        latlng = row.get("latlng") or [None, None]
        capital = row.get("capital") or []
        records.append(
            {
                "iso3": row.get("cca3"),
                "capital": capital[0] if capital else "",
                "latitude": latlng[0] if len(latlng) > 0 else None,
                "longitude": latlng[1] if len(latlng) > 1 else None,
            }
        )

    geography = pd.DataFrame.from_records(records)
    geography["latitude"] = pd.to_numeric(geography["latitude"], errors="coerce")
    geography["longitude"] = pd.to_numeric(geography["longitude"], errors="coerce")
    return geography.drop_duplicates("iso3")


def main() -> None:
    countries, indicators = fetch_wdi_bulk()
    geography = fetch_country_geography()

    data = countries.merge(geography, on="iso3", how="left")
    data["year"] = YEAR

    for indicator in indicators.values():
        data = data.merge(indicator, on="iso3", how="left")

    data["trade_pct_gdp"] = data["exports_pct_gdp"] + data["imports_pct_gdp"]

    columns = [
        "country",
        "iso3",
        "iso2",
        "region",
        "income_level",
        "capital",
        "longitude",
        "latitude",
        "year",
        "population",
        "gdp_usd",
        "gdp_per_capita_usd",
        "inflation_pct",
        "unemployment_pct",
        "exports_pct_gdp",
        "imports_pct_gdp",
        "trade_pct_gdp",
    ]

    data = data[columns]
    data.to_csv(OUT, index=False)
    print(f"Done: {len(data)} countries saved to {OUT}")


if __name__ == "__main__":
    main()