from __future__ import annotations

import csv
from pathlib import Path


LABELS = Path("labels.csv")
OUT = Path("uk_hfrs_label_mapping.csv")
MISSING_OUT = Path("uk_hfrs_missing_label_codes.csv")


# Source: Gilbert T, Neuburger J, Kraindler J, et al. Development and
# validation of a Hospital Frailty Risk Score focusing on older people in
# acute care settings using electronic hospital records. Lancet. 2018.
# Supplementary appendix, Table A2.
UK_HFRS_WEIGHTS = {
    "F00": 7.1,
    "G81": 4.4,
    "G30": 4.0,
    "I69": 3.7,
    "R29": 3.6,
    "N39": 3.2,
    "F05": 3.2,
    "W19": 3.2,
    "S00": 3.2,
    "R31": 3.0,
    "B96": 2.9,
    "R41": 2.7,
    "R26": 2.6,
    "I67": 2.6,
    "R56": 2.6,
    "R40": 2.5,
    "T83": 2.4,
    "S06": 2.4,
    "S42": 2.3,
    "E87": 2.3,
    "M25": 2.3,
    "E86": 2.3,
    "R54": 2.2,
    "Z50": 2.1,
    "F03": 2.1,
    "W18": 2.1,
    "Z75": 2.0,
    "F01": 2.0,
    "S80": 2.0,
    "L03": 2.0,
    "H54": 1.9,
    "E53": 1.9,
    "Z60": 1.8,
    "G20": 1.8,
    "R55": 1.8,
    "S22": 1.8,
    "K59": 1.8,
    "N17": 1.8,
    "L89": 1.7,
    "Z22": 1.7,
    "B95": 1.7,
    "L97": 1.6,
    "R44": 1.6,
    "K26": 1.6,
    "I95": 1.6,
    "N19": 1.6,
    "A41": 1.6,
    "Z87": 1.5,
    "J96": 1.5,
    "X59": 1.5,
    "M19": 1.5,
    "G40": 1.5,
    "M81": 1.4,
    "S72": 1.4,
    "S32": 1.4,
    "E16": 1.4,
    "R94": 1.4,
    "N18": 1.4,
    "R33": 1.3,
    "R69": 1.3,
    "N28": 1.3,
    "R32": 1.2,
    "G31": 1.2,
    "Y95": 1.2,
    "S09": 1.2,
    "R45": 1.2,
    "G45": 1.2,
    "Z74": 1.1,
    "M79": 1.1,
    "W06": 1.1,
    "S01": 1.1,
    "A04": 1.1,
    "A09": 1.1,
    "J18": 1.1,
    "J69": 1.0,
    "R47": 1.0,
    "E55": 1.0,
    "Z93": 1.0,
    "R02": 1.0,
    "R63": 0.9,
    "H91": 0.9,
    "W10": 0.9,
    "W01": 0.9,
    "E05": 0.9,
    "M41": 0.9,
    "R13": 0.8,
    "Z99": 0.8,
    "U80": 0.8,
    "M80": 0.8,
    "K92": 0.8,
    "I63": 0.8,
    "N20": 0.7,
    "F10": 0.7,
    "Y84": 0.7,
    "R00": 0.7,
    "J22": 0.7,
    "Z73": 0.6,
    "R79": 0.6,
    "Z91": 0.5,
    "S51": 0.5,
    "F32": 0.5,
    "M48": 0.5,
    "E83": 0.4,
    "M15": 0.4,
    "D64": 0.4,
    "L08": 0.4,
    "R11": 0.3,
    "K52": 0.3,
    "R50": 0.1,
}


def _read_labels(path: Path) -> list[dict[str, str | int]]:
    rows: list[dict[str, str | int]] = []
    with path.open(encoding="utf-8") as f:
        for i, line in enumerate(f):
            line = line.strip()
            if not line:
                continue
            parts = line.split(maxsplit=1)
            code = parts[0].strip().upper()
            name = parts[1].strip() if len(parts) > 1 else ""
            rows.append({"token_id": i + 3, "label_code": code, "label_name": name})
    return rows


def main() -> None:
    labels = _read_labels(LABELS)
    label_codes = {str(row["label_code"]) for row in labels}
    missing = sorted(set(UK_HFRS_WEIGHTS) - label_codes)

    rows = []
    for row in labels:
        code = str(row["label_code"])
        weight = float(UK_HFRS_WEIGHTS.get(code, 0.0))
        rows.append(
            {
                **row,
                "hfrs_dimension_id": "hfrs_weighted_disease_expression",
                "hfrs_dimension": "DeepHealth HFRS-weighted disease expression",
                "hfrs_key_area": "UK-HFRS",
                "hfrs_weight": weight,
                "hfrs_source": (
                    "Gilbert et al. Lancet 2018 supplementary appendix Table A2"
                ),
                "match_source": "exact_three_character_icd10" if weight else "not_in_hfrs",
            }
        )

    fieldnames = [
        "token_id",
        "label_code",
        "label_name",
        "hfrs_dimension_id",
        "hfrs_dimension",
        "hfrs_key_area",
        "hfrs_weight",
        "hfrs_source",
        "match_source",
    ]
    with OUT.open("w", newline="", encoding="utf-8-sig") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    with MISSING_OUT.open("w", newline="", encoding="utf-8-sig") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "hfrs_source_code",
                "hfrs_weight",
                "missing_reason",
                "hfrs_source",
            ],
        )
        writer.writeheader()
        for code in missing:
            writer.writerow(
                {
                    "hfrs_source_code": code,
                    "hfrs_weight": UK_HFRS_WEIGHTS[code],
                    "missing_reason": "not_present_in_labels_csv",
                    "hfrs_source": (
                        "Gilbert et al. Lancet 2018 supplementary appendix Table A2"
                    ),
                }
            )

    nonzero = sum(1 for row in rows if float(row["hfrs_weight"]) != 0.0)
    print(f"labels: {len(rows)}")
    print(f"uk_hfrs_codes: {len(UK_HFRS_WEIGHTS)}")
    print(f"matched_nonzero_labels: {nonzero}")
    print(f"missing_hfrs_codes: {len(missing)}")
    print(f"output: {OUT}")
    print(f"missing_output: {MISSING_OUT}")


if __name__ == "__main__":
    main()