import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import torch
from threading import Thread
import re

phi4_model_path = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
device = "cuda:0" if torch.cuda.is_available() else "cpu"

phi4_model = AutoModelForCausalLM.from_pretrained(phi4_model_path, device_map="auto", torch_dtype="auto")
phi4_tokenizer = AutoTokenizer.from_pretrained(phi4_model_path)

def format_math(text):
    text = re.sub(r"\[(.*?)\]", r"$$\1$$", text, flags=re.DOTALL)
    text = text.replace(r"\(", "$").replace(r"\)", "$")
    return text

#@spaces.GPU(duration=60)
def generate_response(user_message, max_tokens, temperature, top_p, history_state):
    if not user_message.strip():
        return history_state, history_state

    model = phi4_model
    tokenizer = phi4_tokenizer
    start_tag = "<|im_start|>"
    sep_tag = "<|im_sep|>"
    end_tag = "<|im_end|>"

    system_message = "Your role as an assistant..."
    prompt = f"{start_tag}system{sep_tag}{system_message}{end_tag}"
    for message in history_state:
        if message["role"] == "user":
            prompt += f"{start_tag}user{sep_tag}{message['content']}{end_tag}"
        elif message["role"] == "assistant" and message["content"]:
            prompt += f"{start_tag}assistant{sep_tag}{message['content']}{end_tag}"
    prompt += f"{start_tag}user{sep_tag}{user_message}{end_tag}{start_tag}assistant{sep_tag}"

    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    generation_kwargs = {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
        "max_new_tokens": int(max_tokens),
        "do_sample": True,
        "temperature": temperature,
        "top_k": 50,
        "top_p": top_p,
        "repetition_penalty": 1.0,
        "pad_token_id": tokenizer.eos_token_id,
        "streamer": streamer,
    }

    try:
        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start() 
    except Exception:
        yield history_state + [{"role": "user", "content": user_message}, {"role": "assistant", "content": "⚠️ Generation failed."}], history_state
        return

    assistant_response = ""
    new_history = history_state + [
        {"role": "user", "content": user_message},
        {"role": "assistant", "content": ""}
    ]

    try:
        for new_token in streamer:
            if "<|end" in new_token:
                continue
            cleaned_token = new_token.replace("<|im_start|>", "").replace("<|im_sep|>", "").replace("<|im_end|>", "")
            assistant_response += cleaned_token
            new_history[-1]["content"] = assistant_response.strip()
            yield new_history, new_history
    except Exception:
        pass

    yield new_history, new_history

example_messages = {
    "Combinatorics": "From all the English alphabets, five letters are chosen and are arranged in alphabetical order. The total number of ways, in which the middle letter is 'M', is?",
    "Co-ordinate Geometry": "A circle \\(C\\) of radius 2 lies in the second quadrant and touches both the coordinate axes. Let \\(r\\) be the radius of a circle that has centre at the point \\((2, 5)\\) and intersects the circle \\(C\\) at exactly two points. If the set of all possible values of \\(r\\) is the interval \\((\\alpha, \\beta)\\), then \\(3\\beta - 2\\alpha\\) is?",
    "Prob-Stats": "A coin is tossed three times. Let \(X\) denote the number of times a tail follows a head. If \\(\\mu\\) and \\(\\sigma^2\\) denote the mean and variance of \\(X\\), then the value of \\(64(\\mu + \\sigma^2)\\) is?"
}

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Ramanujan Ganit R1 14B V1 Chatbot 
    
        Welcome to the Ramanujan Ganit R1 14B V1 Chatbot, developed by Fractal AI Research! 
        
        Our model excels at reasoning tasks in mathematics and science.
        
        Try the example problems below from JEE Main 2025 or type in your own problems to see how our model breaks down complex reasoning problems.
        """
    )

    history_state = gr.State([])

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Settings")
            max_tokens_slider = gr.Slider(minimum=6144, maximum=32768, step=1024, value=16384, label="Max Tokens")
            with gr.Accordion("Advanced Settings", open=False):
                temperature_slider = gr.Slider(minimum=0.1, maximum=2.0, value=0.6, label="Temperature")
                top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p")

        with gr.Column():
            user_input = gr.Textbox(
            label="User Input",
            placeholder="Type your question here...",
            lines=2
    )

            with gr.Row():
                submit_button = gr.Button("Send", variant="primary", scale=1)
                clear_button = gr.Button("Clear", scale=1)
            gr.Markdown("**Try these examples:**")
            with gr.Row():
                example1_button = gr.Button("JEE Main 2025 Combinatorics")
                example2_button = gr.Button("JEE Main 2025 Co-ordinate Geometry")
                example3_button = gr.Button("JEE Main 2025 Prob-Stats")

    submit_button.click(
        fn=generate_response,
        inputs=[user_input, max_tokens_slider, temperature_slider, top_p_slider, history_state],
        outputs=[chatbot, history_state]
    ).then(
        fn=lambda: gr.update(value=""),
        inputs=None,
        outputs=user_input
    )

    clear_button.click(
        fn=lambda: ([], []),
        inputs=None,
        outputs=[chatbot, history_state]
    )

    example1_button.click(
        fn=lambda: gr.update(value=example_messages["Combinatorics"]),
        inputs=None,
        outputs=user_input
    )
    example2_button.click(
        fn=lambda: gr.update(value=example_messages["Co-ordinate Geometry"]),
        inputs=None,
        outputs=user_input
    )
    example3_button.click(
        fn=lambda: gr.update(value=example_messages["Prob-Stats"]),
        inputs=None,
        outputs=user_input
    )

demo.launch(share=True, ssr_mode=False)