import json
import math
import os
import time
from argparse import ArgumentParser
from collections import defaultdict

import matplotlib.pyplot as plt
import numpy as np
import torch

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer


# ! CHECK: if expandable segments can improve memory here
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)


def save_bar_chart(title, x, y, ylabel, xlabel, output_path):
    try:
        plt.style.use("ggplot")
        width = 0.4
        xs = np.arange(len(x))
        plt.figure(figsize=(10, 6))
        plt.bar(xs, height=y, width=width, color="skyblue")
        plt.title(title)
        plt.xticks(xs, x)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.savefig(output_path)
    except Exception as e:
        print(f"Error saving chart {title}: {str(e)}")
    finally:
        plt.close()


def format_response(dialog, response):
    formatted_dialog = dialog.copy()
    formatted_dialog.append({"role": "assistant", "content": response})
    return formatted_dialog


parser = ArgumentParser("chat_with_llama")
# parser.add_argument("--llama", type=int, default=3, choices=[2, 3])
parser.add_argument(
    "--llama", type=str, default="3-instruct", choices=["2", "3-instruct"]
)
parser.add_argument("--prompts_path", type=str, default="chats_sys_none.json")
# parser.add_argument('--batch_size', type=int, default=8)
parser.add_argument("--model_size", type=int, default=8, choices=[7, 8, 13])
parser.add_argument("--num_new_tokens", type=int, default=512)
parser.add_argument(
    "--temperature", type=float, default=0.4, help="Temperature for sampling"
)
parser.add_argument("--window_length", type=int, default=32)
parser.add_argument("--kv_bits", type=int, default=1)
parser.add_argument("--output_path", type=str, default="./output")
parser.add_argument(
    "--dtype", type=str, default="fp16", choices=["fp16", "fp32", "bf16"]
)
args = parser.parse_args()
bits = args.kv_bits

