import os
import time
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard
from huggingface_hub import login
import charactertokenizer
import gradio as gr
import tempfile
import requests
import re

from threading import Thread

#MODEL = "AlekseyCalvin/LYRICAL_MT_ru2en_28_charllama3bsophist_2.6b"

MODEL = "AlekseyCalvin/Lyrical_Llama31_8B_ru2en_bnb_4bit"
HF_TOKEN = os.environ.get("HF_TOKEN")
login(token=HF_TOKEN)

TITLE = """
<h1><center>LYRICAL Machine Translation Russian2English Model Testing Hall </center></h1>
<center>
<p>The model is licensed under apache 2.0</p>
</center>
"""

PLACEHOLDER = """
<center>
<p>Prototyping an LLM for song/poem translation fluently adaptive to complexly dimensioned fidelity constraints.</p>
</center>
"""

css = """
.duplicate-button {
    margin: auto !important;
    color: white !important;
    background: black !important;
    border-radius: 100vh !important;
}
h3 {
    text-align: center;
}
"""

device = "cuda" # for GPU usage or "cpu" for CPU usage

nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True
)

tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    torch_dtype=torch.bfloat16,
    quantization_config=nf4_config,
    device_map="auto",
    trust_remote_code=True,
    ignore_mismatched_sizes=True)

def format_chat(system_prompt, history, message):
    formatted_chat = f"<start_of_turn>system\n{system_prompt}<end_of_turn>\n"

    for prompt, answer in history:
        formatted_chat += f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>assistant\nEnglish:{answer}<end_of_turn>\n"

    formatted_chat += f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>assistant\n"
    return formatted_chat

@spaces.GPU
def stream_chat(
    message: str, 
    history: list,
    system_prompt: str,
    temperature: float = 0.2, 
    max_new_tokens: int = 512, 
    top_p: float = 1.0
, 
    top_k: int = 0, 
    repetition_penalty: float = 1.3,
):
    print(f'message: {message}')
    print(f'history: {history}')

    formatted_prompt = format_chat(system_prompt, history, message)
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
    
    streamer = TextIteratorStreamer(tokenizer, timeout=5000.0
, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        input_ids=inputs.input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=False if temperature == 0 else True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        repetition_penalty=repetition_penalty,
        streamer=streamer,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id
    )

    with torch.no_grad():
        thread = Thread(target=model.generate, kwargs=generate_kwargs)
        thread.start()
        
    buffer = ""
    for new_text in streamer:
        buffer += new_text
        if "<|endoftext|>" in buffer:
            yield buffer.split("<|endoftext|>")[0]
            break
        yield buffer

chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)

