import gradio as gr
import sys
import html

sys.path.append("C4_200M-synthetic-dataset-for-grammatical-error-correction/multilingual")
from merrant import api

TITLE = "Errant Multilingual-es"

DESCRIPTION = """
**Demo interactiva basada en la herramienta de anotación multilingüe desarrollada por Felix Stahlberg y Shankar Kumar**
Artículo original: Synthetic Data Generation for Low-resource Grammatical Error Correction with Tagged Corruption Models (https://aclanthology.org/2024.bea-1.2/)

Repositorio de la herramienta de anotacion:
https://github.com/google-research-datasets/C4_200M-synthetic-dataset-for-grammatical-error-correction/tree/main

**Funcionamiento:**

La herramienta recibe dos oraciones: **Texto original** (contiene errores) y el **texto corregido** (versión correcta).

A partir de ambas, la herramienta detecta los cambios y **genera una anotación automática** indicando: qué palabra fue corregida, cuál es la corrección aplicada y el tipo de error (por ejemplo: R:DET reemplazo de determinante)

_Las etiquetas de error se basan en las categorías gramaticales (POS) definidas bajo el estándar Universal Dependencies (UD)._
"""

EXAMPLES = [
    ["Pepito jugar en el parque", "Pepito juega en el parque"],
    ["muchos perros juegan en la parque tambien muchos niño salen. de la escuels","Muchos perros juegan en el parque, tambien muchos niños salen de la escuela"],
    ["Las gata pasear por el jardin bello", "La gata pasea por el bello jardin"]
]

annotator = api.Annotator("es_core_news_sm", aspell_lang="es")
annotator.initialize()

COLOR_R="#86E3CE"
COLOR_M="#FFDD94"
COLOR_U="#CCABD8"
COLOR_TRG="#D0E6A5"

CSS = f"""
.correction-box-R,
.correction-box-M,
.correction-box-U,
.tgt-word {{
    display: inline-flex;
    align-items: center;
    vertical-align: middle;
    border-radius: 6px;
    padding: 2px 6px;
    margin: 1px 2px;
}}

.correction-box-R {{
    background-color: {COLOR_R};
}}

.correction-box-M {{
    background-color: {COLOR_M};
}}

.correction-box-U {{
    background-color: {COLOR_U};
}}

.tgt-word {{
    background-color: {COLOR_TRG};
}}

/* Palabra original tachada */
.src-word {{
    padding: 1px 4px;
    text-decoration: line-through;
    margin-right: 4px;
    opacity: 0.8;
    color:#686868;
    font-weight:bold;
}}

/* Etiqueta tipo de error */
.label {{
    background-color: white;
    padding: 0px 4px;
    border-radius: 4px;
    margin-left: 4px;
    font-size: 0.75em;
    color: #686868;
    font-weight: 500;
}}
"""

CSS += """
/* Contenedor principal al 70% del ancho */
.gradio-container {
    width: 70% !important;
    margin: auto;
}

/* Título centrado y grande */
.gradio-title {
    text-align: center !important;
    font-size: 2.5em !important;
    font-weight: bold;
    margin-bottom: 20px;
}

/* Subtítulos Markdown centrados */
.gradio-markdown {
    text-align: center;
}

/* Ajuste de filas y spacing */
.gr-row {
    margin-bottom: 15px;
    justify-content: center; /* centra los elementos de la fila */
}

/* TextBoxes uniformes */
.gr-textbox {
    width: 48% !important; /* dos cajas lado a lado */
}

/* Botones centrados y uniformes */
.gr-button {
    min-width: 150px;
    margin: 0 5px;
}

/* HTML output centrado */
.gr-html {
    width: 100%;
    text-align: left;
    margin-top: 10px;
}

/* Leyenda de colores */
.leyenda {
    margin-top: 10px;
    font-size: 0.9em;
    text-align: center;
}
"""


LEYEND=f"""
<div style="margin: 8px 0 18px; display:flex; gap:10px; font-size:14px; align-items:center; flex-wrap:wrap;">

    <div style="display:flex; align-items:center; gap:6px;">
        <span style="width:16px; height:16px; background:{COLOR_R}; border-radius:4px; display:inline-block;"></span>
        <span>Replacement (R)</span>
    </div>

    <div style="display:flex; align-items:center; gap:6px;">
        <span style="width:16px; height:16px; background:{COLOR_M}; border-radius:4px; display:inline-block;"></span>
        <span>Missing (M)</span>
    </div>

    <div style="display:flex; align-items:center; gap:6px;">
        <span style="width:16px; height:16px; background:{COLOR_U}; border-radius:4px; display:inline-block;"></span>
        <span>Unnecessary (U)</span>
    </div>

    <div style="display:flex; align-items:center; gap:6px;">
        <span style="width:16px; height:16px; background:{COLOR_TRG}; border-radius:4px; display:inline-block;"></span>
        <span>Correct</span>
    </div>

</div>
"""
def get_annotation(source_text, corrected_text):
    """Obtiene la anotación completa desde Merrant"""
    annotation = annotator.annotate(source_text, [corrected_text])
    return annotation


