smart_match.py

import os
import json
import re
import logging
from typing import Dict, List, Tuple, Optional
import numpy as np
from sentence_transformers import SentenceTransformer
#from langchain_openai import ChatOpenAI
#from langchain.prompts import ChatPromptTemplate
from dataclasses import dataclass
from template_db_creation import TemplateInfo, MedicalTemplateParser
from dotenv import load_dotenv
from openai import AzureOpenAI


# IMPORTANT: Configuration pour éviter les segfaults sur Mac
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'

# Configuration
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
load_dotenv()

# Azure OpenAI Configuration
AZURE_OPENAI_KEY = os.getenv("AZURE_OPENAI_KEY")
AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
AZURE_OPENAI_DEPLOYMENT = os.getenv("AZURE_OPENAI_DEPLOYMENT")
AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION", "2024-12-01-preview")
AZURE_OPENAI_MODEL = os.getenv("AZURE_OPENAI_MODEL", "gpt-5")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class MatchResult:
    """Résultat du matching"""
    template_id: str
    template_content: str
    similarity_score: float
    filled_template: str
    sections_filled: Dict[str, str]
    confidence_score: float

class TranscriptionMatcher:
    """Système de matching et remplissage automatique de templates"""
    
    def __init__(self, parser_instance):
        """
        Initialise le matcher avec une instance de MedicalTemplateParser
        
        Args:
            parser_instance: Instance de MedicalTemplateParser avec DB chargée
        """
        self.parser = parser_instance
        # Initialiser le modèle avec des paramètres Mac-friendly
        self.model = SentenceTransformer(
            EMBEDDING_MODEL, 
            device='cpu'
        )
        # Désactiver le multiprocessing pour éviter les segfaults
        self.model.encode = self._safe_encode
        self.llm = None
        self.filler = None
        self._initialize_gpt()
    
    def _safe_encode(self, *args, **kwargs):
        """Wrapper pour encode qui force show_progress_bar=False"""
        kwargs['show_progress_bar'] = False
        kwargs['batch_size'] = 1  # Traiter un par un pour plus de stabilité
        return SentenceTransformer.encode(self.model, *args, **kwargs)
    
    def _initialize_gpt(self):
        """Initialise GPT pour le remplissage des templates (Azure OpenAI)"""
        if not AZURE_OPENAI_KEY:
            logger.warning("AZURE_OPENAI_KEY non définie")
            return
        if not AZURE_OPENAI_ENDPOINT:
            logger.warning("AZURE_OPENAI_ENDPOINT non défini")
            return
        if not AZURE_OPENAI_DEPLOYMENT:
            logger.warning("AZURE_OPENAI_DEPLOYMENT non défini")
            return

        try:
            # Initialiser le client AzureOpenAI
            self.client = AzureOpenAI(
                api_key=AZURE_OPENAI_KEY,
                api_version=AZURE_OPENAI_API_VERSION,
                azure_endpoint=AZURE_OPENAI_ENDPOINT,
            )
            logger.info("✅ GPT initialisé avec AzureOpenAI client")
        except Exception as e:
            logger.error(f"❌ Erreur initialisation GPT: {e}")
            self.client = None

    
    def preprocess_transcription(self, transcription: str) -> str:
        """
        Prétraite la transcription pour améliorer le matching
        
        Args:
            transcription: Texte brut de la transcription
            
        Returns:
            str: Transcription nettoyée
        """
        # Nettoyer les espaces multiples
        text = re.sub(r'\s+', ' ', transcription)
        
        # Normaliser les points
        text = re.sub(r'\.{2,}', '.', text)
        
        # Assurer une majuscule après chaque point
        sentences = text.split('. ')
        text = '. '.join(s.capitalize() for s in sentences if s.strip())
        
        return text.strip()
    
    def extract_exam_type(self, transcription: str) -> str:
        """
        Extrait le type d'examen de la transcription
        
        Args:
            transcription: Texte de la transcription
            
        Returns:
            str: Type d'examen détecté
        """
        exam_patterns = [
            (r'\b(IRM|irm)\s+([a-zéèàçê\s]+)', 'IRM'),
            (r'\b(SCANNER|scanner|TDM|tdm)\s+([a-zéèàçê\s]+)', 'Scanner'),
            (r'\b(ECHOGRAPHIE|échographie|écho)\s+([a-zéèàçê\s]+)', 'Échographie'),
            (r'\b(RADIO|radio|radiographie)\s+([a-zéèàçê\s]+)', 'Radiographie'),
        ]
        
        text_lower = transcription.lower()
        for pattern, exam_type in exam_patterns:
            match = re.search(pattern, transcription, re.IGNORECASE)
            if match:
                return match.group(0).strip()
        
        # Fallback : chercher juste le type sans localisation
        for pattern, exam_type in exam_patterns:
            if exam_type.lower() in text_lower:
                return exam_type
        
        return "Examen médical"
    
    def create_search_query(self, transcription: str) -> str:
        """
        Crée une requête optimisée pour la recherche vectorielle
        
        Args:
            transcription: Texte de la transcription
            
        Returns:
            str: Requête optimisée
        """
        # Extraire les mots-clés importants
        exam_type = self.extract_exam_type(transcription)
        
        # Extraire les premières phrases pour le contexte
        sentences = transcription.split('.')[:3]
        context = '. '.join(sentences).strip()
        
        # Détecter les sections mentionnées
        sections = []
        section_keywords = {
            'technique': ['technique', 'acquisition', 'séquence'],
            'résultats': ['résultats', 'présence', 'mesure', 'aspect'],
            'conclusion': ['conclusion', 'diagnostic', 'aspect de']
        }
        
        text_lower = transcription.lower()
        for section, keywords in section_keywords.items():
            if any(kw in text_lower for kw in keywords):
                sections.append(section)
        
        # Construire la requête
        query_parts = [
            f"Type: {exam_type}",
            f"Sections: {', '.join(sections)}",
            f"Contexte: {context[:200]}"
        ]
        
        return ' | '.join(query_parts)
    
    def find_best_template(self, transcription: str, top_k: int = 3) -> List[Tuple[str, float]]:
        """
        Trouve les meilleurs templates correspondant à la transcription
        
        Args:
            transcription: Texte de la transcription
            top_k: Nombre de templates à retourner
            
        Returns:
            List[Tuple[str, float]]: Liste de (template_id, score)
        """
        logger.info("🔍 Recherche du meilleur template...")
        
        # Prétraiter la transcription
        clean_transcription = self.preprocess_transcription(transcription)
        
        # Créer la requête de recherche
        search_query = self.create_search_query(clean_transcription)
        logger.info(f"📝 Requête de recherche: {search_query[:100]}...")
        
        try:
            # Rechercher dans la base vectorielle avec gestion d'erreur
            results = self.parser.search_similar_templates(search_query, k=top_k)
            
            if results:
                logger.info(f"✅ {len(results)} templates trouvés")
                for i, (template_id, score) in enumerate(results, 1):
                    logger.info(f"   {i}. {template_id} (score: {score:.3f})")
            else:
                logger.warning("⚠️  Aucun template trouvé")
            
            return results
        except Exception as e:
            logger.error(f"❌ Erreur lors de la recherche: {e}")
            logger.info("🔄 Tentative avec recherche simplifiée...")
            return self._fallback_search(clean_transcription, top_k)
    
    def _fallback_search(self, transcription: str, top_k: int = 3) -> List[Tuple[str, float]]:
        """
        Méthode de recherche fallback en cas d'erreur
        
        Args:
            transcription: Texte de la transcription
            top_k: Nombre de résultats
            
        Returns:
            List[Tuple[str, float]]: Résultats de recherche
        """
        logger.info("⚠️ Utilisation de la recherche par mots-clés (fallback)")
        
        exam_type = self.extract_exam_type(transcription).lower()
        results = []
        
        # Recherche simple par mots-clés dans les IDs
        for template_id in self.parser.templates.keys():
            template_type = template_id.lower()
            score = 0.0
            
            # Correspondance exacte du type d'examen
            if exam_type in template_type or any(word in template_type for word in exam_type.split()):
                score = 0.8
                results.append((template_id, score))
        
        # Trier par score et retourner top_k
        results.sort(key=lambda x: x[1], reverse=True)
        return results[:top_k] if results else []
    
    def fill_template_with_gpt(self, template_content: str, transcription: str) -> Dict:
        """
        Remplit le template avec GPT (via Azure OpenAI)
        """
        if not self.client:
            logger.error("❌ GPT non disponible pour le remplissage")
            return self._fallback_fill(template_content, transcription)

        try:
            logger.info("🤖 Remplissage du template avec GPT...")

            system_prompt = """
    Vous êtes un expert médical spécialisé dans le formatage de comptes-rendus médicaux.

    Votre tâche est de remplir un template médical avec le contenu d'une transcription.

    **RÈGLES ABSOLUES DE FORMATAGE** :
1. JAMAIS commencer le contenu d'une section par son nom
2. JAMAIS écrire "TITRE:", "CLINIQUE:", "TECHNIQUE:", "RESULTATS:", "CONCLUSION:" dans le contenu
3. JAMAIS écrire "Titre.", "Clinique.", "Technique.", "Résultats.", "Conclusion." dans le contenu
4. Le contenu doit commencer DIRECTEMENT par le texte médical, sans label


**INSTRUCTIONS** :
1. Analysez la transcription pour identifier les différentes sections
2. Mappez chaque partie de la transcription aux sections du template
3. CONSERVER ABSOLUMENT la casse exacte des termes médicaux
4. Ne JAMAIS convertir en minuscules les acronymes médicaux
5. Si une section n'a pas de contenu, laissez-la vide

**RÈGLES DE MAPPING** :
- TITRE : première phrase ou mention de l'examen (IRM, Scanner, Échographie, etc.)
- CLINIQUE : informations sur l'indication clinique, motif d'examen
- TECHNIQUE : description de la technique d'acquisition, séquences, protocole
- RESULTATS : toutes les observations et mesures détaillées
- CONCLUSION : synthèse finale, diagnostic, recommandations

**FORMAT DE RÉPONSE** :
Retournez UNIQUEMENT un objet JSON avec cette structure exacte :
{{
    "sections": {{
        "TITRE": "contenu direct sans label",
        "CLINIQUE": "contenu direct sans label",
        "TECHNIQUE": "contenu direct sans label",
        "RESULTATS": "contenu direct sans label",
        "CONCLUSION": "contenu direct sans label"
    }},
    "confidence": 0.95
}}


    **IMPORTANT** : 
    - Répondez UNIQUEMENT avec du JSON valide
    - Pas de texte avant ou après le JSON
    - Pas de balises markdown
    - Conservez les sauts de ligne pour la lisibilité des RESULTATS
    """

            user_prompt = f"""
    Template à remplir :
    {template_content}

    Transcription médicale :
    {transcription}

    Remplissez le template en mappant correctement chaque section. Retournez UNIQUEMENT le JSON.
    """

            response = self.client.chat.completions.create(
                model=AZURE_OPENAI_DEPLOYMENT,  # ⚠️ deployment name, not model family
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                #temperature=0,
                #max_tokens=1000,
            )

            result = response.choices[0].message.content.strip()

            # Nettoyer le JSON
            if result.startswith("```json"):
                result = result[7:]
            if result.endswith("```"):
                result = result[:-3]
            result = result.strip()

            filled_data = json.loads(result)

            logger.info("✅ Template rempli avec succès")
            return filled_data

        except json.JSONDecodeError as e:
            logger.error(f"❌ Erreur parsing JSON: {e}")
            return self._fallback_fill(template_content, transcription)
        except Exception as e:
            logger.error(f"❌ Erreur remplissage GPT: {e}")
            return self._fallback_fill(template_content, transcription)

    
    def _fallback_fill(self, template_content: str, transcription: str) -> Dict:
        """
        Méthode de fallback pour remplir le template sans GPT
        
        Args:
            template_content: Template vide
            transcription: Transcription
            
        Returns:
            Dict: Sections remplies (basique)
        """
        logger.info("⚠️  Utilisation du remplissage basique (fallback)")
        
        sections = {
            "TITRE": self.extract_exam_type(transcription),
            "CLINIQUE": "",
            "TECHNIQUE": "",
            "RESULTATS": "",
            "CONCLUSION": ""
        }
        
        # Découpage basique par mots-clés
        text_lower = transcription.lower()
        
        # Technique
        if 'technique' in text_lower:
            tech_match = re.search(r'technique[:\.]?\s*([^\.]+(?:\.[^\.]+){0,3})', transcription, re.IGNORECASE)
            if tech_match:
                sections["TECHNIQUE"] = tech_match.group(1).strip()
        
        # Conclusion
        if 'conclusion' in text_lower:
            concl_match = re.search(r'conclusion[:\.]?\s*(.+)$', transcription, re.IGNORECASE | re.DOTALL)
            if concl_match:
                sections["CONCLUSION"] = concl_match.group(1).strip()
        
        # Résultats (tout ce qui reste)
        if not sections["TECHNIQUE"] and not sections["CONCLUSION"]:
            sections["RESULTATS"] = transcription
        
        return {
            "sections": sections,
            "confidence": 0.5
        }
    
    def format_filled_template(self, template_structure: str, filled_sections: Dict[str, str]) -> str:
        """
        Formate le template rempli
        
        Args:
            template_structure: Structure du template
            filled_sections: Sections remplies
            
        Returns:
            str: Template formaté
        """
        output_lines = []
        
        for section_name, content in filled_sections.items():
            if content and content.strip():
                # Nettoyer le contenu pour enlever toute duplication
                cleaned_content = self._clean_section_content(content, section_name)
                
                # Ne rien ajouter si le contenu est vide après nettoyage
                if not cleaned_content:
                    continue
                
                # Ajouter la section avec son nom
                output_lines.append(f"{section_name} :")
                output_lines.append(cleaned_content)
                output_lines.append("")  # Ligne vide entre sections
        
        return "\n".join(output_lines)



    def _clean_section_content(self, content: str, section_name: str) -> str:
        """
        Nettoie le contenu en enlevant les répétitions de labels
        
        Args:
            content: Contenu de la section
            section_name: Nom de la section (ex: "CONCLUSION", "TECHNIQUE")
            
        Returns:
            str: Contenu nettoyé
        """
        import re
        
        content = content.strip()
        
        # 1. Enlever "SECTION :" ou "SECTION:" au début (majuscules exactes)
        pattern1 = rf'^{re.escape(section_name)}\s*:\s*'
        content = re.sub(pattern1, '', content, flags=re.IGNORECASE)
        
        # 2. Enlever les variations avec/sans S (ex: RESULTAT/RESULTATS)
        section_variations = [
            section_name,
            section_name.rstrip('S'),  # RESULTATS → RESULTAT
        ]
        
        for variation in section_variations:
            # Enlever "Variation :" ou "Variation."
            pattern = rf'^{re.escape(variation)}\s*[.:]\s*'
            content = re.sub(pattern, '', content, flags=re.IGNORECASE)
        
        # 3. Cas spéciaux pour les variations françaises courantes
        special_cases = {
            'RESULTATS': ['résultat', 'résultats', 'resultat', 'resultats'],
            'TECHNIQUE': ['technique', 'techniques'],
            'CONCLUSION': ['conclusion', 'conclusions'],
            'CLINIQUE': ['clinique', 'cliniques', 'indication'],
        }
        
        if section_name in special_cases:
            for variation in special_cases[section_name]:
                pattern = rf'^{re.escape(variation)}\s*[.:]\s*'
                content = re.sub(pattern, '', content, flags=re.IGNORECASE)
        
        # 4. Nettoyer les espaces multiples résultants
        content = re.sub(r'\s+', ' ', content)
        
        return content.strip()
        
    def match_and_fill(self, transcription: str, return_top_k: int = 1) -> List[MatchResult]:
        """
        Fonction principale : trouve et remplit le(s) meilleur(s) template(s)
        
        Args:
            transcription: Texte de la transcription
            return_top_k: Nombre de templates à retourner
            
        Returns:
            List[MatchResult]: Liste des templates remplis
        """
        logger.info("\n" + "="*80)
        logger.info("🚀 DÉMARRAGE DU MATCHING ET REMPLISSAGE")
        logger.info("="*80)
        
        # Prétraiter
        clean_transcription = self.preprocess_transcription(transcription)
        logger.info(f"📝 Transcription prétraitée ({len(clean_transcription)} caractères)")
        
        # Trouver les meilleurs templates
        candidates = self.find_best_template(clean_transcription, top_k=return_top_k * 2)
        
        if not candidates:
            logger.error("❌ Aucun template trouvé")
            return []
        
        results = []
        
        for template_id, similarity_score in candidates[:return_top_k]:
            logger.info(f"\n{'─'*60}")
            logger.info(f"📋 Traitement du template: {template_id}")
            logger.info(f"{'─'*60}")
            
            # Récupérer le template
            template_info = self.parser.get_template_info(template_id)
            if not template_info:
                logger.warning(f"⚠️  Template {template_id} introuvable")
                continue
            
            logger.info(f"📄 Type: {template_info.type}")
            logger.info(f"📊 Score de similarité: {similarity_score:.3f}")
            
            # Remplir le template
            filled_data = self.fill_template_with_gpt(
                template_info.content,
                clean_transcription
            )
            
            sections_filled = filled_data.get("sections", {})
            confidence = filled_data.get("confidence", 0.0)
            
            # Formater le résultat
            filled_template = self.format_filled_template(
                template_info.content,
                sections_filled
            )
            
            result = MatchResult(
                template_id=template_id,
                template_content=template_info.content,
                similarity_score=similarity_score,
                filled_template=filled_template,
                sections_filled=sections_filled,
                confidence_score=confidence
            )
            
            results.append(result)
            
            logger.info(f"✅ Template rempli - Confiance: {confidence:.2%}")
        
        logger.info("\n" + "="*80)
        logger.info(f"✅ TRAITEMENT TERMINÉ - {len(results)} résultat(s)")
        logger.info("="*80 + "\n")
        
        return results
    
    def display_result(self, result: MatchResult):
        """
        Affiche un résultat de manière formatée
        
        Args:
            result: Résultat à afficher
        """
        print("\n" + "="*80)
        print(f"📋 RÉSULTAT - Template: {result.template_id}")
        print("="*80)
        print(f"📊 Score de similarité: {result.similarity_score:.3f}")
        print(f"🎯 Confiance: {result.confidence_score:.2%}")
        print("\n" + "─"*80)
        print("📄 TEMPLATE REMPLI:")
        print("─"*80)
        print(result.filled_template)
        print("="*80 + "\n")


