import json

import gradio as gr
import numpy as np
from fastchrf import aggregate_chrf, pairwise_chrf


def parse_lines(text: str) -> list[str]:
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    if not lines:
        raise gr.Error("Provide at least one non-empty line.")
    return lines


def compute_chrf(
    mode: str,
    hypotheses_text: str,
    references_text: str,
    char_order: int,
    beta: float,
    remove_whitespace: bool,
    eps_smoothing: bool,
) -> str:
    hypotheses = parse_lines(hypotheses_text)
    references = parse_lines(references_text)

    if mode == "default":
        scores = pairwise_chrf(
            [hypotheses],
            [references],
            char_order=char_order,
            beta=beta,
            remove_whitespace=remove_whitespace,
            eps_smoothing=eps_smoothing,
        )
        array = np.array(scores[0])
        payload = {
            "mode": "default",
            "shape": list(array.shape),
            "scores": array.tolist(),
            "hypotheses": hypotheses,
            "references": references,
        }
    else:
        scores = aggregate_chrf(
            [hypotheses],
            [references],
            char_order=char_order,
            beta=beta,
            remove_whitespace=remove_whitespace,
            eps_smoothing=eps_smoothing,
        )
        array = np.array(scores[0])
        payload = {
            "mode": "aggregate",
            "shape": list(array.shape),
            "scores": array.tolist(),
            "hypotheses": hypotheses,
            "references": references,
        }

    return json.dumps(payload, indent=2, ensure_ascii=False)


EXAMPLE_HYPOTHESES = """The cat sat on the mat.
The cat sat on the hat."""

EXAMPLE_REFERENCES = """The cat sat on the mat.
The fat cat sat on the mat.
A cat sat on a mat."""

with gr.Blocks(title="fastChrF") as demo:
    gr.Markdown(
        """
        # fastChrF

        Fast sentence-level ChrF for Minimum Bayes Risk decoding.

        - **Default (pairwise):** ChrF between each hypothesis and each reference.
        - **Aggregate:** Streamlined variant that aggregates across references.

        [GitHub](https://github.com/jvamvas/fastChrF) · [Paper](https://arxiv.org/abs/2402.04251)
        """
    )

    with gr.Row():
        mode = gr.Radio(
            choices=["default", "aggregate"],
            value="default",
            label="Mode",
            info="Default (pairwise) returns a hypothesis × reference matrix; aggregate returns one score per hypothesis.",
        )

    with gr.Row():
        hypotheses = gr.Textbox(
            label="Hypotheses (one per line)",
            lines=8,
            placeholder="Enter one hypothesis per line",
            value=EXAMPLE_HYPOTHESES,
        )
        references = gr.Textbox(
            label="References (one per line)",
            lines=8,
            placeholder="Enter one reference per line",
            value=EXAMPLE_REFERENCES,
        )

    with gr.Accordion("Advanced options", open=False):
        with gr.Row():
            char_order = gr.Slider(
                minimum=1,
                maximum=10,
                value=6,
                step=1,
                label="Character n-gram order",
            )
            beta = gr.Slider(
                minimum=0.1,
                maximum=5.0,
                value=2.0,
                step=0.1,
                label="Beta (F-score weight)",
            )
        with gr.Row():
            remove_whitespace = gr.Checkbox(value=True, label="Remove whitespace")
            eps_smoothing = gr.Checkbox(value=False, label="Epsilon smoothing")

    compute_button = gr.Button("Compute ChrF", variant="primary")
    output = gr.Code(label="Results (JSON)", language="json", lines=20)

    compute_button.click(
        fn=compute_chrf,
        inputs=[
            mode,
            hypotheses,
            references,
            char_order,
            beta,
            remove_whitespace,
            eps_smoothing,
        ],
        outputs=output,
    )

    gr.Examples(
        examples=[
            [EXAMPLE_HYPOTHESES, EXAMPLE_REFERENCES, "default"],
            [EXAMPLE_HYPOTHESES, EXAMPLE_REFERENCES, "aggregate"],
        ],
        inputs=[hypotheses, references, mode],
        label="Examples",
    )

if __name__ == "__main__":
    demo.launch()