import gradio as gr
import torch
import torch.nn as nn
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from peft import PeftModel
from huggingface_hub import hf_hub_download
import os
import time

BASE_MODEL = os.getenv("BASE_MODEL", "facebook/nllb-200-distilled-600M")

EXPERT_LANGS = ["en", "de", "fr", "nl"]

EXPERT_REPOS = {
    "en": os.getenv("ADAPTER_EN", "entropy25/moe_en"),
    "de": os.getenv("ADAPTER_DE", "entropy25/moe_de"),
    "fr": os.getenv("ADAPTER_FR", "entropy25/moe_fr"),
    "nl": os.getenv("ADAPTER_NL", "entropy25/moe_nl"),
}

ROUTER_REPO = os.getenv("ROUTER_REPO", "entropy25/moe_router")

LANG_CODES = {
    "en": "eng_Latn",
    "de": "deu_Latn",
    "fr": "fra_Latn",
    "nl": "nld_Latn",
    "no": "nob_Latn",
}

LANG_LABELS = {
    "en": "English",
    "de": "German",
    "fr": "French",
    "nl": "Dutch",
}

MAX_LENGTH = 256
NUM_BEAMS = 3
ROUTER_HIDDEN_DIM = 256

device = "cuda" if torch.cuda.is_available() else "cpu"


class GatedRouter(nn.Module):
    """Mirrors the router architecture used during training:
    Linear(hidden->256) -> ReLU -> Dropout -> Linear(256->num_experts)
    """

    def __init__(self, input_dim, hidden_dim, num_experts):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, num_experts),
        )

    def forward(self, hidden_states, attention_mask=None):
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).float()
            pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
        else:
            pooled = hidden_states.mean(dim=1)
        logits = self.network(pooled)
        weights = torch.softmax(logits, dim=-1)
        return weights, logits


print("Loading tokenizer and backbone...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

backbone = AutoModelForSeq2SeqLM.from_pretrained(
    BASE_MODEL,
    low_cpu_mem_usage=True,
).to(device)
backbone.eval()

print("Loading expert LoRA adapters...")
first_lang = EXPERT_LANGS[0]
peft_model = PeftModel.from_pretrained(
    backbone, EXPERT_REPOS[first_lang], adapter_name=first_lang
)
for lang in EXPERT_LANGS[1:]:
    peft_model.load_adapter(EXPERT_REPOS[lang], adapter_name=lang)
peft_model.eval()

print("Loading router...")
router_path = hf_hub_download(repo_id=ROUTER_REPO, filename="router.pt")
router = GatedRouter(
    input_dim=backbone.config.hidden_size,
    hidden_dim=ROUTER_HIDDEN_DIM,
    num_experts=len(EXPERT_LANGS),
).to(device)
router.load_state_dict(torch.load(router_path, map_location=device))
router.eval()

print("All models loaded.")

EXAMPLES = {
    "en": "Mud weight adjusted to 1.82 specific gravity at 3,247 meters depth.",
    "de": "Das Schlammgewicht wurde bei 3.247 Metern Tiefe auf ein spezifisches Gewicht von 1,82 angepasst.",
    "fr": "Le poids de la boue a été ajusté à une densité de 1,82 à 3 247 mètres de profondeur.",
    "nl": "Het slikgewicht werd aangepast naar een soortelijk gewicht van 1,82 op 3.247 meter diepte.",
}


@torch.inference_mode()
def route_and_translate(text, source_mode):
    if not text.strip():
        return "", ""

    start = time.time()

    # IMPORTANT: tokenizer.src_lang is stateful and persists across calls.
    # Reset it to a fixed, neutral value before the routing pass so the
    # router sees a consistent language-tag prefix regardless of what was
    # translated previously.
    tokenizer.src_lang = LANG_CODES["en"]
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=MAX_LENGTH)
    inputs = {k: v.to(device) for k, v in inputs.items()}

    # Compute router decision from encoder hidden states
    encoder_outputs = peft_model.get_encoder()(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        return_dict=True,
    )
    router_weights, router_logits = router(encoder_outputs.last_hidden_state, inputs["attention_mask"])

    if source_mode == "Auto-detect (MoE router)":
        expert_idx = router_weights.mean(dim=0).argmax().item()
        chosen_lang = EXPERT_LANGS[expert_idx]
    else:
        chosen_lang = {v: k for k, v in LANG_LABELS.items()}[source_mode]

    # Re-tokenize with the correct src_lang set for NLLB
    tokenizer.src_lang = LANG_CODES[chosen_lang]
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=MAX_LENGTH)
    inputs = {k: v.to(device) for k, v in inputs.items()}

    peft_model.set_adapter(chosen_lang)
    output = peft_model.generate(
        **inputs,
        forced_bos_token_id=tokenizer.convert_tokens_to_ids(LANG_CODES["no"]),
        max_length=MAX_LENGTH,
        num_beams=NUM_BEAMS,
        early_stopping=True,
    )
    translation = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
    elapsed = time.time() - start

    weight_str = ", ".join(
        f"{LANG_LABELS[l]}: {router_weights[0, i].item():.1%}"
        for i, l in enumerate(EXPERT_LANGS)
    )
    logit_str = ", ".join(
        f"{LANG_LABELS[l]}: {router_logits[0, i].item():.3f}"
        for i, l in enumerate(EXPERT_LANGS)
    )
    info = (
        f"Routed to: {LANG_LABELS[chosen_lang]} expert  |  {elapsed:.2f}s\n"
        f"Router weights — {weight_str}\n"
        f"Raw logits — {logit_str}"
    )

    return translation, info


