import gradio as gr
import nltk
nltk.download('punkt_tab')
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from IndicTransToolkit import IndicProcessor
import torch

# Load IndicTrans2 model
model_name = "ai4bharat/indictrans2-indic-indic-dist-320M"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
ip = IndicProcessor(inference=True)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model.to(DEVICE)

def split_text_into_batches(text, max_tokens_per_batch):
    sentences = nltk.sent_tokenize(text)  # Tokenize text into sentences
    batches = []
    current_batch = ""
    for sentence in sentences:
        if len(current_batch) + len(sentence) + 1 <= max_tokens_per_batch:  # Add 1 for space
            current_batch += sentence + " "  # Add sentence to current batch
        else:
            batches.append(current_batch.strip())  # Add current batch to batches list
            current_batch = sentence + " "  # Start a new batch with the current sentence
    if current_batch:
        batches.append(current_batch.strip())  # Add the last batch
    return batches

def run_translation(file_uploader, input_text, source_language, target_language):
    if file_uploader is not None:
        with open(file_uploader.name, "r", encoding="utf-8") as file:
            input_text = file.read()

    # Language mapping
    lang_code_map = {
        "Hindi": "hin_Deva",
        "Punjabi": "pan_Guru",
        "English": "eng_Latn",
    }

    src_lang = lang_code_map[source_language]
    tgt_lang = lang_code_map[target_language]

    max_tokens_per_batch = 256
    batches = split_text_into_batches(input_text, max_tokens_per_batch)
    translated_text = ""

    for batch in batches:
        batch_preprocessed = ip.preprocess_batch([batch], src_lang=src_lang, tgt_lang=tgt_lang)
        inputs = tokenizer(
            batch_preprocessed,
            truncation=True,
            padding="longest",
            return_tensors="pt",
            return_attention_mask=True,
        ).to(DEVICE)

        with torch.no_grad():
            generated_tokens = model.generate(
                **inputs,
                use_cache=True,
                min_length=0,
                max_length=256,
                num_beams=5,
                num_return_sequences=1,
            )

        with tokenizer.as_target_tokenizer():
            decoded_tokens = tokenizer.batch_decode(
                generated_tokens.detach().cpu().tolist(),
                skip_special_tokens=True,
                clean_up_tokenization_spaces=True,
            )

        translations = ip.postprocess_batch(decoded_tokens, lang=tgt_lang)
        translated_text += " ".join(translations) + " "

    output = translated_text.strip()
    _output_name = "result.txt"
    with open(_output_name, "w", encoding="utf-8") as out_file:
        out_file.write(output)

    return output, _output_name

# Define Gradio UI
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            file_uploader = gr.File(label="Upload a text file (Optional)")
            input_text = gr.Textbox(label="Input text", lines=5, placeholder="Enter text here...")
            source_language = gr.Dropdown(
                label="Source language",
                choices=["Hindi", "Punjabi", "English"],
                value="Hindi",
            )
            target_language = gr.Dropdown(
                label="Target language",
                choices=["Hindi", "Punjabi", "English"],
                value="English",
            )
            btn = gr.Button("Translate")
        with gr.Column():
            output_text = gr.Textbox(label="Translated text", lines=5)
            output_file = gr.File(label="Translated text file")

    btn.click(
        fn=run_translation,
        inputs=[file_uploader, input_text, source_language, target_language],
        outputs=[output_text, output_file],
    )

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