import spaces  # Must precede torch for ZeroGPU.
import difflib
import logging
import gradio as gr
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration

MODEL_ID = "aimgo/Emendator"
MODEL_REVISION = "52040eff93121a80d63e3ee25a56714f578a21d2"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MAX_CHARACTERS = 250

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.info("Loading %s on %s", MODEL_ID, DEVICE)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=MODEL_REVISION)
model = T5ForConditionalGeneration.from_pretrained(
    MODEL_ID, revision=MODEL_REVISION, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True
).to(DEVICE).eval()
logger.info("Emendator loaded successfully.")


def highlight_changes(original: str, corrected: str) -> list[tuple[str, str | None]]:
    differences = []
    for operation, i1, i2, j1, j2 in difflib.SequenceMatcher(None, original, corrected, autojunk=False).get_opcodes():
        if operation == "equal":
            differences.append((original[i1:i2], None))
        else:
            if operation in ("delete", "replace"):
                differences.append((original[i1:i2], "Removed"))
            if operation in ("insert", "replace"):
                differences.append((corrected[j1:j2], "Added"))
    return differences


@spaces.GPU(duration=120)
def correct_latin(text: str) -> tuple[str, list[tuple[str, str | None]]]:
    if not text or not text.strip():
        raise gr.Error("Input a Latin passage or choose an example.")
    if len(text) > MAX_CHARACTERS:
        raise gr.Error(f"Your passage contains {len(text)} characters. Shorten it to 250 characters or fewer.")

    # ByT5 tokenizes UTF-8 bytes; avoid silently truncating accented text.
    inputs = tokenizer(text, return_tensors="pt", truncation=False).to(DEVICE)
    with torch.inference_mode():
        output = model.generate(
            **inputs, max_new_tokens=inputs["input_ids"].shape[1] + 32,
            num_beams=4, do_sample=False, early_stopping=True, repetition_penalty=1.15, use_cache=True
        )

    token_ids = output[0].detach().cpu().tolist()
    if tokenizer.eos_token_id not in token_ids:
        raise gr.Error("Generation reached its length limit. Try a shorter passage.")
    corrected = tokenizer.decode(token_ids, skip_special_tokens=True)
    if not corrected.strip():
        raise gr.Error("The model returned an empty result. Try another passage.")
    return corrected, highlight_changes(text, corrected)


EXAMPLES = [
    ["Nil igirur rnors cft ad nos ncq;pcrtinct hilurn»"],
    ["Vt quod ali cibus eft aliis fuat acre uenenurn."],
    ["atqu optimo viro, peterem, superavi tamen dignitate Catilinam, qatia Gdbam. Quod si id crimen homini novo esse deberet, protecto"],
    ["/Egidinslufitanusolufiponc natusS.Thcol.Mag.fcripfitCorn rncntariurninlibros>Scntcntiarurn«"],
    ["Diligamtc Do rnincfortitudomca:Dominusfirma rncntumrncurn&rcfugiurnmcurn&libcratorrnc9«"],
]

with gr.Blocks(title="Emendator — Latin OCR Correction") as demo:
    gr.Markdown(
        "# Latin OCR Correction\n"
        "Correct Latin OCR errors with [Emendator](https://huggingface.co/aimgo/Emendator).\n\n"
        "Use passages of 250 characters. Shorter or longer sequences lower performance."
        "Carefully review the output, as the model may change spelling, punctuation, or wording."
    )
    if DEVICE == "cpu":
        gr.Markdown("**Running on CPU.** Inference will be slow.")

    with gr.Row():
        source = gr.Textbox(label="Latin OCR input", placeholder="Paste your Latin OCR text here...", lines=7)
        result = gr.Textbox(label="Corrected Latin", lines=7, interactive=False)
    with gr.Row():
        submit = gr.Button("Correct Latin", variant="primary")
        clear = gr.ClearButton(value="Clear")

    differences = gr.HighlightedText(
        label="Changes", show_legend=True, combine_adjacent=True,
        color_map={"Removed": "#fecaca", "Added": "#bbf7d0"}
    )
    gr.Examples(examples=EXAMPLES, inputs=[source], cache_examples=False, label="Examples from the model card")
    submit.click(fn=correct_latin, inputs=[source], outputs=[result, differences], api_name="correct", concurrency_limit=1)
    clear.add([source, result, differences])
    gr.Markdown(
        "Model: Aidan McCarthy, *Emendator* (2026). Model and examples licensed under "
        "[CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/)."
    )

demo.queue(max_size=16, default_concurrency_limit=1)

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