import json
import os
import time
import random
import torch
import re
import math
import gradio as gr
import numpy as np
from collections import defaultdict
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

os.environ["TOKENIZERS_PARALLELISM"] = "0"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

class TorchTracemalloc:
    track_memory_consumption = []

    def __enter__(self):
        self.begin = torch.cuda.memory_allocated()
        torch.cuda.reset_max_memory_allocated()
        return self

    def __exit__(self, *exc):
        peak = torch.cuda.max_memory_allocated()
        peaked = (peak - self.begin) // 1024 ** 2
        TorchTracemalloc.track_memory_consumption.append(peaked)
        print(f"Memory consumed: {peaked} MB")  # Debugging print

def format_response(dialog, response):
    question = next((turn['content'] for turn in dialog if turn['role'] == 'user'), 'No question found')
    return {"question": question, "answer": response}

# Global variables to store the model and tokenizer
global_model = None
global_tokenizer = None

def load_model_and_tokenizer(model_name, dtype, kv_bits):
    global global_model, global_tokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    special_tokens = {"pad_token": "<PAD>"}
    tokenizer.add_special_tokens(special_tokens)

    config = AutoConfig.from_pretrained(model_name)
    if kv_bits != "unquantized":
        quantizer_path = f"codebooks/{model_name.split('/')[-1]}_{kv_bits}bit.xmad"
        setattr(config, "quantizer_path", quantizer_path)

    if dtype == "bf16":
        dtype = torch.bfloat16
    elif dtype == "fp16":
        dtype = torch.float16
    elif dtype == "fp32":
        dtype = torch.float32

    model = AutoModelForCausalLM.from_pretrained(model_name, config=config, torch_dtype=dtype, device_map="auto")

    if len(tokenizer) > model.get_input_embeddings().weight.shape[0]:
        model.resize_token_embeddings(len(tokenizer))

    tokenizer.padding_side = "left"
    model.config.pad_token_id = tokenizer.pad_token_id

    global_model = model
    global_tokenizer = tokenizer

def load_questions(prompts_path, custom_questions):
    with open(prompts_path, "r") as file:
        dialogs = json.load(file)

    selected_dialogs = []
    if custom_questions:
        for question in custom_questions:
            if question.strip():
                custom_dialog = [{"role": "user", "content": question}]
                selected_dialogs.append(custom_dialog)

    num_questions = 60 - len(selected_dialogs)
    random.shuffle(dialogs)
    selected_dialogs.extend(dialogs[:num_questions])

    return selected_dialogs[:60]

def markdown_to_plain_text(markdown_text):
    # Convert markdown bold (**) to plain text uppercase
    markdown_text = re.sub(r'\*\*(.*?)\*\*', r'\1'.upper(), markdown_text)
    # Convert markdown italics (*) to plain text
    markdown_text = re.sub(r'\*(.*?)\*', r'\1', markdown_text)
    # Remove markdown headers (###)
    markdown_text = re.sub(r'### ', '', markdown_text)
    # Convert markdown lists (- or *)
    markdown_text = re.sub(r'^\s*[-*]\s+', '', markdown_text, flags=re.MULTILINE)
    # Remove remaining markdown formatting
    markdown_text = re.sub(r'[`~>]', '', markdown_text)
    return markdown_text