def load_example(lang_label):
    lang = {v: k for k, v in LANG_LABELS.items()}.get(lang_label, "en")
    return EXAMPLES[lang]


custom_css = """
.gradio-container {
    max-width: 1000px !important;
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
}
.translate-box {
    background: white !important;
    border-radius: 5px !important;
    box-shadow: 0 2px 4px rgba(0,0,0,0.08) !important;
    margin: 20px 0 !important;
}
.text-area textarea {
    border: none !important;
    font-size: 17px !important;
    line-height: 1.7 !important;
    padding: 20px !important;
    min-height: 200px !important;
}
.translate-btn {
    background: #ff8c00 !important;
    color: white !important;
    border: none !important;
    padding: 12px 24px !important;
    font-size: 15px !important;
    font-weight: 500 !important;
    border-radius: 4px !important;
}
.time-info {
    text-align: center !important;
    color: #666 !important;
    font-size: 13px !important;
    padding: 10px !important;
    font-style: italic !important;
    white-space: pre-line !important;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
    gr.HTML(
        "<div style='text-align:center;padding:20px 0 0 0'>"
        "<h2>Modular Mixture-of-Experts Translation</h2>"
        "<p style='color:#888'>EN / DE / FR / NL → Norwegian &nbsp;·&nbsp; Petroleum Domain</p>"
        "</div>"
    )

    with gr.Row():
        source_mode = gr.Dropdown(
            choices=["Auto-detect (MoE router)"] + [LANG_LABELS[l] for l in EXPERT_LANGS],
            value="Auto-detect (MoE router)",
            label="Source language",
        )

    with gr.Row():
        with gr.Column():
            with gr.Group(elem_classes="translate-box"):
                input_text = gr.Textbox(
                    placeholder="Type text in English, German, French, or Dutch",
                    show_label=False,
                    lines=8,
                    container=False,
                    elem_classes="text-area",
                )
        with gr.Column():
            with gr.Group(elem_classes="translate-box"):
                output_text = gr.Textbox(
                    placeholder="Norwegian translation",
                    show_label=False,
                    lines=8,
                    container=False,
                    elem_classes="text-area",
                    interactive=False,
                )

    with gr.Row():
        translate_btn = gr.Button("Translate", variant="primary", elem_classes="translate-btn", size="lg")

    with gr.Row():
        info_display = gr.Textbox(show_label=False, container=False, interactive=False, elem_classes="time-info")

    with gr.Accordion("Example Sentences", open=True):
        with gr.Row():
            for label in [LANG_LABELS[l] for l in EXPERT_LANGS]:
                gr.Button(label, size="sm").click(
                    lambda l=label: load_example(l), outputs=input_text
                )

    translate_btn.click(
        fn=route_and_translate,
        inputs=[input_text, source_mode],
        outputs=[output_text, info_display],
    )

demo.queue().launch()