import time
import os
import re
import json
import random
from contextlib import nullcontext

import numpy as np
from tokenizers import Tokenizer
import torch

from modules.semantic import GPTConfig, GPT

from torch.profiler import profile, record_function, ProfilerActivity


activities = [ProfilerActivity.CPU]
if torch.cuda.is_available():
    activities.append(ProfilerActivity.CUDA)


data_dir = None
do_full_profile = True
eval_steps = 50

dataset_names = "youtube,foreign,genius,podcasts"
filename_pattern = "semantic_{dataset}_{set}.bin"
filename_pattern_meta = "semantic_{dataset}_meta_{set}.json"
filename_tokenizer = "bert_tokenizer.json"

text_codebook_size = 119_547
text_vocab_size = 119_552
assert text_vocab_size - text_codebook_size >= 3  # to accommodate pad tokens etc
text_pad_token = text_codebook_size
semantic_codebook_size = 10_000
semantic_vocab_size = 10_048
n_semantic_codebooks = 8
assert semantic_vocab_size - semantic_codebook_size >= 3  # to accommodate pad tokens etc
semantic_pad_token = semantic_codebook_size
infer_token = semantic_codebook_size + 1
custom_seed_offset = 0

gradient_accumulation_steps = 1  # used to simulate larger batch sizes
batch_size = 12  # if gradient_accumulation_steps > 1, this is the micro-batch size
block_size = 4096

# model
n_layer = 12
n_head = 12
n_embd = 768
dropout = 0.0  # for pretraining 0 is good, for finetuning try 0.1+
bias = False  # do we use bias inside LayerNorm and Linear layers?

# adamw optimizer
learning_rate = 1e-3  # max learning rate
max_iters = 100000  # total number of training iterations
weight_decay = 1e-1
beta1 = 0.9
beta2 = 0.95
grad_clip = 1.0  # clip gradients at this value, or disable if == 0.0

# system
device = "cuda"  # examples: "cpu", "cuda", "cuda:0", "cuda:1" etc., or try "mps" on macbooks
dtype = "bfloat16"  # "float32", "bfloat16", or "float16" (implements a GradScaler)
compile = True  # use PyTorch 2.0 to compile the model to be faster

# -----------------------------------------------------------------------------
config_keys = [
    k for k, v in globals().items() if not k.startswith("_") and isinstance(v, (int, float, bool, str))
]
exec(open("configurator.py").read())  # overrides from command line or config file
config = {k: globals()[k] for k in config_keys}  # will be useful for logging
# -----------------------------------------------------------------------------


seed_offset = 0
seed_offset += custom_seed_offset
torch.manual_seed(1337 + seed_offset)
random.seed(6006 + seed_offset)
torch.backends.cuda.matmul.allow_tf32 = True  # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = True  # allow tf32 on cudnn
device_type = "cuda" if "cuda" in device else "cpu"  # for later use in torch.autocast
# note: float16 data type will automatically use a GradScaler
ptdtype = {
    "float32": torch.float32,
    "bfloat16": torch.bfloat16,
    "float16": torch.float16,
}[dtype]
ctx = (
    nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
)


# load data
print("loading data...")
dataset_names = [s.strip() for s in dataset_names.split(",")]

val_data = []
val_data_metas = []
val_data_weights = []
for name in dataset_names:
    fn = filename_pattern.replace("{dataset}", name).replace("{set}", "val")
    mm = np.memmap(os.path.join(data_dir, fn), dtype=np.uint16, mode="r")
    assert mm[: 10 * block_size].max() <= semantic_codebook_size
    mm = mm.reshape(-1, 8, block_size)
    val_data.append(mm)
    val_data_weights.append(len(mm))
    fn_metas = filename_pattern_meta.replace("{dataset}", name).replace("{set}", "val")
    with open(os.path.join(data_dir, fn_metas)) as f:
        metas = json.load(f)
    val_data_metas.append(metas)
    assert len(mm) == len(metas)

tokenizer = Tokenizer.from_file(os.path.join(data_dir, filename_tokenizer))


def tokenize(text):
    return tokenizer.encode(text, add_special_tokens=False).ids


def get_sample(dataset_idx=None, n_text_context=256, use_fake_data=False):
    if use_fake_data:
        x_text_arr = np.zeros(256).astype(np.int64) + 119_547
        y_semantic_arr = np.zeros((8, 4096 - 256)).astype(np.int64) + 10_000
        x_semantic_arr = y_semantic_arr.copy()
        x_semantic_arr[:, 0] = 10_001
        return x_text_arr, x_semantic_arr, y_semantic_arr

    if dataset_idx is None:
        dataset_idx = random.choices(list(range(len(val_data))), weights=val_data_weights, k=1)[0]
    data = val_data[dataset_idx]
    data_metas = val_data_metas[dataset_idx]
    idx = random.randint(0, len(data) - 1)
    # build semantic
    y_semantic_arr_list = []
    for n in range(n_semantic_codebooks):
        arr = np.pad(
            data[idx][n, : block_size - n_text_context - n],
            (n, 0),
            constant_values=semantic_pad_token,
            mode="constant",
        ).astype(np.uint16)
        y_semantic_arr_list.append(arr)
    y_semantic_arr = np.vstack(y_semantic_arr_list)
    x_semantic_arr = np.hstack(
        [np.array([[infer_token]] * n_semantic_codebooks), y_semantic_arr[:, :-1]]
    )
    assert x_semantic_arr.shape == (n_semantic_codebooks, block_size - n_text_context)
    assert y_semantic_arr.shape == (n_semantic_codebooks, block_size - n_text_context)
    # build text
    meta = data_metas[idx]
    x_text_arr = np.array([text_pad_token] * n_text_context)
    # metas can be text, score, tags, language
    # and also dataset
    p_cond = 0.3
    pretext_segments = []
    if random.random() < p_cond:
        pretext_segments.append(f"dataset:{dataset_names[dataset_idx].lower()}")
    if "language" in meta and random.random() < p_cond:
        pretext_segments.append(f"language:{meta['language'].lower()}")
    if "score" in meta and random.random() < p_cond:
        pretext_segments.append(f"score:{meta['score']}")
    if "tags" in meta and len(meta["tags"]) > 0 and random.random() < p_cond:
        tags = meta["tags"]
        if len(tags) > 0 and len(tags[0]) == 1:
            # bugfix for ["p", "o", "p"]
            tags = "".join(tags).split(",")
        random.shuffle(tags)
        tags = tags[: random.randint(1, len(tags))]
        tag_str = ",".join(tags)
        pretext_segments.append(f"tags:{tag_str.lower()}")
    text = ""
    if len(pretext_segments) > 0:
        text = "{" + ";".join(pretext_segments) + "}"
    if "text" in meta and random.random() < p_cond:
        if random.random() < 0.75:
            # remove newlines completely (eg lyrics)
            tts_text = re.sub(r"\s+", " ", meta["text"])
        else:
            tts_text = meta["text"]
        text += " " + tts_text.strip()
    text = text.strip()
    text_tokens = tokenize(text)
    x_text_arr[: len(text_tokens)] = text_tokens[: len(x_text_arr)]
    assert x_text_arr.shape == (n_text_context,)
    return x_text_arr, x_semantic_arr, y_semantic_arr


