import os
from suno_utils.audio import Audio
import fire
import gc



from suno_utils.gpt.generation import (
    generate,
    models,
    load_dummy_model,
    GPTConfig,
)
import torch

import torch
import contextlib
import time
import pandas as pd
from tqdm import tqdm

CFG = 3  # multiply batch size by CFG
REP_RATE = 75


def run(n_batch=2, max_duration_s=20, profile=False):
    prof = (
        contextlib.nullcontext()
        if not profile
        else torch.profiler.profile(with_flops=True)
    )
    # get max memory allocated
    # print(f"Running with batch size {n_batch * CFG}")
    print(f"Current memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.synchronize()
    t0 = time.time()

    with prof:
        y = generate(
            GenerationConfig(
                text="hello",
                text_tags="edm",
                n_batch=n_batch,
                max_gen_duration_s=max_duration_s,
                allow_eos=False,
            )
        )
    t1 = time.time()

    max_mem = torch.cuda.max_memory_allocated() / 1e9
    print(f"Max memory allocated: {max_mem:.2f} GB")

    # if profile:
    #     prof.export_chrome_trace(f"profile_bs_{n_batch}_{max_duration_s}.json")

    time_elapsed = t1 - t0
    cfg = models["main_model"].config
    tokens_generated = (
        cfg.semantic_n_codebooks
        + (cfg.coarse_n_codebooks - 1) * cfg.coarse_shift_factor
        + len(y[1])
    )
    tps = tokens_generated / time_elapsed
    throughput = tps * n_batch
    return time_elapsed, tps, throughput, max_mem
    # min_bandwidth = n_params * tps  # Bandwidth to load model weights once per token
    # return y, time_elapsed, throughput, min_bandwidth, tps


tokenizer_path = "/home/victor/data/models/chirp_v2/tokenizer_60k.json"


def run_dummy(model_config, n_batch=2, max_duration_s=20):
    load_dummy_model(
        GPTConfig(
            n_head=model_config["n_head"],
            n_layer=model_config["n_layer"],
            n_embd=model_config["n_embd"],
        ),
        tokenizer_path,
    )
    try:
        return run(max_duration_s=max_duration_s, n_batch=n_batch)
    except RuntimeError:
        print(f"{model_config}, OOM")
        return None


def profile_shapes(gpu: str):
    data = []
    configs = [
        (40, 40, 5120),
        (25, 48, 1600),
        (32, 48, 2048),
        (32, 32, 4096),
    ]
    for n_head, n_layer, n_embd in tqdm(configs, desc="configs"):
        for duration in tqdm([30, 60, 80, 100, 120], desc="duration"):
            for batch_size in [1, 2]:
                print(
                    f"Running with n_head={n_head}, n_layer={n_layer}, n_embd={n_embd}. Generating...."
                )
                try:
                    # clear cache
                    torch.cuda.empty_cache()
                    load_dummy_model(
                        GPTConfig(n_head=n_head, n_layer=n_layer, n_embd=n_embd),
                        tokenizer_path,
                    )
                    time_elapsed, tps, throughput, max_mem = run(
                        max_duration_s=duration, n_batch=batch_size
                    )
                except:
                    print(f"n_head={n_head}, n_layer={n_layer}, n_embd={n_embd}, OOM")
                    # free up memory
                    del models["main_model"]
                    torch.cuda.empty_cache()

                    break
                data.append(
                    {
                        "gpu": gpu,
                        "n_head": n_head,
                        "n_layer": n_layer,
                        "n_embd": n_embd,
                        "duration": duration,
                        "batch_size": batch_size,
                        "time_elapsed": time_elapsed,
                        "tps": tps,
                        "throughput": throughput,
                        "max_mem": max_mem,
                    }
                )
                print(data[-1])
        #         break
        #     break
        # break

    df = pd.DataFrame(data)
    df.to_csv(f"gpt_throughput_{gpu}.csv")
    print(df)


def run_7b(n_batch=2, max_duration_s=20, profile=False):
    print(f"Current memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.synchronize()
    t0 = time.time()

    y = generate_audio(
        GenerationConfig(
            text="hello " * 3000,
            text_tags="edm",
            text_neg_control_tags="classical",
            n_batch=n_batch,
            max_gen_duration_s=max_duration_s,
            allow_eos=False,
        ),
        return_raw_arrays=True,
    )
    t1 = time.time()

    max_mem = torch.cuda.max_memory_allocated() / 1e9
    print(f"Max memory allocated: {max_mem:.2f} GB")

    time_elapsed = t1 - t0
    cfg = models["main_model"]["model"].config
    tokens_generated = (
        cfg.semantic_shift_factor
        + (cfg.coarse_n_codebooks - 1) * cfg.coarse_shift_factor
        + y[1][0].shape[0]
    )
    # print(f"tokens generated: {tokens_generated}")
    tps = tokens_generated / time_elapsed
    throughput = tps * n_batch
    return time_elapsed, tps, throughput, max_mem


def profile(gpu: str, duration: int, model: str):
    global generate_audio, GenerationConfig

    assert model in ["7B", "3B"]
    if model == "7B":
        from suno_utils.gpt.chirp_v2_5 import (
            preload_models,
            generate_audio,
            GenerationConfig,
        )
        preload_models(
            codec_ckpt_path="s3://suno-data/georg/models/codec/dac_2c_25x12.pt",
            gpt_ckpt_path="s3://suno-data/georg/checkpoints/chirp_v2_5/7b.pt",
        )
    elif model == "3B":
        from suno_utils.gpt.chirp_v2 import (
            preload_models,
            generate_audio,
            GenerationConfig,
        )
        preload_models(
            codec_ckpt_path="/home/victor/data/models/chirp_v2/dac_2c_25x8.pt",
            gpt_ckpt_path="/app/suno/checkpoints/2023-10-17_16-01-32/last_ckpt_infer.pt",
        )


    # for batch_size in [2]:
    for batch_size in [1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 16, 32, 64, 128]:
        try:
            gc.collect()
            torch.cuda.empty_cache()
            time_elapsed, tps, throughput, max_mem = run_7b(
                max_duration_s=duration, n_batch=batch_size
            )
        except Exception as e:
            print(e)
            print(f"OOM with batch size {batch_size}, duration {duration}")
            break

        j = {
            "model": model,
            "gpu": gpu,
            "duration": duration,
            "batch_size": batch_size,
            "time_elapsed": time_elapsed,
            "tps": tps,
            "throughput": throughput,
            "max_mem": max_mem,
        }
        print(j)

        with open(f"7b_throughput.jsonl", "a+") as f:
            f.write(str(j) + "\n")



if __name__ == "__main__":
    fire.Fire(profile)