from __future__ import annotations

import csv
import re
from pathlib import Path


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


ICD10_CHAPTERS = [
    ("A00", "B99", "I", "Certain infectious and parasitic diseases", "infection_immune"),
    ("C00", "D48", "II", "Neoplasms", "neoplasm"),
    ("D50", "D89", "III", "Diseases of the blood and immune mechanism", "hematologic_immune"),
    ("E00", "E90", "IV", "Endocrine, nutritional and metabolic diseases", "metabolic_endocrine"),
    ("F00", "F99", "V", "Mental and behavioural disorders", "mental_behavioral"),
    ("G00", "G99", "VI", "Diseases of the nervous system", "neurologic"),
    ("H00", "H59", "VII", "Diseases of the eye and adnexa", "sensory_eye"),
    ("H60", "H95", "VIII", "Diseases of the ear and mastoid process", "sensory_ear"),
    ("I00", "I99", "IX", "Diseases of the circulatory system", "cardiovascular"),
    ("J00", "J99", "X", "Diseases of the respiratory system", "respiratory"),
    ("K00", "K93", "XI", "Diseases of the digestive system", "digestive_liver"),
    ("L00", "L99", "XII", "Diseases of the skin and subcutaneous tissue", "skin_soft_tissue"),
    ("M00", "M99", "XIII", "Diseases of the musculoskeletal system", "musculoskeletal"),
    ("N00", "N99", "XIV", "Diseases of the genitourinary system", "genitourinary_renal"),
    ("O00", "O99", "XV", "Pregnancy, childbirth and puerperium", "pregnancy_reproductive"),
    ("P00", "P96", "XVI", "Certain conditions originating in the perinatal period", "perinatal"),
    ("Q00", "Q99", "XVII", "Congenital malformations", "congenital"),
    ("R00", "R99", "XVIII", "Symptoms, signs and abnormal findings", "symptoms_findings"),
    ("S00", "T98", "XIX", "Injury, poisoning and external causes", "injury_poisoning"),
    ("V01", "Y98", "XX", "External causes of morbidity and mortality", "external_causes"),
    ("Z00", "Z99", "XXI", "Factors influencing health status", "health_status_factors"),
    ("U00", "U99", "XXII", "Codes for special purposes", "special_purpose"),
]


ORGAN_OVERRIDES = [
    ("I60", "I69", "brain_vascular"),
    ("N17", "N19", "kidney"),
    ("K70", "K77", "liver"),
    ("E10", "E14", "diabetes_metabolic"),
    ("F00", "F09", "cognition_neuropsychiatric"),
    ("G30", "G32", "cognition_neurodegenerative"),
    ("H53", "H54", "sensory_vision"),
    ("H90", "H91", "sensory_hearing"),
    ("J40", "J47", "chronic_respiratory"),
    ("C00", "D48", "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 _code_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 _chapter_for_code(code: str) -> tuple[str, str, str]:
    if not re.match(r"^[A-Z]\d{2}", code.strip().upper()):
        return "", "", "unmapped"
    for start, end, chapter_id, chapter_name, default_system in ICD10_CHAPTERS:
        if _code_in_range(code, start, end):
            return chapter_id, chapter_name, default_system
    return "", "", "unmapped"


def _organ_system_for_code(code: str, default_system: str) -> str:
    if default_system == "unmapped":
        return "unmapped"
    for start, end, organ_system in ORGAN_OVERRIDES:
        if _code_in_range(code, start, end):
            return organ_system
    return default_system


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 ""
            chapter_id, chapter_name, default_system = _chapter_for_code(code)
            organ_system = _organ_system_for_code(code, default_system)
            rows.append(
                {
                    "token_id": i + 3,
                    "label_code": code,
                    "label_name": name,
                    "icd10_chapter_id": chapter_id,
                    "icd10_chapter_name": chapter_name,
                    "organ_system": organ_system,
                    "organ_weight": 1.0 if organ_system != "unmapped" else 0.0,
                    "match_source": "icd10_chapter_range",
                }
            )

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

    print(f"labels: {len(rows)}")
    print(f"output: {OUT}")
    systems = sorted({row["organ_system"] for row in rows})
    print("organ_systems:", ", ".join(systems))


if __name__ == "__main__":
    main()