"""
ALMA-7B-Ja-V2 Japanese to English Translation Interface
"""
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the ALMA-7B-Ja-V2 model and tokenizer
MODEL_NAME = "webbigdata/ALMA-7B-Ja-V2"
print(f"Loading {MODEL_NAME}...")

try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    # Check if GPU is available
    if torch.cuda.is_available():
        print(f"GPU detected: {torch.cuda.get_device_name(0)}")
        device = "cuda"
        dtype = torch.float16
    else:
        print("No GPU detected, using CPU (inference will be slow)")
        device = "cpu"
        dtype = torch.float32

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        dtype=dtype,  # Use 'dtype' instead of deprecated 'torch_dtype'
        device_map="auto",
        low_cpu_mem_usage=True
    )
    print(f"Model loaded successfully on device: {model.device}")
    print(f"Model dtype: {model.dtype}")
except Exception as e:
    print(f"Error loading model: {e}")
    raise


def translate_japanese_to_english(japanese_text):
    """
    Translate Japanese text to English using ALMA-7B-Ja-V2 model.
    """
    if not japanese_text or not japanese_text.strip():
        return "Please enter some Japanese text to translate."

    try:
        # Format the prompt for translation (official format)
        prompt = (
            f"Translate this from Japanese to English:\n"
            f"Japanese: {japanese_text}\n"
            f"English:"
        )

        # Tokenize and generate
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=512,
                num_beams=3,
                temperature=0.6,
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )

        # Decode the output
        full_output = tokenizer.decode(
            outputs[0],
            skip_special_tokens=True
        )

        # Extract only the English translation (after "English:")
        if "English:" in full_output:
            translation = full_output.split("English:")[-1].strip()
        else:
            translation = full_output.replace(prompt, "").strip()

        return translation if translation else "Translation failed. Please try again."

    except Exception as e:
        error_msg = f"Translation error: {str(e)}"
        print(error_msg)
        return f"An error occurred during translation. Please try again."


# Create Gradio interface
with gr.Blocks(
    title="ALMA-7B-Ja-V2 Translation",
    theme=gr.themes.Soft()
) as demo:
    gr.Markdown("""
    # 🎬 ALMA-7B-Ja-V2: Japanese Dialogue Translation

    Test Japanese to English translation using the ALMA-7B-Ja-V2 model.
    Perfect for translating dialogue sections from Japanese animated films
    (4-6 lines).

    Paste your Japanese dialogue below or try one of the example
    conversations!
    """)

    with gr.Row():
        with gr.Column():
            japanese_input = gr.Textbox(
                label="Japanese Dialogue",
                placeholder=(
                    "Paste your Japanese dialogue here (4-6 lines)...\n"
                    "Example:\n"
                    "キャラクター1: セリフ\n"
                    "キャラクター2: セリフ"
                ),
                lines=8
            )
            translate_btn = gr.Button("Translate", variant="primary")

        with gr.Column():
            english_output = gr.Textbox(
                label="English Translation",
                lines=8
            )

    # Example dialogue sections for testing
    gr.Examples(
        examples=[
            [
                "キャラクターA: おはよう！今日はどこに行くの？\n"
                "キャラクターB: 学校に行くよ。一緒に行かない？\n"
                "キャラクターA: うん、行こう！\n"
                "キャラクターB: じゃあ、準備しよう。"
            ],
            [
                "少年: ここは本当に美しい場所だね。\n"
                "少女: そうね。いつも来たかったの。\n"
                "少年: 一緒に来られて良かった。\n"
                "少女: 私も嬉しいわ。"
            ],
            [
                "主人公: 諦めたらそこで終わりだ。\n"
                "友人: でも、もう無理だよ。\n"
                "主人公: 無理じゃない。まだチャンスはある。\n"
                "友人: 分かった。最後までやってみよう。"
            ],
            [
                "キャラクターA: これは何？\n"
                "キャラクターB: 不思議な石だよ。\n"
                "キャラクターA: どこで見つけたの？\n"
                "キャラクターB: 森の奥で見つけたんだ。\n"
                "キャラクターA: 本当に神秘的ね。"
            ],
        ],
        inputs=japanese_input,
    )

    translate_btn.click(
        fn=translate_japanese_to_english,
        inputs=japanese_input,
        outputs=english_output
    )

# Launch the interface
demo.launch()