#!/usr/bin/env python3
"""
Phase 1A — Génération de paires ChatML CIM-10 depuis le FHIR JSON.

Sources : smt_cim10_fhir.json (19 161 concepts)
Produit : data/processed/cim10_chatml.jsonl

Types d'exemples générés :
  1. code → description (lookup)
  2. description → code (codage)
  3. discrimination entre codes frères (même parent)
  4. inclusions/exclusions (ce qui est compris / exclu d'un code)
"""

import json
import random
from pathlib import Path

random.seed(42)

BASE = Path(__file__).resolve().parent.parent
RAW = BASE / "data" / "raw"
OUT = BASE / "data" / "processed"
OUT.mkdir(parents=True, exist_ok=True)

SYSTEM_MSG = "Tu es un médecin DIM expert en codage CIM-10 pour le PMSI français."

# --- Chapitres CIM-10 ---
CHAPTERS = {
    "A": "Certaines maladies infectieuses et parasitaires",
    "B": "Certaines maladies infectieuses et parasitaires",
    "C": "Tumeurs",
    "D": "Tumeurs / Maladies du sang",
    "E": "Maladies endocriniennes, nutritionnelles et métaboliques",
    "F": "Troubles mentaux et du comportement",
    "G": "Maladies du système nerveux",
    "H": "Maladies de l'œil et de l'oreille",
    "I": "Maladies de l'appareil circulatoire",
    "J": "Maladies de l'appareil respiratoire",
    "K": "Maladies de l'appareil digestif",
    "L": "Maladies de la peau et du tissu cellulaire sous-cutané",
    "M": "Maladies du système ostéo-articulaire",
    "N": "Maladies de l'appareil génito-urinaire",
    "O": "Grossesse, accouchement et puerpéralité",
    "P": "Certaines affections dont l'origine se situe dans la période périnatale",
    "Q": "Malformations congénitales et anomalies chromosomiques",
    "R": "Symptômes, signes et résultats anormaux",
    "S": "Lésions traumatiques et empoisonnements",
    "T": "Lésions traumatiques et empoisonnements",
    "V": "Causes externes de morbidité et de mortalité",
    "W": "Causes externes de morbidité et de mortalité",
    "X": "Causes externes de morbidité et de mortalité",
    "Y": "Causes externes de morbidité et de mortalité",
    "Z": "Facteurs influant sur l'état de santé et motifs de recours aux services de santé",
    "U": "Codes d'utilisation particulière",
}


def load_fhir():
    """Charger et indexer les concepts FHIR."""
    with open(RAW / "smt_cim10_fhir.json") as f:
        data = json.load(f)
    concepts = data["concept"]

    # Index par code
    by_code = {}
    for c in concepts:
        by_code[c["code"]] = c

    return concepts, by_code


def get_props(concept):
    """Extraire les propriétés d'un concept sous forme de dict (multi-valeurs en listes)."""
    props = {}
    for p in concept.get("property", []):
        key = p["code"]
        val = p.get("valueCode", p.get("valueString", ""))
        if key in props:
            if isinstance(props[key], list):
                props[key].append(val)
            else:
                props[key] = [props[key], val]
        else:
            props[key] = val
    return props


def clean_display(display):
    """Nettoyer le libellé (enlever les codes entre crochets type [G31.0])."""
    import re
    # Retirer les références entre crochets comme [G31.0†]
    cleaned = re.sub(r'\s*\[[\w.†*+]+\]\s*', ' ', display)
    # Retirer les guillemets décoratifs
    cleaned = cleaned.replace('"', '').replace('"', '').replace('"', '')
    # Nettoyer les espaces multiples et les tirets isolés
    cleaned = re.sub(r'\s*-\s*-\s*', ' - ', cleaned)
    cleaned = re.sub(r'\s+', ' ', cleaned).strip()
    return cleaned


def make_chatml(system, user, assistant):
    """Créer un exemple ChatML."""
    return {
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ]
    }


def generate_lookup_pairs(concepts, by_code):
    """Type 1 : code → description."""
    pairs = []
    for c in concepts:
        props = get_props(c)
        if props.get("type") not in ("category",):
            continue

        code = c["code"]
        display = clean_display(c["display"])

        if not display or len(display) < 3:
            continue

        chapter = CHAPTERS.get(code[0], "")
        parent_code = props.get("parent", "")
        parent_display = ""
        if parent_code and parent_code in by_code:
            parent_display = clean_display(by_code[parent_code]["display"])

        # Construire la réponse enrichie
        answer_parts = [f"{code} — {display}"]
        if chapter:
            answer_parts.append(f"Chapitre : {chapter}")
        if parent_display and parent_code != code:
            answer_parts.append(f"Catégorie parente : {parent_code} ({parent_display})")

        # Ajouter inclusions si présentes
        incl = props.get("inclusionNote", "")
        if incl and len(incl) < 300:
            answer_parts.append(f"Comprend : {incl}")

        # Varier les formulations de question
        templates = [
            f"Que désigne le code CIM-10 {code} ?",
            f"Quel est le libellé du code {code} ?",
            f"Décris le code CIM-10 {code}.",
        ]
        question = random.choice(templates)

        pairs.append(make_chatml(SYSTEM_MSG, question, "\n".join(answer_parts)))

    return pairs


