post_processing.py

import re
import logging
from typing import List, Tuple

logger = logging.getLogger(__name__)

class MedicalReportPostProcessor:
    """Post-traitement pour nettoyer les rapports médicaux mappés"""
    
    def __init__(self):
        # Patterns pour détecter les lignes avec choix multiples
        self.choice_patterns = [
            # Position utérus
            r'L\'utérus est(.*?)de taille',
            # Adénomyose
            r'Adénomyose associée :(.*?)(?:\n|$)',
            r'&x oui :(.*?)(?:\n|Col utérin)',
            # Accessibilité ovaires
            r'Accessibilité :(.*?)(?:\n|\t|$)',
            # Compartiments
            r'Signe du glissement \(sliding\) :(.*?)(?:\n|$)',
            # Nodules
            r'Présence d\'un nodule :(.*?)(?:\n|$)',
            r'Présence d\'un nodule hypoéchogène :(.*?)(?:\n|$)',
            # Aspect
            r'Aspect du torus\s*:(.*?)(?:\n|$)',
            r'Ligament utéro-.*?:(.*?)(?:\n|$)',
            # Infiltration
            r'Infiltration digestive:(.*?)(?:\n|$)',
            # Endométriose
            r'Endométriose(.*?)(?:\n|Absence)',
            # Épanchement
            r'- (Pas d\'épanchement.*?|Faible épanchement.*?)(?:\n|$)',
            # Vessie
            r'- (.*?Vessie.*?)(?:\n|$)',
            # Dilatation
            r'- (.*?dilatation.*?)(?:\n|$)',
            # Calcifications
            r'Présence de micro-calcifications(.*?)(?:\n|L\'échostructure)',
            # Ovaires dimensions supérieures
            r', (.*?est de dimensions supérieures.*?)(?:,|\n)',
            # Endométriome
            r'\. (.*?endométriome\.)(?:\n|$)',
        ]
    
    def process_report(self, report: str) -> str:
        """
        Traite le rapport complet pour nettoyer les choix multiples
        """
        logger.info("🧹 Début du post-traitement du rapport")
        
        processed_report = report
        
        # Étape 1: Nettoyer les lignes avec choix multiples
        processed_report = self._clean_multiple_choices(processed_report)
        
        # Étape 2: Nettoyer les placeholders isolés restants
        processed_report = self._clean_isolated_placeholders(processed_report)
        
        # Étape 3: Nettoyer les espaces et formatage
        processed_report = self._clean_formatting(processed_report)
        
        logger.info("✅ Post-traitement terminé")
        
        return processed_report
    
    def _clean_multiple_choices(self, text: str) -> str:
        """
        Nettoie les lignes contenant plusieurs choix (&x ou X)
        Ne garde que les options cochées (X)
        """
        lines = text.split('\n')
        cleaned_lines = []
        
        for line in lines:
            # Vérifier si la ligne contient des choix multiples
            if self._has_multiple_choices(line):
                cleaned_line = self._extract_checked_choices(line)
                cleaned_lines.append(cleaned_line)
            else:
                cleaned_lines.append(line)
        
        return '\n'.join(cleaned_lines)
    
    def _has_multiple_choices(self, line: str) -> bool:
        """
        Détecte si une ligne contient plusieurs choix (au moins 2 occurrences de &x ou X suivi d'un mot)
        """
        # Compter les patterns de choix: &x ou X suivi d'un mot
        pattern = r'(?:&x|(?<!\w)X(?=\s+\w))\s+\w+'
        matches = re.findall(pattern, line)
        return len(matches) >= 2
    
    def _extract_checked_choices(self, line: str) -> str:
        """
        Extrait uniquement les choix cochés (X) d'une ligne
        """
        # Séparer la partie avant les choix et après
        parts = self._split_line_by_choices(line)
        
        if not parts:
            return line
        
        prefix = parts['prefix']
        choices = parts['choices']
        suffix = parts['suffix']
        
        # Extraire les choix cochés
        checked_choices = []
        for choice in choices:
            if choice.strip().startswith('X '):
                # Enlever le X et garder le texte
                checked_text = choice.strip()[2:].strip()
                checked_choices.append(checked_text)
        
        # Reconstruire la ligne
        if checked_choices:
            result = prefix
            if len(checked_choices) == 1:
                result += checked_choices[0]
            else:
                result += ', '.join(checked_choices)
            result += suffix
            return result
        else:
            # Si aucun choix coché, retourner la ligne d'origine
            return line
    
    def _split_line_by_choices(self, line: str) -> dict:
        """
        Sépare une ligne en: préfixe, choix, suffixe
        """
        # Cas spécifiques avec patterns connus
        
        # Position utérus
        match = re.search(r'(L\'utérus est\s+)((?:[X&]x?\s+\w+[,\s]+)+)(de taille.*)', line)
        if match:
            return {
                'prefix': match.group(1),
                'choices': self._parse_choices(match.group(2)),
                'suffix': ' ' + match.group(3)
            }
        
        # Adénomyose associée
        match = re.search(r'(Adénomyose associée\s*:\s*)((?:[X&]x?\s+\w+\s*)+)(.*)', line)
        if match:
            return {
                'prefix': match.group(1),
                'choices': self._parse_choices(match.group(2)),
                'suffix': match.group(3)
            }
        
        # Type d'adénomyose
        match = re.search(r'([X&]x?\s+oui\s*:\s*)((?:[X&]x?\s+\w+\s*)+)(.*)', line)
        if match:
            return {
                'prefix': '',
                'choices': self._parse_choices(match.group(2)),
                'suffix': match.group(3)
            }
        
        # Accessibilité
        match = re.search(r'(Accessibilité\s*:\s*)((?:[X&]x?\s+[\w-]+\s*)+)(.*)', line)
        if match:
            return {
                'prefix': match.group(1),
                'choices': self._parse_choices(match.group(2)),
                'suffix': match.group(3)
            }
        
        # Signe du glissement
        match = re.search(r'(.*?Signe du glissement.*?:\s*)((?:[X&]x?\s*\w+\s*)+)(.*)', line)
        if match:
            return {
                'prefix': match.group(1),
                'choices': self._parse_choices(match.group(2)),
                'suffix': match.group(3)
            }
        
        # Présence nodule
        match = re.search(r'(.*?Présence d\'un nodule.*?:\s*)((?:[X&]x?\s*\w+\s*)+)(.*)', line)
        if match:
            return {
                'prefix': match.group(1),
                'choices': self._parse_choices(match.group(2)),
                'suffix': match.group(3)
            }
        
        # Aspect
        match = re.search(r'(.*?Aspect.*?:\s*)((?:[X&]x?\s+\w+\s*)+)(.*)', line)
        if match:
            return {
                'prefix': match.group(1),
                'choices': self._parse_choices(match.group(2)),
                'suffix': match.group(3)
            }
        
        # Infiltration digestive
        match = re.search(r'(.*?Infiltration digestive:\s*)((?:[X&]x?\s+\w+\s*)+:\s*)((?:[X&]x?\s+[\w\s-]+)+)(.*)', line)
        if match:
            # Gérer le cas spécial avec "non/oui :"
            first_choice = self._parse_choices(match.group(2))
            second_choices = self._parse_choices(match.group(3))
            return {
                'prefix': match.group(1),
                'choices': first_choice + second_choices,
                'suffix': match.group(4)
            }
        
        # Calcifications
        match = re.search(r'(.*?micro-calcifications.*?)((?:[X&]x?\s+\w+\s*)+)(.*)', line)
        if match:
            return {
                'prefix': match.group(1),
                'choices': self._parse_choices(match.group(2)),
                'suffix': match.group(3)
            }
        
        # Endométriose
        match = re.search(r'([X&]x?\s+Endométriose\s+)((?:[X&]x?\s+\w+\s*)+)(.*)', line)
        if match:
            return {
                'prefix': 'Endométriose ',
                'choices': self._parse_choices(match.group(2)),
                'suffix': match.group(3)
            }
        
        # Épanchement (début par -)
        match = re.search(r'^(\s*-?\s*)([X&]x?-?\s*(?:Pas|Faible).*?)$', line)
        if match:
            text = match.group(2)
            if text.strip().startswith('X'):
                return {
                    'prefix': match.group(1),
                    'choices': ['X ' + text[1:].strip()],
                    'suffix': ''
                }
        
        # Vessie
        match = re.search(r'^(\s*-\s*)([X&]x?\s*Vessie.*?)([X&]x?\s*Vessie.*?)$', line)
        if match:
            choices = []
            if match.group(2).strip().startswith('X'):
                choices.append('X ' + match.group(2)[1:].strip())
            if match.group(3).strip().startswith('&x'):
                pass  # Ne rien ajouter
            elif match.group(3).strip().startswith('X'):
                choices.append('X ' + match.group(3)[1:].strip())
            return {
                'prefix': match.group(1),
                'choices': choices,
                'suffix': ''
            }
        
        return None
    
    def _parse_choices(self, choices_text: str) -> List[str]:
        """
        Parse le texte des choix pour extraire chaque option
        """
        # Séparer par &x ou X en début de mot
        parts = re.split(r'(?=[X&]x?\s+)', choices_text)
        return [p.strip() for p in parts if p.strip()]
    
    def _clean_isolated_placeholders(self, text: str) -> str:
        """
        Nettoie les placeholders &x isolés qui restent
        """
        # Supprimer les &x en début de ligne ou après espace
        text = re.sub(r'^\s*&x\s*', '', text, flags=re.MULTILINE)
        text = re.sub(r'\s+&x\s+', ' ', text)
        text = re.sub(r'\s+&x$', '', text, flags=re.MULTILINE)
        
        # Nettoyer les lignes qui ne contiennent que des &x
        lines = text.split('\n')
        cleaned_lines = []
        for line in lines:
            # Si la ligne ne contient que des &x et espaces, la supprimer
            if re.match(r'^\s*(?:&x\s*)+$', line):
                continue
            cleaned_lines.append(line)
        
        return '\n'.join(cleaned_lines)
    
    def _clean_formatting(self, text: str) -> str:
        """
        Nettoie le formatage général
        """
        # Supprimer les espaces multiples
        text = re.sub(r' +', ' ', text)
        
        # Supprimer les lignes vides multiples
        text = re.sub(r'\n\s*\n\s*\n+', '\n\n', text)
        
        # Nettoyer les espaces avant ponctuation
        text = re.sub(r' +([,.])', r'\1', text)
        
        # Nettoyer les espaces après tirets en début de ligne
        text = re.sub(r'^(\s*-)\s+', r'\1 ', text, flags=re.MULTILINE)
        
        return text.strip()