def infer(model_name, dialogs, num_new_tokens, temperature, dtype, kv_bits, progress=gr.Progress()):
    print("Starting inference...")
    global global_model, global_tokenizer
    
    model = global_model
    tokenizer = global_tokenizer
    
    batch_inputs = [
        tokenizer.apply_chat_template(dialog, tokenize=False, add_generation_prompt=True)
        for dialog in dialogs
    ]

    responses = []
    start_time = time.time()
    batch_size = 60  # Adjust batch size based on GPU capacity
    num_dialogs = len(dialogs)
    # total_time = 0
    # total_tokens = 0
    # total_ttft = 0
    # num_batches = (num_dialogs + batch_size - 1) // batch_size

    actual_batch_size = min(batch_size, num_dialogs)
    total_time = 0
    total_tokens = 0
    total_ttft = 0
    num_batches = math.ceil(num_dialogs / actual_batch_size)

    memory_avg = []
    tokens_per_sec_avg = []
    time_to_first_token_avg = []
    responses_by_batch_size = defaultdict(list)
    batch_generation_time = 0
    total_generation_time  = 0

    terminators = [
        tokenizer.eos_token_id,
        tokenizer.convert_tokens_to_ids("<|eot_id|>"),
    ]

    with TorchTracemalloc() as tt:
        for i in range(0, num_dialogs, actual_batch_size):
        # for batch_idx in range(num_batches):
            batch = batch_inputs[i : i + actual_batch_size]
            try:
                encoded_inputs = tokenizer(
                    batch,
                    padding=True,
                    truncation=False,
                    return_tensors="pt",
                )

                input_ids = encoded_inputs["input_ids"].to(model.device)
                attention_mask = encoded_inputs["attention_mask"].to(
                    model.device
                )

                torch.cuda.synchronize()
                start_time = time.perf_counter()

                with torch.no_grad():
                        output_tokens = model.generate(
                            input_ids,
                            attention_mask=attention_mask,
                            max_new_tokens=num_new_tokens,
                            num_return_sequences=1,
                            do_sample=True,
                            temperature=temperature,
                            pad_token_id=tokenizer.pad_token_id,
                            eos_token_id=terminators,
                        )

                torch.cuda.synchronize()
                end_time = time.perf_counter()

                batch_time = end_time - start_time
                total_time += batch_time
                batch_generation_time += (
                    batch_time  # Add to batch generation time
                )
                total_generation_time += (
                    batch_time  # Add to total generation time
                )
                total_tokens += output_tokens.numel()

                if i == 0:
                    total_ttft = batch_time

            # if batch_idx == 0:
            #     total_ttft = batch_time

                decoded_outputs = tokenizer.batch_decode(
                        output_tokens, skip_special_tokens=True
                )
            # decoded_outputs = tokenizer.batch_decode(output_tokens, skip_special_tokens=True)

                for j, response in enumerate(decoded_outputs):
                    original_dialog = dialogs[i + j]
                    formatted_responses = format_response(
                        original_dialog, response
                    )
                    responses.append(formatted_responses)
                    # responses_by_batch_size[batch_size].append(
                    #     formatted_response
                    # )
                    # Format the responses
                    formatted_responses = "\n\n---\n\n".join([f"**Question**: {res['question']}\n\n**Answer**: {res['answer']}" for res in responses])
                    plain_text_responses = markdown_to_plain_text(formatted_responses)
                    yield plain_text_responses
                    progress(i, desc="Processing batches")

                    torch.cuda.empty_cache()

            except Exception as e:
                print(
                    f"Error processing batch {i//batch_size + 1}: {str(e)}"
                )
                continue


    elapsed_time = total_time
    tokens_per_second = total_tokens / total_time if total_time > 0 else 0
    # avg_memory_consumption = np.mean(TorchTracemalloc.track_memory_consumption)
    total_memory_consumption = np.sum(TorchTracemalloc.track_memory_consumption)
    avg_memory_consumption = total_memory_consumption/num_dialogs

    # Use actual_batch_size in calculations
    ttft = (
        total_ttft / actual_batch_size if actual_batch_size > 0 else 0
    )

    print(f"Inference completed in {elapsed_time:.2f} seconds.")
    
    yield {
        "Time Taken (seconds)": elapsed_time,
        "Tokens per Second": tokens_per_second,
        "Time to First Token (TTFT, seconds)": ttft,
        # "Formatted Responses": formatted_responses,
        "Formatted Responses": plain_text_responses,
        "Average Memory Consumption per Question (MB)": avg_memory_consumption,
        "Total Memory Consumption (MB)": total_memory_consumption
    }

# Demo function
def demo(num_new_tokens, temperature, custom_questions_text, kv_bits=1, progress=gr.Progress()):
    custom_questions = custom_questions_text.split("\n")
    print("Loading questions...")
    dialogs = load_questions("chats_sys_none.json", custom_questions)
    print(f"{len(dialogs)} questions loaded. Starting inference...")
    
    result_gen = infer("NousResearch/Meta-Llama-3-8B-Instruct", dialogs, num_new_tokens, temperature, "fp16", kv_bits, progress=progress)
    
    formatted_responses = ""
    for result in result_gen:
        if isinstance(result, str):
            formatted_responses = result
            yield None, None, None, None, None, None, None, formatted_responses
        else:
            time_taken = result["Time Taken (seconds)"]
            tokens_per_second = result["Tokens per Second"]
            ttft = result["Time to First Token (TTFT, seconds)"]
            avg_memory_consumption = result["Average Memory Consumption per Question (MB)"]
            total_memory_consumption = result["Total Memory Consumption (MB)"]
            formatted_responses = result["Formatted Responses"]
            yield time_taken, tokens_per_second, ttft, avg_memory_consumption, total_memory_consumption, formatted_responses

