"""
Cuneiform Translator
Built for the Build Small Hackathon 2026.

Translate between English and the oldest writing systems on Earth.
4,000 years ago someone pressed a reed into wet clay to record a beer recipe.
Now you can read it.
"""

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import re

# --- Models ---
MODELS = {
    "thalesian": {
        "repo": "Thalesian/cuneiformBase-400m",
        "name": "cuneiformBase-400m",
    },
    "praeclarum": {
        "repo": "praeclarum/cuneiform",
        "name": "cuneiform-T5",
    },
}

# Load primary model
print("Loading Thalesian/cuneiformBase-400m...")
try:
    tokenizer = AutoTokenizer.from_pretrained(MODELS["thalesian"]["repo"], use_fast=False)
    model = AutoModelForSeq2SeqLM.from_pretrained(MODELS["thalesian"]["repo"])
    MODEL_LOADED = True
    print("Model loaded.")
except Exception as e:
    print(f"Model failed: {e}")
    MODEL_LOADED = False
    tokenizer = None
    model = None

# Unicode cuneiform lookup
CUNEIFORM_MAP = {
    "an": "𒀭", "ki": "𒆠", "en": "𒂗", "me": "𒈨", "nam": "𒉆",
    "lu": "𒇻", "mu": "𒈬", "du": "𒆕", "ga": "𒂵", "ma": "𒈠",
    "ni": "𒉌", "pa": "𒉺", "ra": "𒊏", "sa": "𒊓", "ta": "𒋫",
    "ti": "𒋾", "zu": "𒍪", "ba": "𒁀", "bi": "𒁉", "da": "𒁕",
    "di": "𒁲", "e": "𒂊", "gi": "𒄀", "gu": "𒄖", "ha": "𒄩",
    "hu": "𒄷", "li": "𒇷", "kal": "𒅗", "im": "𒅎", "ab": "𒀊",
    "ug": "𒌑", "lugal": "𒈗", "nin": "𒊩", "dam": "𒁮", "dumu": "𒌉",
    "ama": "𒂼", "dingir": "𒀭", "nanna": "𒀭𒂗𒍪", "inanna": "𒀭𒈹",
    "enlil": "𒀭𒂗𒍪", "enki": "𒀭𒂗𒆠", "utu": "𒀭𒌓", "ag": "𒀝",
    "nam-ti": "𒉆𒋾", "nam-lu": "𒉆𒇻", "ki-ag": "𒆠𒀝",
    "me-en": "𒈨𒂗", "za": "𒍝", "nun": "𒉣", "ur": "𒌨", "ub": "𒌒",
    "su": "𒍪", "szu": "𒋗", "ig": "𒅅", "ib": "𒅁",
    "a": "𒀀", "u": "𒌋", "i": "𒄿", "la": "𒆷", "na": "𒈾",
    "nu": "𒉡", "si": "𒋛", "ka": "𒅗", "gal": "𒃲", "tur": "𒌉",
    "nig": "𒐊", "sag": "𒊕", "kur": "𒆳", "e2": "𒂍", "ud": "𒌓",
}

REVERSE_MAP = {v: k for k, v in CUNEIFORM_MAP.items() if len(v) == 1}


def to_cuneiform_script(transliteration):
    """Convert transliteration to Unicode cuneiform signs."""
    text = transliteration.lower().strip()
    chunks = re.split(r'[\s\-]+', text)
    result = []
    for chunk in chunks:
        if not chunk:
            continue
        matched = False
        for k in sorted(CUNEIFORM_MAP.keys(), key=len, reverse=True):
            if chunk == k:
                result.append(CUNEIFORM_MAP[k])
                matched = True
                break
        if not matched:
            for k in sorted(CUNEIFORM_MAP.keys(), key=len, reverse=True):
                if chunk.startswith(k):
                    result.append(CUNEIFORM_MAP[k])
                    matched = True
                    break
        if not matched:
            result.append(f"[{chunk}]")
    return " ".join(result)


DIRECTIONS = [
    "Translate English to Akkadian cuneiform:",
    "Translate English to Sumerian cuneiform:",
    "Translate Akkadian cuneiform to English:",
    "Translate Sumerian cuneiform to English:",
    "Translate Hittite cuneiform to English:",
    "Translate Linear B to English:",
    "Transliterate Akkadian cuneiform to plain:",
    "Transliterate Sumerian cuneiform to plain:",
    "Convert Akkadian transliteration to cuneiform:",
    "Convert Sumerian transliteration to cuneiform:",
]


