#!/usr/bin/env python3
"""Generate the RunCalcs VO2 reference table from public NHANES files.

The script downloads the 1999-2000 and 2001-2002 cardiovascular-fitness
and demographic SAS transport files, merges them by respondent identifier,
and calculates survey-weighted percentiles of estimated VO2 max.
"""

from __future__ import annotations

import csv
import hashlib
import json
import tempfile
import urllib.request
from datetime import date
from pathlib import Path

import numpy as np
import pandas as pd


ROOT = Path(__file__).resolve().parents[1]
OUTPUT = ROOT / "articles/data/vo2-max-by-age-complete-dataset.csv"
METADATA = ROOT / "articles/data/vo2-max-by-age-source-metadata.json"

SOURCES = {
    "1999-2000 cardiovascular fitness": "https://wwwn.cdc.gov/Nchs/Data/Nhanes/Public/1999/DataFiles/CVX.XPT",
    "1999-2000 demographics": "https://wwwn.cdc.gov/Nchs/Data/Nhanes/Public/1999/DataFiles/DEMO.XPT",
    "2001-2002 cardiovascular fitness": "https://wwwn.cdc.gov/Nchs/Data/Nhanes/Public/2001/DataFiles/CVX_B.XPT",
    "2001-2002 demographics": "https://wwwn.cdc.gov/Nchs/Data/Nhanes/Public/2001/DataFiles/DEMO_B.XPT",
}

AGE_GROUPS = (
    (12, 19, "12-19"),
    (20, 29, "20-29"),
    (30, 39, "30-39"),
    (40, 49, "40-49"),
)


def download(url: str, destination: Path) -> str:
    with urllib.request.urlopen(url) as response:
        payload = response.read()
    destination.write_bytes(payload)
    return hashlib.sha256(payload).hexdigest()


def weighted_quantile(values: np.ndarray, weights: np.ndarray, q: float) -> float:
    order = np.argsort(values)
    values = values[order]
    weights = weights[order]
    cumulative = np.cumsum(weights) - (0.5 * weights)
    cumulative = cumulative / np.sum(weights)
    return float(np.interp(q, cumulative, values))


def load_cycle(cv_path: Path, demo_path: Path, cycle: str) -> pd.DataFrame:
    cv = pd.read_sas(cv_path, format="xport")[["SEQN", "CVDEXSTS", "CVDESVO2"]]
    demo = pd.read_sas(demo_path, format="xport")[["SEQN", "RIDAGEYR", "RIAGENDR", "WTMEC4YR"]]
    merged = cv.merge(demo, on="SEQN", how="inner", validate="one_to_one")
    merged["cycle"] = cycle
    return merged


def main() -> None:
    with tempfile.TemporaryDirectory(prefix="runcalcs-nhanes-") as temp_name:
        temp = Path(temp_name)
        files: dict[str, Path] = {}
        hashes: dict[str, str] = {}
        for label, url in SOURCES.items():
            destination = temp / url.rsplit("/", 1)[-1]
            hashes[label] = download(url, destination)
            files[label] = destination

        data = pd.concat(
            [
                load_cycle(
                    files["1999-2000 cardiovascular fitness"],
                    files["1999-2000 demographics"],
                    "1999-2000",
                ),
                load_cycle(
                    files["2001-2002 cardiovascular fitness"],
                    files["2001-2002 demographics"],
                    "2001-2002",
                ),
            ],
            ignore_index=True,
        )

        eligible = data[
            (data["CVDEXSTS"] == 1)
            & data["CVDESVO2"].notna()
            & data["RIDAGEYR"].between(12, 49)
            & data["RIAGENDR"].isin([1, 2])
            & data["WTMEC4YR"].notna()
            & (data["WTMEC4YR"] > 0)
        ].copy()

        rows: list[list[object]] = []
        for minimum, maximum, label in AGE_GROUPS:
            for sex_code, sex_label in ((1, "Male"), (2, "Female")):
                group = eligible[
                    eligible["RIDAGEYR"].between(minimum, maximum)
                    & (eligible["RIAGENDR"] == sex_code)
                ]
                values = group["CVDESVO2"].to_numpy(dtype=float)
                weights = group["WTMEC4YR"].to_numpy(dtype=float)
                percentiles = [weighted_quantile(values, weights, q) for q in (0.10, 0.25, 0.50, 0.75, 0.90)]
                rows.append(
                    [label, sex_label, len(group), *[round(value, 1) for value in percentiles]]
                )

        OUTPUT.parent.mkdir(parents=True, exist_ok=True)
        with OUTPUT.open("w", newline="", encoding="utf-8") as output_file:
            writer = csv.writer(output_file, lineterminator="\n")
            writer.writerow(
                [
                    "Age group",
                    "Sex",
                    "Unweighted sample size",
                    "P10 estimated VO2 max (mL/kg/min)",
                    "P25 estimated VO2 max (mL/kg/min)",
                    "P50 estimated VO2 max (mL/kg/min)",
                    "P75 estimated VO2 max (mL/kg/min)",
                    "P90 estimated VO2 max (mL/kg/min)",
                ]
            )
            writer.writerows(rows)

        metadata = {
            "dataset": "RunCalcs NHANES estimated VO2 max percentiles by age and sex",
            "version": "1.0.0",
            "generated_on": date.today().isoformat(),
            "geography": "United States non-institutionalized population",
            "survey_cycles": ["1999-2000", "2001-2002"],
            "source_publisher": "US CDC National Center for Health Statistics",
            "source_urls": SOURCES,
            "source_sha256": hashes,
            "measure": "CVDESVO2, estimated maximal oxygen uptake from the NHANES submaximal treadmill protocol",
            "units": "mL/kg/min",
            "inclusion_criteria": [
                "Age 12 through 49 years",
                "NHANES cardiovascular fitness exam status equals 1 (VO2 max estimated)",
                "Non-missing CVDESVO2",
                "RIAGENDR coded 1 or 2 in the public-use demographic file",
                "Positive non-missing four-year MEC examination weight (WTMEC4YR)",
            ],
            "analysis": "Weighted P10, P25, P50, P75, and P90 within age group and recorded sex using WTMEC4YR.",
            "eligible_unweighted_sample_size": int(len(eligible)),
            "limitations": [
                "VO2 max was estimated from a submaximal treadmill test, not directly measured by gas analysis.",
                "The public component is limited to participants aged 12-49.",
                "Medical and protocol exclusions mean the tested sample differs from all age-eligible participants.",
                "Percentile estimates shown here do not calculate complex-survey confidence intervals.",
                "NHANES records sex as male or female in RIAGENDR for these survey cycles.",
            ],
        }
        METADATA.write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")

        print(f"Wrote {OUTPUT.relative_to(ROOT)} ({len(rows)} rows)")
        print(f"Wrote {METADATA.relative_to(ROOT)}")
        print(f"Eligible unweighted sample: {len(eligible)}")


if __name__ == "__main__":
    main()