# Load JSON data
with open("chats_sys_none.json", "r") as file:
    json_data = json.load(file)

# Load 50 random questions into the input area by default
def load_default_questions():
    random.shuffle(json_data)
    default_questions = [dialog[0]['content'] for dialog in json_data[:50] if 'content' in dialog[0]]
    return "\n".join(default_questions)

# Load default questions on button click
def load_questions_action():
    return load_default_questions()

# Gradio interface
css = """
body, html {
    height: 100vh;
    margin: 0;
}

.gradio-container {
    height: 100vh;
}

#main-row {
    height: 90vh;
    display: flex;
}

#control-panel, #formatted-responses-container {
    height: 90vh;
    box-sizing: border-box;
    display: flex;
    flex-direction: column;
    overflow: hidden;
    flex: 1; /* Ensure equal width */
}

#control-panel {
    flex: 1; /* Ensure equal height */
}

#custom-questions-text {
    flex-grow: 1;
    overflow-y: auto;
    max-height: 30vh; /* Limit height of custom questions text */
}

#metrics-panel {
    display: flex;
    flex-wrap: wrap;
    gap: 1vh;
    margin-bottom: 1vh;
    flex-shrink: 0;
    height: auto; /* Let the panel size adjust based on its content */
}

#metrics-panel .metric {
    flex: 1 1 48%;
    min-width: 10vw;
    box-sizing: border-box;
}

#buttons-container {
    display: flex;
    justify-content: space-between;
    min-height: 6vh; /* Minimum height for buttons container */
    flex-shrink: 0;
}
"""

with gr.Blocks(css=css) as app:
    with gr.Row(elem_id="main-row", equal_height=True):
        with gr.Column(elem_id="control-panel"):
            num_new_tokens = gr.Slider(label="Number of New Tokens", minimum=128, maximum=2048, step=128, value=512)
            temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, step=0.1, value=0.4)
            custom_questions_text = gr.Textbox(
                label="Custom Questions",
                placeholder="Type your custom questions here, one per line...",
                autoscroll=False,
                container=False,
                lines=5,
                elem_id="custom-questions-text"
            )
            with gr.Row(elem_id="metrics-panel"):
                time_taken = gr.Number(label="Time Taken (seconds)", interactive=False, elem_classes=["metric"])
                tokens_per_second = gr.Number(label="Tokens per Second", interactive=False, elem_classes=["metric"])
                ttft = gr.Number(label="Time to First Token (TTFT, seconds)", interactive=False, elem_classes=["metric"])
                total_memory_consumption = gr.Number(label="Total Memory Consumption (MB)", interactive=False, elem_classes=["metric"])
                avg_memory_consumption = gr.Number(label="Average Memory Consumption per Question (MB)", interactive=False, elem_classes=["metric"])
            with gr.Row(elem_id="buttons-container"):
                load_questions_btn = gr.Button("Load Default Questions")
                demo_btn = gr.Button("Run Inference", elem_id="run-inference-btn")

        formatted_responses = gr.Textbox(
            label="Formatted Responses",
            elem_id="formatted-responses",
            value="No responses yet. Run the inference to see results.",
            lines=37,
            container=False,
            autoscroll=False,
            show_copy_button=True
        )

        load_questions_btn.click(fn=load_questions_action, inputs=[], outputs=custom_questions_text)
        demo_btn.click(demo, inputs=[num_new_tokens, temperature, custom_questions_text], outputs=[time_taken, tokens_per_second, ttft, avg_memory_consumption, total_memory_consumption, formatted_responses])

if __name__ == "__main__":
    print("Loading model and tokenizer on startup...")
    # load_model_and_tokenizer("NousResearch/Meta-Llama-3-8B-Instruct", "fp16", "1")
    print("Model and tokenizer loaded. Starting Gradio interface...")
    app.launch()