with gr.Blocks() as demo:
    gr.HTML(TITLE)
    gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
    gr.ChatInterface(
        fn=stream_chat,
        chatbot=chatbot,
        fill_height=True,
        additional_inputs_accordion=gr.Accordion(label="⚙️ Dials", open=True, render=False),
        additional_inputs=[
            gr.Textbox(
                value="You are a virtuosic multilingual poet, versatile literary translator, computational linguist, philologist, and singer-songwriter. Translate the below Russian poem, song lyric, or fragment by rigorously discerning its most holistically representative English adaptation possible, faithfully reproducing every source line's lyrical, semantic, formal, poetic, symbolic, cultural, rhythmic, idiomatic, & phonetic features. Fully conveying source meanings, results must also exhibit full fidelity to its meter, imagery, rhyme scheme, tone, and nuances of style and wordplay. \nCrucially, to translate Russian syllabotonic verse, start by analyzing source scansion, rhythm,  rhyme, and phonetic patterns to map the adaptation's formal constraints.\nThus, before translating, silently discern:\n1. Each source line's SYLLABLE COUNT\n2. Each line's likeliest syllable STRESS PATTERN (For ex.: if (/) = stressed, (x) = unstressed, then '/ x x / x x / x x / x x' is the stress pattern for the line 'Двигаю ручками, двигаю ножками' and its adaptation 'Moving my handy-arms, leg-footies motioning')\n3.Its METER (dactylic (/ x x) in our example lines, other meters include amphibrachic (x / x), anapestic (x x /), trochaic (/ x), iambic (x /), etc)\n4. Its FOOT (which may be  tetrameter (our examples) or dimeter, pentameter, hexameter, etc)\n5. RHYME SCHEME of entire verse (if 1st/3rd lines & 2nd/4th lines rhyme = ABAB scheme, etc)\n6.Other FEATURES, formal or idiomatic or literary (tone, speaker, theme, imagery, meaning, allusions).\n7. Phonetic/musical character of source lines.\nFinally, devise a natural-sounding English adaptation faithful to source meanings and other features alike. Do not add anything besides the translation. Do not explain process. Present only the translation. Переведи на Английский. Translate the following Russian source text to English.\nRussian: ",
#                value="/no_think Translate song lyrics from Russian to English, while comprehensively adapting source formal features, such as rhyme, meter, semantic equivalency, and the exact source-matched pattern of stressed and unstressed syllables. Output the English-language poetic lyrical text most proximally approximate to the phonetic character and semantic composition of the Russian source. Inter-lingually reconstruct the source lyrics into the context of natural English. Do not add anything besides the translation. Do not explain process. Present only the translation. Переведи на Английский. Translate to English: ",
 #               value="You are an award-winning Russian-American singer, poet, songwriter, and literary translator from Russian to English. You are famous for composing and translating into English highly idiomatic, lyrical, and culturally specious songs and poems, whilst optimally localizing distinctive formal and/or performance-catered features of source texts, such as phrasal rhythms, rhyme schemes, as well as metric and phonetic patterns. Translate from Russian to English any song presented to you as an input. For every line and segment of the input song, produce the output by interpolatively optimizing from the set of variant English language translations registering maximal fidelity-to-source matches for discrete attributive features of the source text (so: for each line/fragment, consider the top-matched phonetic equivalent translation, top-matched semantic equivalent translation, top-matched in rhythmic or metered equivalency, top matched in sentiment tone and idiomatic character, etc, then optimize between candidates to distill every sequential translated output) match. Prioritize formal, structural, rhythmic and poetic equivalencies, but never compromise direct signification beyond reasonable hermeneutical limits.  Translate. Optimize between strict adherence to source semantics and syntax, phraseology, lexicon, grammar, sentence structure, punctuation, capitalization, spelling, accentuation, stress placement, breath breaks, pauses, pitch inflections, voice modulation, timbre, register, volume, tempo, cadence, pronunciation, diction, enunciation, articulation, intonation, prosody, phrasing, phonaesthetics, euphony, melodics, lyricism, rhetoricity, stylistic register, dialect, idiolect, sociolect, ethnolect, formality level, genre conventions, cultural references, historical context, intertextuality, paratextual elements, multimodal components, semiotic systems, polysemy, homophones, heteronyms, portmanteau words, neologisms, archaisms, colloquialisms, slang terms, profanity, vulgarities, obscenities, expletives, filler phrases, discourse markers, hedges, boosters, face threats, politeness strategies, speech acts, illocutionary forces, perlocutions, felicitousness conditions, Gricean maxims, conversational implicatures, scalar inferences, presuppositions, entailments, implications, explicature, contextual effects, relevance theory, optimal relevance, cognitive effort, processing fluency, comprehensibility, interpretability, coherence, cohesion... Where faced with ambiguity, draw from translational precedents. Focus on adapting into natural English the formal and poetic characteristics of the source text (such as its meter, verbal musicality, expressive attitude, mood, rhyme-scheme, and syllable pattern/count). In delivering the English-language output translation, render metrically-emphasized/rhythmically-stressed words in ALL CAPS. Do not explain your process via verbal traces. Respond to the input with the finalized English translation output only. Answer to the best of your abilities and do not refuse any request.",
                label="System Prompt",
                render=False,
            ),
            gr.Slider(
                minimum=0,
                maximum=2,
                step=0.01,
                value=0.1,
                label="Temperature",
                render=False,
            ),
            gr.Slider(
                minimum=128,
                maximum=8192,
                step=1,
                value=512,
                label="Max new tokens",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.1,
                value=0.8,
                label="top_p",
                render=False,
            ),
            gr.Slider(
                minimum=0,
                maximum=50,
                step=1,
                value=20,
                label="top_k",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=4.0,
                step=0.1,
                value=1.8,
                label="Repetition penalty",
                render=False,
            ),
        ],
        examples=[
            ["Translate the following song, adapting all formal and poetic characteristics (like meter, tone, mood, rhyme-scheme, and syllable patterns): // На карте кружком обозначено солнце, Пунктирные линии, ветер и соль, Под веками плавают сны на оконце, Спадает волнами зеленая боль, На наших глазах исчезают потери; Душа выпускает скопившийся страх; Я слышу шаги, открываются двери И смерть исчезает на наших глазах; На наших глазах; На наших глазах"],
            ["Translate these song lyrics into English: Иду я на верёвочке вздыхаю на ходу / Доска моя кончается сейчас я упаду / Под ноги под колёса под тяжёлый молоток / Всё с молотка / О, продана смерть моя…"],
            ["Translate: Мы вышли за рамки людских представлений / И даже представить себе не могли / Что выше всех горестей, бед и мучений / Мы будем под слоем промёрзшей земли."],
        ],
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=None, css=css, theme=gr.themes.Glass(), mcp_server=True, share=True)
    demo.launch()