try:
    if args.llama == 2:
        model_name = "NousResearch/Llama-2-7b-hf"
    else:
        model_name = "NousResearch/Meta-Llama-3-8B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Add pad token to the tokenizer for batched inference
    special_tokens = {"pad_token": "<PAD>"}
    tokenizer.add_special_tokens(special_tokens)

    config = AutoConfig.from_pretrained(model_name)

    if isinstance(bits, int):
        if args.llama == 2:
            setattr(
                config,
                "quantizer_path",
                f"codebooks/llama-2-7b_{bits}bit.xmad",
            )
        else:
            setattr(
                config,
                "quantizer_path",
                f"codebooks/llama-3-8b_{bits}bit.xmad",
            )
    if isinstance(args.window_length, int):
        setattr(config, "window_length", args.window_length)

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

    # ! CHECK: if using accelerate's "auto" device map can save more on memory
    model = AutoModelForCausalLM.from_pretrained(
        model_name, config=config, torch_dtype=dtype, device_map="auto"
    )

    # TODO: Figure out why this is kinda slow vs. using accelerate...?
    # model = AutoModelForCausalLM.from_pretrained(
    #     model_name, config=config, torch_dtype=dtype, device_map="cuda:0"
    # )

    if len(tokenizer) > model.get_input_embeddings().weight.shape[0]:
        print(
            "WARNING: Resizing the embedding matrix to match the tokenizer vocab size."
        )
        model.resize_token_embeddings(len(tokenizer))

    # Set padding side and pad token ID
    tokenizer.padding_side = "left"
    model.config.pad_token_id = tokenizer.pad_token_id

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

    num_dialogs = len(dialogs)
    print(f"Loaded {num_dialogs} dialogues...")

    # ! CHECK: if add_generation_prompt=True improves answer generation
    batch_inputs = [
        tokenizer.apply_chat_template(
            dialog, tokenize=False, add_generation_prompt=True
        )
        for dialog in dialogs
    ]

    # TODO: Add terminators to model.generate()
    terminators = [
        tokenizer.eos_token_id,
        tokenizer.convert_tokens_to_ids("<|eot_id|>"),
    ]

    # batch_sizes = [
    #     5,
    #     10,
    #     20,
    #     30,
    #     40,
    #     50,
    #     60,
    #     70,
    #     80,
    #     90,
    #     100,
    # ]  # working range on t8
    batch_sizes = [
        30
    ]

    memory_avg = []
    tokens_per_sec_avg = []
    time_to_first_token_avg = []
    responses_by_batch_size = defaultdict(list)

    os.makedirs(args.output_path, exist_ok=True)

    for batch_size in batch_sizes:
        print(f"\nProcessing with batch size: {batch_size}")

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

        # ! CHECK: if dynamic padding is better vs. pre-allocated tensors; no need to truncate b/c left-padding enabled
        with TorchTracemalloc() as tt:
            for i in range(0, num_dialogs, actual_batch_size):
                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()

                    # Generate responses
                    with torch.no_grad():
                        output_tokens = model.generate(
                            input_ids,
                            attention_mask=attention_mask,
                            max_new_tokens=args.num_new_tokens,
                            num_return_sequences=1,
                            do_sample=True,
                            temperature=args.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
                    total_tokens += output_tokens.numel()

                    if i == 0:
                        total_ttft = batch_time

                    # Decode the generated responses
                    decoded_outputs = tokenizer.batch_decode(
                        output_tokens, skip_special_tokens=True
                    )

                    # Store the responses
                    for j, response in enumerate(decoded_outputs):
                        original_dialog = dialogs[i + j]
                        formatted_response = format_response(
                            original_dialog, response
                        )
                        responses_by_batch_size[batch_size].append(
                            formatted_response
                        )

                    # ! CHECK: if Clearing CUDA cache after each batch works, but this can also cause fragmentation errors
                    torch.cuda.empty_cache()

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

        avg_memory = np.mean(TorchTracemalloc.track_memory_consumption)
        memory_avg.append(avg_memory)

        tokens_per_sec = total_tokens / total_time if total_time > 0 else 0
        tokens_per_sec_avg.append(tokens_per_sec)

        # Use actual_batch_size in calculations
        time_to_first_token = (
            total_ttft / actual_batch_size if actual_batch_size > 0 else 0
        )
        time_to_first_token_avg.append(time_to_first_token)

        print(f"Actual Batch Size Used: {actual_batch_size}")
        print(f"GPU Memory Consumption (MiB): {avg_memory:.2f} MiB")
        print(f"Tokens per Second: {tokens_per_sec:.2f}")
        print(f"TTFT (seconds): {time_to_first_token:.4f} seconds")

    for batch_size, responses in responses_by_batch_size.items():
        output_file = os.path.join(
            args.output_path, f"batch_{batch_size}_responses.json"
        )
        with open(output_file, "w") as f:
            json.dump(responses, f, indent=2)

    save_bar_chart(
        title="GPU Memory Consumption as a Function of Batch Size",
        x=batch_sizes,
        y=memory_avg,
        xlabel="Batch Size",
        ylabel="GPU Memory Consumption (MiB)",
        output_path=f"{args.output_path}/memory_usage.png",
    )

    save_bar_chart(
        title="Number of Tokens per Second as a Function of Batch Size",
        x=batch_sizes,
        y=tokens_per_sec_avg,
        xlabel="Batch Size",
        ylabel="Tokens per Second",
        output_path=f"{args.output_path}/tokens_per_second.png",
    )

    save_bar_chart(
        title="Time to First Token (TTFT) as a Function of Batch Size",
        x=batch_sizes,
        y=time_to_first_token_avg,
        xlabel="Batch Size",
        ylabel="TTFT (seconds)",
        output_path=f"{args.output_path}/time_to_first_token.png",
    )

    print("\nBenchmarking Results:")
    print(f"Batch Sizes: {batch_sizes}")
    print(f"GPU Memory Consumption (MiB): {memory_avg}")
    print(f"Tokens per Second: {tokens_per_sec_avg}")
    print(f"Time to First Token (seconds): {time_to_first_token_avg}")

    print(
        f"\nModel size: {args.model_size}, Max New Tokens: {args.num_new_tokens}, KV bits: {bits}"
    )
    print(f"Results and responses saved in: {args.output_path}")

except Exception as e:
    print(f"An error occurred during script execution: {str(e)}")