def main():
    """Fonction principale pour tester le matching"""
    
    # Charger la base de données
    db_path = input("Chemin vers la base de données (.pkl): ").strip()
    
    if not os.path.exists(db_path):
        print(f"❌ Fichier {db_path} introuvable")
        return
    
    print("📂 Chargement de la base de données...")
    parser = MedicalTemplateParser()
    parser.load_database(db_path)
    
    print(f"✅ Base chargée: {len(parser.templates)} templates")
    
    # Initialiser le matcher
    matcher = TranscriptionMatcher(parser)
    
    # Exemple de transcription
    transcription_example = """
IRM pelvienne. Indication clinique. Technique. Acquisition sagittale, axiale et coronale T2, saturation axiale, diffusion axiale T1. Résultats. Présence d'un utérus antéversé médio-pelvien dont le grand axe mesure 72 mm sur 40 mm sur 40 mm. La zone jonctionnelle apparaît floue. Elle est épaissie de façon diffuse, asymétrique, avec une atteinte de plus de 50% de l'épaisseur du myomètre et comporte des spots en hypersignal T2, l'ensemble traduisant une adénomyose.
Pas d'épaississement cervical. À noter la présence d'un petit kyste liquidien de type Naboth.
Les 2 ovaires sont repérés, porteurs de formations folliculaires communes en hypersignal homogène T2 de petite taille. L'ovaire droit mesure 30 x 25 mm. L'ovaire gauche mesure 25 x 23 mm. Pas d'épanchement dans le cul-de-sac de Douglas.
Absence de foyer d'endométriose profonde. Conclusion.
Aspect d'adénomyose diffuse, symétrique, profonde.
Pas d'épaississement endométrial. Absence d'endométriome. Absence d'épanchement dans le cul-de-sac de Douglas.
"""
    
    # Mode interactif
    while True:
        print("\n" + "="*80)
        print("OPTIONS:")
        print("1. Utiliser l'exemple de transcription")
        print("2. Entrer une nouvelle transcription")
        print("3. Charger depuis un fichier")
        print("4. Quitter")
        print("="*80)
        
        choice = input("\nVotre choix: ").strip()
        
        if choice == "1":
            transcription = transcription_example
        elif choice == "2":
            print("\nEntrez la transcription (Ctrl+D ou Ctrl+Z pour terminer):")
            lines = []
            try:
                while True:
                    line = input()
                    lines.append(line)
            except EOFError:
                transcription = "\n".join(lines)
        elif choice == "3":
            filepath = input("Chemin du fichier: ").strip()
            try:
                with open(filepath, 'r', encoding='utf-8') as f:
                    transcription = f.read()
            except Exception as e:
                print(f"❌ Erreur lecture fichier: {e}")
                continue
        elif choice == "4":
            print("👋 Au revoir!")
            break
        else:
            print("❌ Choix invalide")
            continue
        
        # Effectuer le matching
        results = matcher.match_and_fill(transcription, return_top_k=3)
        
        if results:
            # Afficher tous les résultats
            for i, result in enumerate(results, 1):
                print(f"\n{'#'*80}")
                print(f"# RÉSULTAT {i}/{len(results)}")
                print(f"{'#'*80}")
                matcher.display_result(result)
            
            # Sauvegarder le meilleur résultat
            save_choice = input("\nSauvegarder le meilleur résultat? (o/n): ").strip().lower()
            if save_choice == 'o':
                output_file = input("Nom du fichier de sortie: ").strip()
                try:
                    with open(output_file, 'w', encoding='utf-8') as f:
                        f.write(results[0].filled_template)
                    print(f"✅ Résultat sauvegardé dans {output_file}")
                except Exception as e:
                    print(f"❌ Erreur sauvegarde: {e}")
        else:
            print("❌ Aucun résultat trouvé")


if __name__ == "__main__":
    main()