"""Create a compact calendar-dated disease-event index from ``ukb_data.csv``.

Unlike ``prepare_data.py``, this ETL does not create model-ready relative-time
sequences or other-information tokens.  It writes one structured ``.npy`` file
with exactly three fields:

    eid         int64
    event_date  datetime64[D]
    token       int32

``token`` follows the existing ``labels.csv`` convention used by
``prepare_data.py``: padding=0, checkup=1 (not emitted here), and the first
label in ``labels.csv`` receives token 2.  Each ``(eid, token)`` is deduplicated
to the first known event date.

The output is intended for calendar-indexed temperature and air-pollution
queries.  It contains no date of birth, sex, covariates, or checkup events.

Usage
-----
python prepare_event_dates.py --output ukb_disease_event_dates.npy
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
try:
    from tqdm import tqdm
except ImportError:  # Keep the ETL runnable in minimal Python environments.
    tqdm = None


EVENT_DTYPE = np.dtype(
    [
        ("eid", "<i8"),
        ("event_date", "datetime64[D]"),
        ("token", "<i4"),
    ]
)


def load_label_tokens(labels_file: str | Path) -> dict[str, int]:
    """Return the label-code -> token mapping shared with ``prepare_data.py``."""
    token_map: dict[str, int] = {}
    with Path(labels_file).open(encoding="utf-8") as handle:
        for index, line in enumerate(handle):
            code = line.strip().split(" ", maxsplit=1)[0]
            if code:
                token_map[code] = index + 2
    return token_map


def build_raw_column_map(
    field_map_file: str | Path,
    icd_map_file: str | Path,
) -> tuple[dict[str, str], list[str]]:
    """Build raw-column renames and the calendar-date event columns to inspect."""
    field_df = pd.read_csv(field_map_file, low_memory=False)
    required = {"field_instance", "var_name"}
    missing = required - set(field_df.columns)
    if missing:
        raise ValueError(f"{field_map_file} is missing columns: {sorted(missing)}")

    raw_to_name = dict(
        zip(field_df["field_instance"].astype(str), field_df["var_name"].astype(str))
    )
    icd_date_columns: list[str] = []
    with Path(icd_map_file).open(encoding="utf-8") as handle:
        for line in handle:
            parts = line.strip().split()
            if len(parts) >= 6:
                raw_to_name[parts[0]] = parts[5]
                icd_date_columns.append(parts[5])

    for slot in range(17):
        raw_to_name[f"40005-{slot}.0"] = f"cancer_date_{slot}"
        raw_to_name[f"40006-{slot}.0"] = f"cancer_type_{slot}"
    return raw_to_name, icd_date_columns


def _records_from_icd_columns(
    chunk: pd.DataFrame,
    event_columns: list[str],
    token_map: dict[str, int],
) -> list[pd.DataFrame]:
    frames: list[pd.DataFrame] = []
    for column in event_columns:
        token = token_map.get(column)
        if token is None or column not in chunk.columns:
            continue
        event_date = pd.to_datetime(chunk[column], format="%Y-%m-%d", errors="coerce")
        valid = event_date.notna()
        if valid.any():
            frames.append(
                pd.DataFrame(
                    {
                        "eid": chunk.index[valid].astype("int64"),
                        "event_date": event_date.loc[valid].to_numpy(),
                        "token": np.full(valid.sum(), token, dtype=np.int32),
                    }
                )
            )
    return frames


def _records_from_cancer_columns(chunk: pd.DataFrame, token_map: dict[str, int]) -> list[pd.DataFrame]:
    frames: list[pd.DataFrame] = []
    for slot in range(17):
        date_column = f"cancer_date_{slot}"
        type_column = f"cancer_type_{slot}"
        if date_column not in chunk.columns or type_column not in chunk.columns:
            continue
        event_date = pd.to_datetime(chunk[date_column], format="%Y-%m-%d", errors="coerce")
        code = chunk[type_column].astype("string").str.slice(0, 3)
        token = code.map(token_map)
        valid = event_date.notna() & token.notna()
        if valid.any():
            frames.append(
                pd.DataFrame(
                    {
                        "eid": chunk.index[valid].astype("int64"),
                        "event_date": event_date.loc[valid].to_numpy(),
                        "token": token.loc[valid].astype("int32").to_numpy(),
                    }
                )
            )
    return frames


def prepare_event_dates(
    *,
    ukb_data_file: str | Path,
    field_map_file: str | Path,
    icd_map_file: str | Path,
    labels_file: str | Path,
    output_file: str | Path,
    chunksize: int = 10_000,
) -> int:
    """Stream the raw UKB export, then write a sorted structured event array."""
    token_map = load_label_tokens(labels_file)
    raw_to_name, icd_date_columns = build_raw_column_map(field_map_file, icd_map_file)
    event_columns = [*icd_date_columns, "Death"]
    frames: list[pd.DataFrame] = []

    reader = pd.read_csv(
        ukb_data_file,
        chunksize=chunksize,
        index_col=0,  # UKB participant ID / eid
        low_memory=False,
    )
    chunk_iterator = tqdm(reader, desc="Extracting calendar-dated disease events") if tqdm else reader
    for raw_chunk in chunk_iterator:
        chunk = raw_chunk.rename(columns=raw_to_name)
        frames.extend(_records_from_icd_columns(chunk, event_columns, token_map))
        frames.extend(_records_from_cancer_columns(chunk, token_map))

    if not frames:
        result = np.empty(0, dtype=EVENT_DTYPE)
    else:
        events = pd.concat(frames, ignore_index=True)
        events = events.dropna(subset=["eid", "event_date", "token"])
        events = events.sort_values(["eid", "token", "event_date"], kind="stable")
        # Match prepare_data.py: first occurrence of each disease/death token.
        events = events.drop_duplicates(["eid", "token"], keep="first")
        events = events.sort_values(["eid", "event_date", "token"], kind="stable")

        result = np.empty(len(events), dtype=EVENT_DTYPE)
        result["eid"] = events["eid"].to_numpy(dtype=np.int64)
        result["event_date"] = events["event_date"].to_numpy(dtype="datetime64[D]")
        result["token"] = events["token"].to_numpy(dtype=np.int32)

    np.save(output_file, result)
    return len(result)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--ukb-data", default="ukb_data.csv")
    parser.add_argument("--field-map", default="field_ids_enriched.csv")
    parser.add_argument("--icd-map", default="icd10_codes_mod.tsv")
    parser.add_argument("--labels", default="labels.csv")
    parser.add_argument("--output", default="ukb_disease_event_dates.npy")
    parser.add_argument("--chunksize", type=int, default=10_000)
    args = parser.parse_args()
    if args.chunksize <= 0:
        raise ValueError("chunksize must be positive")
    count = prepare_event_dates(
        ukb_data_file=args.ukb_data,
        field_map_file=args.field_map,
        icd_map_file=args.icd_map,
        labels_file=args.labels,
        output_file=args.output,
        chunksize=args.chunksize,
    )
    print(f"Wrote {count:,} first disease/death events to {args.output}")


if __name__ == "__main__":
    main()
