from __future__ import annotations

import csv
import re
from pathlib import Path


LABELS = Path("labels.csv")
OUT = Path("organ_involvement_label_mapping.csv")


ORGANS = [
    "brain_neurologic",
    "heart",
    "artery_vascular",
    "immune",
    "intestine_digestive",
    "kidney",
    "liver",
    "lung",
    "muscle_musculoskeletal",
    "pancreas_endocrine",
    "adipose_metabolic",
    "female_reproductive",
    "male_reproductive",
    "neoplasm",
]


def _code_key(code: str) -> tuple[str, int, str]:
    code = code.strip().upper()
    match = re.match(r"^([A-Z])(\d{2})(?:\.?([A-Z0-9]+))?", code)
    if not match:
        raise ValueError(f"Invalid ICD-10 code: {code!r}")
    letter, number, suffix = match.groups()
    return letter, int(number), suffix or ""


def _in_range(code: str, start: str, end: str) -> bool:
    c_letter, c_num, _ = _code_key(code)
    s_letter, s_num, _ = _code_key(start)
    e_letter, e_num, _ = _code_key(end)
    if s_letter == e_letter:
        return c_letter == s_letter and s_num <= c_num <= e_num
    return (
        (s_letter < c_letter < e_letter)
        or (c_letter == s_letter and c_num >= s_num)
        or (c_letter == e_letter and c_num <= e_num)
    )


def _matches_any(code: str, ranges: list[tuple[str, str]]) -> bool:
    return any(_in_range(code, start, end) for start, end in ranges)


def organ_for_icd10(code: str) -> tuple[str, str]:
    code = code.strip().upper()
    if not re.match(r"^[A-Z]\d{2}", code):
        return "", "unmapped_non_icd10"

    if _matches_any(code, [("C00", "D48")]):
        return "neoplasm", "neoplasm_c00_d48"

    if _matches_any(code, [("F00", "F09"), ("G00", "G99"), ("I60", "I69")]):
        return "brain_neurologic", "f00_f09_g00_g99_i60_i69"

    if _matches_any(code, [("I00", "I09"), ("I20", "I52")]):
        return "heart", "i00_i09_i20_i52"

    if _matches_any(code, [("I10", "I15"), ("I70", "I89"), ("I95", "I99")]):
        return "artery_vascular", "i10_i15_i70_i89_i95_i99"

    if _matches_any(code, [("A00", "B99"), ("D50", "D89")]):
        return "immune", "a00_b99_d50_d89"

    if _matches_any(code, [("J00", "J99")]):
        return "lung", "j00_j99"

    if _matches_any(code, [("K70", "K77")]):
        return "liver", "k70_k77"

    if _matches_any(code, [("K85", "K86"), ("E10", "E16")]):
        return "pancreas_endocrine", "k85_k86_e10_e16"

    if _matches_any(code, [("K00", "K69"), ("K78", "K84"), ("K87", "K93")]):
        return "intestine_digestive", "k00_k69_k78_k84_k87_k93"

    if _matches_any(code, [("N00", "N39")]):
        return "kidney", "n00_n39"

    if _matches_any(code, [("N70", "N98"), ("O00", "O99")]):
        return "female_reproductive", "n70_n98_o00_o99"

    if _matches_any(code, [("N40", "N53")]):
        return "male_reproductive", "n40_n53"

    if _matches_any(code, [("M00", "M99")]):
        return "muscle_musculoskeletal", "m00_m99"

    if _matches_any(code, [("E00", "E09"), ("E17", "E90")]):
        return "adipose_metabolic", "e00_e09_e17_e90"

    return "", "unmapped_no_organ_rule"


def main() -> None:
    rows = []
    with LABELS.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 ""
            organ_id, match_source = organ_for_icd10(code)
            rows.append(
                {
                    "token_id": i + 3,
                    "label_code": code,
                    "label_name": name,
                    "organ_id": organ_id,
                    "organ_label": organ_id,
                    "organ_weight": 1.0 if organ_id else 0.0,
                    "match_source": match_source,
                    "mapping_source": (
                        "organ-age-inspired clinical systems based on "
                        "Oh et al. Nature 2023; single-label ICD-10 rules"
                    ),
                }
            )

    fieldnames = [
        "token_id",
        "label_code",
        "label_name",
        "organ_id",
        "organ_label",
        "organ_weight",
        "match_source",
        "mapping_source",
    ]
    with OUT.open("w", newline="", encoding="utf-8-sig") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    mapped = [row for row in rows if row["organ_id"]]
    print(f"labels: {len(rows)}")
    print(f"mapped_labels: {len(mapped)}")
    print(f"unmapped_labels: {len(rows) - len(mapped)}")
    print(f"organs: {', '.join(ORGANS)}")
    print(f"output: {OUT}")


if __name__ == "__main__":
    main()