def get_batch(dataset_idx=None):
    # TODO: make this adjustable
    n_text_context = 256
    x_text_list = []
    x_semantic_list = []
    y_list = []
    for _ in range(batch_size):
        x_text, x_semantic, y = get_sample(dataset_idx=dataset_idx, n_text_context=n_text_context)
        x_text_list.append(torch.from_numpy(x_text.astype(np.int64)))
        x_semantic_list.append(torch.from_numpy(x_semantic.astype(np.int64)))
        y_list.append(torch.from_numpy(y.astype(np.int64)))
    x_text = torch.stack(x_text_list)
    x_semantic = torch.stack(x_semantic_list)
    y = torch.stack(y_list)
    if device_type == "cuda":
        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
        x_text, x_semantic, y = (
            x_text.pin_memory().to(device),
            x_semantic.pin_memory().to(device),
            y.pin_memory().to(device),
        )
    else:
        x_text, x_semantic, y = x_text.to(device), x_semantic.to(device), y.to(device)
    del x_text_list, x_semantic_list, y_list
    return x_text, x_semantic, y


# init these up here, can override if init_from="resume" (i.e. from a checkpoint)
iter_num = 0

# model init
model_args = dict(
    n_layer=n_layer,
    n_head=n_head,
    n_embd=n_embd,
    block_size=block_size,
    bias=bias,
    dropout=dropout,
    text_vocab_size=text_vocab_size,
    semantic_vocab_size=semantic_vocab_size,
    n_semantic_codebooks=n_semantic_codebooks,
)

# init a new model from scratch
print("Initializing a new model from scratch")
gptconf = GPTConfig(**model_args)
model = GPT(gptconf)
model.to(device)

# initialize a GradScaler. If enabled=False scaler is a no-op
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == "float16"))

# optimizer
optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)

# compile the model
if compile:
    print("compiling the model... (takes a ~minute)")
    model = torch.compile(model)

X_text, X_semantic, Y = get_batch()

print("Warmup...")
model.train()
for n in range(50):
    for nn in range(gradient_accumulation_steps):
        with ctx:
            loss = model(X_text, X_semantic, targets=Y)
            loss = loss / gradient_accumulation_steps
        scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)
torch.cuda.synchronize()

print("Checking dataload overhead...")
model.train()
t0 = time.time()
for n in range(eval_steps):
    for nn in range(gradient_accumulation_steps):
        with ctx:
            loss = model(X_text, X_semantic, targets=Y)
            loss = loss / gradient_accumulation_steps
        X_text, X_semantic, Y = get_batch()
        scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)
torch.cuda.synchronize()
print(round(time.time() - t0, 1), "seconds with dataload")
t0 = time.time()
for n in range(eval_steps):
    for nn in range(gradient_accumulation_steps):
        with ctx:
            loss = model(X_text, X_semantic, targets=Y)
            loss = loss / gradient_accumulation_steps
        scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)
torch.cuda.synchronize()
print(round(time.time() - t0, 1), "seconds without dataload")

if do_full_profile:
    print("Profiling Infer...")
    model.eval()
    with profile(activities=activities, record_shapes=False) as prof_infer:
        with record_function(" Infer"):
            for n in range(eval_steps):
                for nn in range(gradient_accumulation_steps):
                    with ctx, torch.no_grad():
                        _ = model(X_text, X_semantic)
            torch.cuda.synchronize()
    print(prof_infer.key_averages().table(sort_by="cuda_time_total", row_limit=10))

    print("Profiling Train...")
    model.train()
    with profile(activities=activities, record_shapes=False) as prof_train:
        with record_function(" Train"):
            for n in range(eval_steps):
                for nn in range(gradient_accumulation_steps):
                    with ctx:
                        loss = model(X_text, X_semantic, targets=Y)
                        loss = loss / gradient_accumulation_steps
                    scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad(set_to_none=True)
            torch.cuda.synchronize()
    print(prof_train.key_averages().table(sort_by="cuda_time_total", row_limit=10))

    with open("out.txt", "w") as f:
        f.write(prof_infer.key_averages().table(sort_by="cuda_time_total"))
        f.write("\n" * 5)
        f.write(prof_train.key_averages().table(sort_by="cuda_time_total"))