def post_process_medical_report(filled_template: str) -> str:
    """
    Fonction principale pour post-traiter un rapport médical
    """
    processor = MedicalReportPostProcessor()
    return processor.process_report(filled_template)


# Exemple d'utilisation
if __name__ == "__main__":
    # Exemple de rapport avec choix multiples
    sample_report = """L'utérus est X antéversé, &x rétroversé, &x intermédiaire, &x rétrofléchi, &x antéfléchi, &x fixe de taille normale (7.8 x &x x &x cm).
Hystérométrie : distance orifice externe du col - fond de la cavité utérine : 60 mm.
L'endomètre : mesuré à 3.7 mm.
Myometre : pas de myome.
Zone jonctionnelle : Atteinte de la zone de jonction :  &x non &x oui 
			        Adénomyose associée : &x non  X oui : X diffuse &x focale  &x interne &x externe
Col utérin: pas de kyste de Naboth. Absence de pathologies échographiquement décelable à son niveau. 
Cavité utérine en 3D: morphologie triangulaire.

L'ovaire droit mesure 26 x 20 mm, &x est de dimensions supérieures à la normale il mesure &x x &x mm, &xfolliculaire CFA 15 follicules: (&x mm). &x Absence d'endométriome. &x Présence d'une formation kystique hypoéchogène, uniloculaire, non vascularisé, à contenu ground glass mesurée à &x mm d'allure endométriome.
Accessibilité : &x rétro-utérin &x fixe X aisée.	
L'ovaire gauche mesure 25 x 19 mm, &x est de dimensions supérieures à la normale il mesure &x x &x mm,  &x folliculaire CFA 22 follicules: (&x mm). &x Absence d'endométriome. &x Présence d'une formation kystique hypoéchogène, uniloculaire, non vascularisé, à contenu ground glass mesurée à &x mm d'allure endométriome.
Accessibilité : X rétro-utérin &x fixe &x aisée.	
&x Présence de micro-calcifications sous thécales &x bilatérales &x droites &x gauches pouvant témoigner d'implants endométriosiques  superficiels.
L'échostructure des deux ovaires apparait normale, avec une vascularisation artério-veineuse normale au Doppler, sans formation ou image kystique pathologique échographiquement décelable à leur niveau.

Cavité péritonéale
&x- Pas d'épanchement liquidien dans le cul du sac du Douglas. Pas de douleur à l'écho-palpation.
&x- Faible épanchement corpusculé dans le cul du sac du Douglas qui silhouette des adhérences (soft marqueur d'endométriose?). Pas de douleur à l'écho-palpation.
- XVessie vide pendant l'examen. &x Vessie en semi-réplétion pendant l'examen.
- X Absence de dilatation pyélo-calicielle.
- Artère utérine : IP : 3.24  - IR : 0,91 - Spectre : type 2 avec notch protodiastolique.
- Pas d'image d'hydrosalpinx visible à ce jour.

RECHERCHE ENDOMETRIOSE PELVIENNE

A-Compartiment antérieur (vessie en semi-réplétion)	
-	Signe du glissement (sliding) :  &xprésent &xdiminué &xabsent
-	Présence d'un nodule :	 &xnon &xoui	       
-    Uretères dans la partie pelvienne vus non dilatés.

B-Compartiment postérieur	
- Signe du glissement (sliding) : 
         - Espace recto-vaginal :  &xprésent &xdiminué &xabsent	 
         - Plan sus-péritonéal :   &xprésent &xdiminué &xabsent
- Aspect du torus	:  &x normal &x épaissi	 
- Aspect des ligaments utéro-sacrés :
        - Ligament utéro- sacré droit : &x normal 	&x épaissi
        - Ligament utéro-sacré gauche : &x normal  &x épaissi
- Présence d'un nodule hypoéchogène : &x non	    	
- Infiltration digestive: &x non   X oui : &x bas rectum    &x moyen rectum  &x haut rectum &x jonction recto-sigmoïde    	

Conclusions
Utérus de taille et de morphologie normales. 
Endomètre mesuré à 3.7 mm.
CFA : 15+22 follicules.
Ovaires sans formation ou image kystique pathologique échographiquement décelable à leur niveau.
X Absence d'image d'endométriose visible ce jour, à confronter éventuellement à une IRM.
&x Endométriose &x superficielle &x et profonde.
Absence d'anomalie échographiquement décelable au niveau des trompes.
--> L'ensemble de ces aspects reste à confronter au contexte clinico-thérapeutique."""

    # Appliquer le post-traitement
    cleaned_report = post_process_medical_report(sample_report)
    
    print("=" * 60)
    print("RAPPORT APRÈS POST-TRAITEMENT")
    print("=" * 60)
    print(cleaned_report)