def draw_corrected(error):
    html_out = f"<span class='tgt-word'>{html.escape(error['target'])}</span>"
    return html_out

def draw_error(error):
     html_out = f"""
        <span class='correction-box-{error['label']}'>
            <span class='src-word'>{html.escape(error['source'])}</span>
            <span class='label'>{html.escape(error['tag'])}</span>
        </span>
        """
     return html_out

def draw_merge(error):
    return draw_error(error)+draw_corrected(error)

def render_with_draw(annotation, draw_fn):
  
    doc = annotation.target_sentences[0].doc
    tokens = [token.text for token in doc]
    spaces = [token.whitespace_ for token in doc]
    errors = get_errors(annotation)

    for error in reversed(errors):
        start = error['tgt_start']
        end = error['tgt_end']

        if start == end:
            if 0 <= start <= len(tokens):
                tokens.insert(start, draw_fn(error))
                spaces.insert(start, spaces[start] if start < len(spaces) else "")
        else:
            tokens[start] = draw_fn(error)
            if start + 1 < end:
                del tokens[start+1:end]
                del spaces[start+1:end]

    return "".join([t + s for t, s in zip(tokens, spaces)])

def render_corrected(annotation):
    return render_with_draw(annotation, draw_corrected)

def render_original(annotation):
    return render_with_draw(annotation, draw_error)

def render_merge(annotation):
    return render_with_draw(annotation, draw_merge)

def show_annotation(source_doc, target_sentence, type_draw):
    annotation=get_annotation(source_doc, target_sentence)
    if type_draw == 'corrected':
        output = render_corrected(annotation)
    elif type_draw == 'original':
        output = render_original(annotation)
    else:
        output = render_merge(annotation)

    edits = format_edits(annotation)
    return output, edits

def format_edits(annotation):
    edits =get_errors(annotation)
    lines = ["<h3>Ediciones</h3>"]
    for e in edits: 
        src= e["source"]
        tgt = e["target"]
        tgt_start = e["tgt_start"]
        tgt_end = e["tgt_end"]
        src_start = e["src_start"]
        src_end = e["src_end"]
        label=e["label"]
        tag = e["tag"]

        line = f"""
                <span style="color:#686868;">ORIGINAL:</span>'{src}'({src_start}:{src_end}) → 
                <span style="color:#686868;">CORRECT</span>'{tgt}'({tgt_start}:{tgt_end}) | 
                <span style="color:#686868;">LABEL:</span> {label}:{tag}
                """
        lines.append(line)

    return "<br>".join(lines)

def get_errors(annotation):
    spans =[]
    for target in annotation.target_sentences:
        for edit in target.edit_spans:
            spans.append({
                "source":edit.source_span.text,
                "target":edit.target_span.text,
                "src_start":edit.source_span.start,
                "src_end":edit.source_span.end,
                "tgt_start":edit.target_span.start,
                "tgt_end":edit.target_span.end,
                "label":edit.get_label()[0],
                "tag":edit.tag,
            })
    return spans


def clear():
    return "", "", ""

with gr.Blocks(css=CSS) as demo:
    gr.Markdown(f"<h1 class='gradio-title'>{TITLE}</h1>") 
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        text = gr.Textbox(label="Texto Original", placeholder="Texto con errores...", lines=3)
        text_corrected = gr.Textbox(label="Texto Corregido", placeholder="Texto corregido...", lines=3)

    with gr.Row():
        btn_correct = gr.Button("Procesar", variant="primary")
        btn_clear = gr.Button("Limpiar", variant="secondary")
    type_draw= gr.Radio(choices=["original", "corrected", "merged"],
                                label="Tipo de visualización",
                                value="original"   
                               )
    gr.Markdown("### Errores detectados")
    errors_tags=gr.HTML(label="Errores detectados")
    gr.HTML(LEYEND)
    edits_html = gr.HTML(label="Ediciones detectadas")
   
    btn_correct.click(fn=show_annotation, inputs=[text, text_corrected, type_draw], outputs=[errors_tags, edits_html])
    btn_clear.click(fn=clear, outputs=[text, text_corrected, errors_tags])

   
    gr.Examples(
        examples= EXAMPLES,
        inputs=[text, text_corrected,type_draw],
        outputs=[errors_tags, edits_html],
        fn=show_annotation,
        cache_examples=False,
    )
demo.launch(server_name="0.0.0.0", server_port=7860)