#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
cargar_modelo.py
────────────────
Ejemplo de uso completo de LexiMus-BETO-per-v1.

Carga el modelo NER transformer + Entity Ruler desde CSV y aplica
la inferencia sobre textos de prensa musical española histórica.

Requisitos:
    pip install spacy>=3.8,<3.9 spacy-transformers

Uso:
    python cargar_modelo.py
    python cargar_modelo.py --texto "Anoche actuó la Orquesta Nacional"
    python cargar_modelo.py --archivo mis_textos.txt
"""

import csv
import argparse
from pathlib import Path

import spacy

# ── Configuración ──────────────────────────────────────────────────────────────

ETIQUETAS = {"COMPOSITOR", "INTERPRETE", "CANTANTE", "AGRUPACION"}

COLORES_TERMINAL = {
    "COMPOSITOR": "\033[94m",   # azul
    "INTERPRETE": "\033[93m",   # amarillo
    "CANTANTE":   "\033[92m",   # verde
    "AGRUPACION": "\033[91m",   # rojo
    "RESET":      "\033[0m",
}

EJEMPLOS = [
    "El maestro Arbós dirigió anoche la Orquesta Sinfónica de Madrid interpretando la Sinfonía en re menor de Beethoven.",
    "La soprano Conchita Supervía cosechó grandes ovaciones en el Teatro Real junto al tenor Miguel Fleta.",
    "El Cuarteto Francés, formado por los profesores Francés, Outumuro, Del Campo y Fontova, ofreció una sesión de música de cámara.",
    "Wagner ocupa un lugar central en el programa de la temporada, con la colaboración del Orfeón Pamplonés.",
    "Falla presentó su nuevo ballet ante un público entregado; la coreografía corrió a cargo de la compañía de Diaghilev.",
]

# ── Funciones ──────────────────────────────────────────────────────────────────

def cargar_patrones_csv(ruta_csv: Path) -> list:
    """Lee el CSV del Entity Ruler y devuelve lista de patrones spaCy."""
    patrones = []
    vistos = set()
    with open(ruta_csv, encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for fila in reader:
            etq = fila["etiqueta"]
            if etq not in ETIQUETAS:
                continue
            textos_patron = {fila["texto"].strip()}
            if fila.get("variante_limpia", "").strip():
                textos_patron.add(fila["variante_limpia"].strip())
            notas = fila.get("notas", "")
            if "Apodo:" in notas:
                apodo = notas.split("Apodo:")[-1].split("|")[0].strip()
                if apodo:
                    textos_patron.add(apodo)
            for t in textos_patron:
                if not t or len(t) < 2:
                    continue
                clave = (t.lower(), etq)
                if clave in vistos:
                    continue
                vistos.add(clave)
                patrones.append({"label": etq, "pattern": t})
    return patrones


def cargar_leximus(ruta_modelo: str | Path, ruta_csv: str | Path) -> spacy.language.Language:
    """
    Carga el modelo LexiMus con Entity Ruler integrado.

    Args:
        ruta_modelo: Ruta al directorio del modelo spaCy (model-best/).
        ruta_csv:    Ruta al archivo entidades_ner_leximus.csv.

    Returns:
        nlp: Pipeline spaCy listo para inferencia.
    """
    ruta_modelo = Path(ruta_modelo)
    ruta_csv    = Path(ruta_csv)

    if not ruta_modelo.exists():
        raise FileNotFoundError(f"Modelo no encontrado: {ruta_modelo}")
    if not ruta_csv.exists():
        raise FileNotFoundError(f"CSV de patrones no encontrado: {ruta_csv}")

    print(f"Cargando modelo desde {ruta_modelo}...")
    nlp = spacy.load(ruta_modelo)

    # Eliminar Entity Ruler previo (puede tener patrones desactualizados)
    if "entity_ruler" in nlp.component_names:
        nlp.remove_pipe("entity_ruler")

    patrones = cargar_patrones_csv(ruta_csv)
    ruler = nlp.add_pipe("entity_ruler", last=True, config={"overwrite_ents": False})
    ruler.add_patterns(patrones)
    print(f"Sistema listo: NER transformer + Entity Ruler ({len(ruler)} patrones)")
    return nlp


def analizar_texto(nlp: spacy.language.Language, texto: str, color: bool = True) -> list:
    """
    Analiza un texto y devuelve las entidades musicales detectadas.

    Args:
        nlp:    Pipeline cargado con cargar_leximus().
        texto:  Texto a analizar.
        color:  Si True, imprime con colores ANSI en terminal.

    Returns:
        Lista de dicts con keys: texto, etiqueta, inicio, fin.
    """
    doc = nlp(texto)
    entidades = [
        {"texto": e.text, "etiqueta": e.label_, "inicio": e.start_char, "fin": e.end_char}
        for e in doc.ents if e.label_ in ETIQUETAS
    ]

    if color:
        print(f"\nTexto: {texto[:100]}{'...' if len(texto) > 100 else ''}")
        if entidades:
            for e in entidades:
                c = COLORES_TERMINAL.get(e["etiqueta"], "")
                r = COLORES_TERMINAL["RESET"]
                print(f"  {c}[{e['etiqueta']}]{r} {e['texto']}")
        else:
            print("  (sin entidades detectadas)")

    return entidades


def analizar_archivo(nlp: spacy.language.Language, ruta: str) -> None:
    """Analiza cada línea no vacía de un archivo de texto."""
    ruta = Path(ruta)
    lineas = [l.strip() for l in ruta.read_text(encoding="utf-8").splitlines() if l.strip()]
    print(f"\nAnalizando {len(lineas)} textos de {ruta.name}...")
    total_ents = 0
    for linea in lineas:
        ents = analizar_texto(nlp, linea)
        total_ents += len(ents)
    print(f"\nTotal: {total_ents} entidades detectadas en {len(lineas)} textos.")


# ── Main ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="LexiMus NER - inferencia sobre texto musical español")
    parser.add_argument("--modelo",  default=".", help="Ruta al directorio del modelo (por defecto: directorio actual)")
    parser.add_argument("--csv",     default="entidades_ner_leximus.csv", help="Ruta al CSV del Entity Ruler")
    parser.add_argument("--texto",   help="Texto a analizar directamente")
    parser.add_argument("--archivo", help="Archivo .txt con un texto por línea")
    args = parser.parse_args()

    nlp = cargar_leximus(args.modelo, args.csv)

    if args.texto:
        analizar_texto(nlp, args.texto)
    elif args.archivo:
        analizar_archivo(nlp, args.archivo)
    else:
        # Ejecutar ejemplos por defecto
        print("\n── Ejemplos de inferencia ─────────────────────────────────────────")
        for ejemplo in EJEMPLOS:
            analizar_texto(nlp, ejemplo)


if __name__ == "__main__":
    main()