def translate(text, direction):
    """Run translation through the model."""
    if not text.strip():
        return "", "", ""

    if not MODEL_LOADED:
        return "Model not available.", "", ""

    prompt = f"{direction} {text}"

    try:
        inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
        outputs = model.generate(
            **inputs,
            max_length=256,
            num_beams=4,
            early_stopping=True,
        )
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)

        # Generate cuneiform script if output looks like transliteration
        cuneiform_script = ""
        if direction.startswith("Translate English") or direction.startswith("Convert"):
            cuneiform_script = to_cuneiform_script(result)
        elif direction.startswith("Transliterate"):
            cuneiform_script = to_cuneiform_script(result)

        # Show what went in
        info = f"Direction: {direction}\nInput tokens: {len(inputs['input_ids'][0])}\nOutput tokens: {len(outputs[0])}"

        return result, cuneiform_script, info

    except Exception as e:
        return f"Error: {str(e)}", "", ""


EXAMPLES = [
    ["The king built a great temple", "Translate English to Akkadian cuneiform:"],
    ["The king built a great temple", "Translate English to Sumerian cuneiform:"],
    ["szar-ru ra-bu-u2 e2-gal ib-ni", "Translate Akkadian cuneiform to English:"],
    ["lugal e2 gal mu-un-du3", "Translate Sumerian cuneiform to English:"],
    ["lugal kur gal-e", "Convert Sumerian transliteration to cuneiform:"],
]


CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Quicksand:wght@400;600;700&display=swap');

.cuneiform-output {
    font-size: 2.2em; line-height: 1.8; letter-spacing: 0.15em;
    padding: 20px; text-align: center; color: #8B6914;
    background: #FAF8F0; border-radius: 12px; border: 1px solid #E8E0C8;
    min-height: 60px;
}
.gradio-container {
    background: #F5F4F0 !important; max-width: 800px !important; margin: 0 auto !important;
}
.gr-button-primary {
    background: #8B6914 !important; border: none !important;
    font-family: 'Quicksand', sans-serif !important;
    padding: 12px 36px !important; border-radius: 12px !important;
}
.gr-button-primary:hover { background: #7A5B10 !important; }
footer { display: none !important; }
"""

with gr.Blocks(css=CUSTOM_CSS, title="Cuneiform Translator", theme=gr.themes.Soft()) as app:

    gr.HTML("""
    <div style="text-align:center; padding: 20px 0 8px;">
        <h1 style="font-family: 'Quicksand', sans-serif; color: #333; font-size: 2.4em; font-weight: 700;">
            📜 Cuneiform Translator
        </h1>
        <p style="font-family: 'Quicksand', sans-serif; color: #999; font-size: 0.9em; max-width: 500px; margin: 4px auto 0;">
            Translate between English and the oldest writing systems on Earth.
            Akkadian, Sumerian, Hittite, and Linear B.
        </p>
    </div>
    """)

    text_input = gr.Textbox(
        label="Text",
        placeholder="Enter English text or cuneiform transliteration...",
        lines=3,
    )

    direction = gr.Dropdown(
        choices=DIRECTIONS,
        value=DIRECTIONS[0],
        label="Direction",
    )

    btn = gr.Button("Translate", variant="primary", size="lg")

    transliteration_output = gr.Textbox(label="Transliteration / Translation", lines=3)
    cuneiform_output = gr.HTML(label="Cuneiform Script")
    info_output = gr.Textbox(label="Info", lines=3, visible=True)

    def run(text, direction):
        result, cuneiform, info = translate(text, direction)
        cuneiform_html = f'<div class="cuneiform-output">{cuneiform}</div>' if cuneiform else ""
        return result, cuneiform_html, info

    btn.click(
        fn=run,
        inputs=[text_input, direction],
        outputs=[transliteration_output, cuneiform_output, info_output],
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[text_input, direction],
        label="Try these",
    )

    gr.HTML("""
    <div style="text-align:center; padding: 16px 0 4px; color: #CCC; font-size: 0.7em; font-family: 'Quicksand', sans-serif; line-height: 1.8;">
        Powered by Thalesian/cuneiformBase-400m (396M parameters)<br>
        <span style="color:#8B6914;">Heuremen — Build Small Hackathon 2026</span>
    </div>
    """)

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