def generate_coding_pairs(concepts, by_code):
    """Type 2 : description → code (codage direct)."""
    pairs = []
    for c in concepts:
        props = get_props(c)
        if props.get("type") not in ("category",):
            continue

        code = c["code"]
        display = clean_display(c["display"])

        if not display or len(display) < 5:
            continue

        # Réponse JSON structurée (format du pipeline T2A)
        answer = json.dumps({
            "code": code,
            "confidence": "high",
            "justification": f"Correspondance directe avec le libellé CIM-10 : {code} {display}."
        }, ensure_ascii=False)

        templates = [
            f"Quel est le code CIM-10 pour : {display} ?",
            f"Code CIM-10 pour « {display} » ?",
            f"Codage CIM-10 du diagnostic : {display}",
        ]
        question = random.choice(templates)

        pairs.append(make_chatml(SYSTEM_MSG, question, answer))

    return pairs


def generate_discrimination_pairs(concepts, by_code):
    """Type 3 : discrimination entre codes frères (même parent)."""
    pairs = []

    # Grouper par parent
    children_by_parent = {}
    for c in concepts:
        props = get_props(c)
        if props.get("type") != "category":
            continue
        parent = props.get("parent", "")
        if parent and parent in by_code:
            children_by_parent.setdefault(parent, []).append(c)

    for parent_code, children in children_by_parent.items():
        if len(children) < 2 or len(children) > 15:
            continue

        parent = by_code[parent_code]
        parent_display = clean_display(parent["display"])

        # Construire la question
        question = f"Quels sont les sous-codes de {parent_code} ({parent_display}) et comment les distinguer ?"

        # Construire la réponse
        lines = [f"La catégorie {parent_code} ({parent_display}) comprend les codes suivants :\n"]
        for child in children:
            child_display = clean_display(child["display"])
            child_props = get_props(child)
            line = f"- {child['code']} : {child_display}"
            # Ajouter une note d'inclusion courte si disponible
            incl = child_props.get("inclusionNote", "")
            if incl and len(incl) < 150:
                line += f" (comprend : {incl})"
            lines.append(line)

        lines.append(f"\nLe choix du code dépend de la précision diagnostique disponible. "
                      f"En l'absence de précision, utiliser le code SAI (.9) s'il existe.")

        answer = "\n".join(lines)

        # Limiter la taille
        if len(answer) > 2000:
            continue

        pairs.append(make_chatml(SYSTEM_MSG, question, answer))

    return pairs


def generate_inclusion_exclusion_pairs(concepts, by_code):
    """Type 4 : questions sur les inclusions/exclusions d'un code."""
    pairs = []
    for c in concepts:
        props = get_props(c)
        if props.get("type") not in ("category",):
            continue

        code = c["code"]
        display = clean_display(c["display"])
        incl = props.get("inclusionNote", "")
        excl_note = props.get("exclusionNote", "")
        excl_codes = props.get("exclusion", "")
        note = props.get("note", "")

        # Il faut au moins une inclusion OU exclusion
        if not incl and not excl_note:
            continue

        # Construire la réponse
        answer_parts = [f"Code {code} — {display}\n"]

        if incl:
            answer_parts.append(f"Ce code COMPREND :\n{incl}")

        if excl_note:
            answer_parts.append(f"\nCe code EXCLUT :\n{excl_note}")

        if note:
            answer_parts.append(f"\nNote : {note}")

        answer = "\n".join(answer_parts)
        if len(answer) > 2000:
            continue

        templates = [
            f"Quelles sont les inclusions et exclusions du code {code} ({display}) ?",
            f"Que comprend et que exclut le code CIM-10 {code} ?",
        ]
        question = random.choice(templates)

        pairs.append(make_chatml(SYSTEM_MSG, question, answer))

    return pairs


def main():
    print("Chargement du FHIR JSON...")
    concepts, by_code = load_fhir()
    print(f"  {len(concepts)} concepts chargés")

    print("\nGénération des paires...")

    print("  Type 1 : code → description (lookup)")
    lookup = generate_lookup_pairs(concepts, by_code)
    print(f"    → {len(lookup)} exemples")

    print("  Type 2 : description → code (codage)")
    coding = generate_coding_pairs(concepts, by_code)
    print(f"    → {len(coding)} exemples")

    print("  Type 3 : discrimination codes frères")
    discrim = generate_discrimination_pairs(concepts, by_code)
    print(f"    → {len(discrim)} exemples")

    print("  Type 4 : inclusions / exclusions")
    incl_excl = generate_inclusion_exclusion_pairs(concepts, by_code)
    print(f"    → {len(incl_excl)} exemples")

    # Fusionner et mélanger
    all_pairs = lookup + coding + discrim + incl_excl
    random.shuffle(all_pairs)

    # Écrire en JSONL
    output_path = OUT / "cim10_chatml.jsonl"
    with open(output_path, "w") as f:
        for pair in all_pairs:
            f.write(json.dumps(pair, ensure_ascii=False) + "\n")

    print(f"\nTotal : {len(all_pairs)} exemples → {output_path}")
    print(f"Taille : {output_path.stat().st_size / 1024 / 1024:.1f} Mo")

    # Stats par type
    print("\nRépartition :")
    print(f"  Lookup (code→desc)  : {len(lookup)}")
    print(f"  Codage (desc→code)  : {len(coding)}")
    print(f"  Discrimination      : {len(discrim)}")
    print(f"  Inclusions/Exclus.  : {len(incl_excl)}")


if __name__ == "__main__":
    main()