"""
HuggingFace Space: Turkish Diacritic Restoration
Loads the CRF model from the emircanerol/turkish-diacritic-crf model repo.
"""
import gradio as gr
import torch
from huggingface_hub import hf_hub_download

# ── Load model ────────────────────────────────────────────────────────────────

def _load_model():
    # Download inference code and weights from the model repo
    vocab_path   = hf_hub_download("emircanerol/turkish-diacritic-crf", "ldgc/vocab.py")
    crf_gpu_path = hf_hub_download("emircanerol/turkish-diacritic-crf", "crf_gpu.py")
    weights_path = hf_hub_download("emircanerol/turkish-diacritic-crf", "crf_gpu.safetensors")

    # Make ldgc.vocab importable from the cached path
    import importlib.util, sys, os
    ldgc_dir = os.path.dirname(vocab_path)
    pkg_dir  = os.path.dirname(ldgc_dir)

    # Install ldgc as a package if not already
    if pkg_dir not in sys.path:
        sys.path.insert(0, pkg_dir)
    init = os.path.join(ldgc_dir, "__init__.py")
    if not os.path.exists(init):
        open(init, "w").close()

    # Load crf_gpu module from cached path
    spec = importlib.util.spec_from_file_location("crf_gpu", crf_gpu_path)
    mod  = importlib.util.module_from_spec(spec)
    sys.modules["crf_gpu"] = mod
    spec.loader.exec_module(mod)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model  = mod.CRFGPUModel.from_pretrained(weights_path, device=device)
    return model, mod.predict_stream


MODEL, predict_stream = _load_model()

# ── Gradio interface ──────────────────────────────────────────────────────────

_EXAMPLES = [
    "turkce dogal dil isleme cok onemlidir",
    "bugun hava cok guzel, disari cikmak istiyorum",
    "universite ogrencileri kütüphane de calisıyor",
    "ruzgar bugun cok siddetli esiyor",
    "turkiye'nin baskenti ankara'dir",
]


def restore(text: str) -> str:
    if not text.strip():
        return ""
    lines  = [l for l in text.splitlines() if l.strip()]
    preds  = MODEL.predict(lines, batch_size=128)
    return "\n".join(preds)


with gr.Blocks(title="Turkish Diacritic Restoration") as demo:
    gr.Markdown(
        """
        # Turkish Diacritic Restoration
        Restores missing diacritics in Turkish text
        (ç, ğ, ı, ö, ş, ü and circumflex variants).

        **Model**: Bidirectional CRF with ±2 character context and bigram features,
        trained on 200k Wikipedia sentences.
        · [Model repo](https://huggingface.co/emircanerol/turkish-diacritic-crf)
        · [Code](https://github.com/emircanerol/tr-grammar)
        """
    )

    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(
                label="Noisy Turkish text (one sentence per line)",
                placeholder="turkce cok guzel…",
                lines=6,
            )
            btn = gr.Button("Restore diacritics", variant="primary")
        with gr.Column():
            out = gr.Textbox(label="Restored text", lines=6, interactive=False)

    gr.Examples(
        examples=[[e] for e in _EXAMPLES],
        inputs=inp,
        outputs=out,
        fn=restore,
        cache_examples=True,
    )

    btn.click(restore, inputs=inp, outputs=out)
    inp.submit(restore, inputs=inp, outputs=out)

demo.launch()