import gradio as gr
import spacy
from graphviz import Digraph
import io
from PIL import Image
import pandas as pd

# Cargar modelo de spaCy
nlp = spacy.load("es_dep_news_trf")

def generar_grafico_dependencia(texto):
    if not texto or texto.strip() == "":
        return [], pd.DataFrame()

    doc = nlp(texto)
    raices = [token for token in doc if token.head == token]
    imagenes = []

    for i, raiz in enumerate(raices):
        dot = Digraph(comment=f"Árbol {i+1}")
        dot.attr('node', shape='ellipse')

        def agregar_nodo(palabra):
            nodo_id = f"{palabra.text}_{palabra.i}"
            etiqueta = f"{palabra.text}\n({palabra.pos_})"
            dot.node(nodo_id, etiqueta)
            for hijo in palabra.children:
                hijo_id = f"{hijo.text}_{hijo.i}"
                etiqueta_hijo = f"{hijo.text}\n({hijo.pos_})"
                dot.node(hijo_id, etiqueta_hijo)
                dot.edge(nodo_id, hijo_id, label=hijo.dep_)
                agregar_nodo(hijo)

        agregar_nodo(raiz)

        # Generar imagen en memoria (bytes)
        png_bytes = dot.pipe(format='png')
        
        # Convertir bytes a imagen PIL
        imagen = Image.open(io.BytesIO(png_bytes))
        imagenes.append(imagen)

    # Crear tabla de atributos
    df = pd.DataFrame([{
        "idx": token.i,
        "text": token.text,
        "lemma_": token.lemma_,
        "pos_": token.pos_,
        "tag_": token.tag_,
        "dep_": token.dep_,
        "head": token.head.text,
        "morph": str(token.morph),
        "ent_type_": token.ent_type_,
        "ent_iob_": token.ent_iob_,
        "shape_": token.shape_,
        "is_alpha": token.is_alpha,
        "is_ascii": token.is_ascii,
        "is_digit": token.is_digit,
        "is_punct": token.is_punct,
        "like_num": token.like_num,
        "is_sent_start": str(token.is_sent_start) if token.is_sent_start is not None else "None"
    } for token in doc])

    return imagenes, df

# Ejemplos
ejemplos = [
    ["El gato duerme en el sofá."],
    ["María compró un libro en la librería."],
    ["Aunque llueva, iremos al parque."],
    ["Pedro dijo que vendría mañana."],
    ["Los niños que juegan en el parque están felices."]
]

CSS = """
.contain { display: flex; flex-direction: column; }
.gallery-container { height: calc(100vh - 250px) !important; }
#component-0 { height: 100%; }
#gallery { flex-grow: 1; overflow: auto;}
"""

# Interfaz con Blocks
with gr.Blocks(title="Visualización de Dependencias Sintácticas", theme=gr.themes.Ocean(), css=CSS) as demo:
    gr.Markdown("# 🌐 Visualización de Dependencias Sintácticas en Español")
    gr.Markdown("Introduce un texto en español para ver el árbol de dependencias generado con spaCy y Graphviz.")

    with gr.Row():
        with gr.Column(scale=1):
            texto_input = gr.Textbox(lines=4, label="Texto en español", placeholder="Escribe aquí tu frase...")
            boton = gr.Button("Generar gráfico")
            gr.Examples(
                examples=ejemplos,
                inputs=texto_input,
                fn=generar_grafico_dependencia,
                label="Ejemplos de texto"
            )

        with gr.Column(scale=2):
            with gr.Tab("visualización"):
                galeria = gr.Gallery(label="Gráfico(s) generado(s)", show_label=True, columns=4, height="auto", elem_id="gallery")
            with gr.Tab("tabla"):
                tabla = gr.Dataframe(
                    headers=[
                        "idx", "text", "lemma_", "pos_", "tag_", "dep_", "head", "morph",
                        "ent_type_", "ent_iob_", "shape_", "is_alpha", "is_ascii",
                        "is_digit", "is_punct", "like_num", "is_sent_start"
                    ],
                    label="Atributos de tokens",
                    interactive=False,
                    type="pandas",
                    row_count=5,
                    max_height="600"
                )

    boton.click(fn=generar_grafico_dependencia, inputs=texto_input, outputs=[galeria, tabla])

# Lanzar app
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)