from contextlib import contextmanager, nullcontext
import datetime
import json
import logging
import math
import os
import random
import time

from encodec import EncodecModel
import numpy as np
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group
import torchaudio

from hoot.model import Hoot, HootConfig, Tokenizer, collate_fn


@contextmanager
def suppress_logging(highest_level=logging.CRITICAL):
    previous_level = logging.root.manager.disable
    logging.disable(highest_level)
    try:
        yield
    finally:
        logging.disable(previous_level)


data_dir = None
out_dir = None
dataset_names = "main"
filename_pattern = "hoot_{dataset}_{set}.bin"
filename_pattern_meta = "hoot_{dataset}_metas_{set}.jsonl"
filename_tokenizer = "hoot_tokenizer.model"
max_duration_s = 6 * 60
debug_val_only = False
preload_checkpoint = None
preload_optimizer = False
preload_strict = True
suppress_compile_warnings = True
# eval items
custom_seed_offset = 0
eval_interval = 2000
log_interval = 25
eval_iters = 600
eval_only = False  # if True, script exits right after the first eval
debug_gradients = False
always_save_checkpoint = True  # if True, always save a checkpoint after each eval
init_from = "scratch"  # "scratch" or "resume" or "gpt2*"
# wandb logging
wandb_log = False
wandb_project = "suno-test"
wandb_run_name = "test"
# data
gradient_accumulation_steps = 1  # used to simulate larger batch sizes
batch_size = 12  # if gradient_accumulation_steps > 1, this is the micro-batch size
# model
n_layers = 18
n_embd = 512
n_classes = 1024
# 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
# learning rate decay settings
decay_lr = True  # whether to decay the learning rate
warmup_iters = 1000  # how many steps to warm up for
lr_decay_iters = None  # should be ~= max_iters per Chinchilla
min_lr = 1e-5  # minimum learning rate, should be ~= learning_rate/10 per Chinchilla
# DDP settings
backend = "nccl"  # "nccl", "gloo", etc.
# 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 = False  # 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
# -----------------------------------------------------------------------------

dataset_names = [s.strip() for s in dataset_names.split(",")]
eval_iters = int(eval_iters * gradient_accumulation_steps)
if lr_decay_iters is None:
    lr_decay_iters = max_iters

# various inits, derived attributes, I/O setup
ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
if ddp:
    init_process_group(backend=backend)
    ddp_rank = int(os.environ["RANK"])
    ddp_local_rank = int(os.environ["LOCAL_RANK"])
    device = f"cuda:{ddp_local_rank}"
    torch.cuda.set_device(device)
    master_process = ddp_rank == 0  # this process will do logging, checkpointing etc.
    seed_offset = ddp_rank  # each process gets a different seed
else:
    # if not ddp, we are running on a single gpu, and one process
    master_process = True
    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)
)

# logging
if wandb_log and master_process:
    import wandb

    wandb.init(project=wandb_project, name=wandb_run_name, config=config)

date_time_str = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
out_dir = os.path.join(out_dir, date_time_str)
if master_process:
    os.makedirs(out_dir, exist_ok=True)
    print(f"logging checkpoint here: {out_dir}")

# load data
print(f"loading data...")
train_data = []
train_data_metas = []
train_data_weights = []
for name in dataset_names:
    fn = filename_pattern.replace("{dataset}", name)
    fn = fn.replace("{set}", "val" if debug_val_only else "tr")
    mm = np.memmap(os.path.join(data_dir, fn), dtype=np.uint16, mode="r")
    assert mm[:10_000].max() <= 1024
    mm = mm.reshape(-1, 8).T
    train_data.append(mm)
    train_data_weights.append(mm.shape[-1])
    fn_metas = filename_pattern_meta.replace("{dataset}", name)
    fn_metas = fn_metas.replace("{set}", "val" if debug_val_only else "tr")
    metas = []
    with open(os.path.join(data_dir, fn_metas)) as f:
        for line in f:
            line = line.strip()
            if len(line) == 0:
                continue
            m = json.loads(line)
            if max_duration_s is not None and m["duration_s"] > max_duration_s:
                continue
            metas.append(m)
    train_data_metas.append(metas)
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_000].max() <= 1024
    mm = mm.reshape(-1, 8).T
    val_data.append(mm)
    val_data_weights.append(mm.shape[-1])
    fn_metas = filename_pattern_meta.replace("{dataset}", name).replace("{set}", "val")
    metas = []
    with open(os.path.join(data_dir, fn_metas)) as f:
        for line in f:
            line = line.strip()
            if len(line) == 0:
                continue
            m = json.loads(line)
            if max_duration_s is not None and m["duration_s"] > max_duration_s:
                continue
            metas.append(m)
    val_data_metas.append(metas)

# model init
print("Initializing a new model from scratch")
model_args = dict(
    n_layers=n_layers,
    n_embd=n_embd,
    n_classes=n_classes,
)
conf = HootConfig(**model_args)
model = Hoot(conf)
model.to(device)

encodec_model = EncodecModel.encodec_model_24khz()
encodec_model.set_target_bandwidth(6.0)
encodec_model.eval()
encodec_model.to(device)

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

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

# 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
)

# load checkpoint
if preload_checkpoint is not None:
    print("preloading checkpoint")
    cur_state_dict = model.state_dict()
    checkpoint = torch.load(preload_checkpoint, map_location=device)
    state_dict = checkpoint["model"] if "model" in checkpoint else checkpoint
    # fix the keys of the state dictionary :(
    # honestly no idea how checkpoints sometimes get this prefix, have to debug more
    unwanted_prefix = "_orig_mod."
    for k, v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix) :]] = state_dict.pop(k)
    # do some stuff in case nemo checkpoint
    if "preprocessor.featurizer.window" in state_dict:
        state_dict["featurizer._mel_spec_extractor.spectrogram.window"] = state_dict[
            "preprocessor.featurizer.window"
        ]
        del state_dict["preprocessor.featurizer.window"]
    if "preprocessor.featurizer.fb" in state_dict:
        state_dict["featurizer._mel_spec_extractor.mel_scale.fb"] = torch.swapaxes(
            state_dict["preprocessor.featurizer.fb"][0], 0, 1
        )
        del state_dict["preprocessor.featurizer.fb"]
    model.load_state_dict(state_dict, strict=preload_strict)
    if preload_optimizer:
        print("preloading optimizer")
        optimizer.load_state_dict(checkpoint["optimizer"])
        iter_num = checkpoint["iter_num"]
        best_val_loss = checkpoint["best_val_loss"]
    del cur_state_dict, state_dict, checkpoint

# compile the model
if compile:
    print("compiling the model... (takes a ~minute)")
    compile_ctx = suppress_logging if suppress_compile_warnings else nullcontext
    with compile_ctx():
        model = torch.compile(model)  # requires PyTorch 2.0

# wrap model into DDP container
if ddp:
    model = DDP(model, device_ids=[ddp_local_rank])
raw_model = model.module if ddp else model  # unwrap DDP container if needed


def get_sample(split, dataset_idx=None):
    # x = torch.zeros(16_000*60*5, dtype=torch.float32).to(device)
    # y = torch.zeros(10, dtype=torch.long).to(device)
    # x = torch.from_numpy(
    #     np.load("/home/georg/notebooks/gpt/tmp/halo.npy").astype(np.float32)
    # ).to(device)
    # y = torch.from_numpy(
    #     np.load("/home/georg/notebooks/gpt/tmp/halo_lyrics.npy").astype(np.int64)
    # ).to(device)
    if split == "train":
        if dataset_idx is None:
            dataset_idx = random.choices(
                list(range(len(train_data))), weights=train_data_weights, k=1
            )[0]
        data = train_data[dataset_idx]
        data_metas = train_data_metas[dataset_idx]
    else:
        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_metas) - 1)
    meta = data_metas[idx]
    encodec_data = torch.from_numpy(
        data[:, meta["start_idx"] : meta["end_idx"]].astype(np.int64)
    )
    y = tokenizer.encode(meta["text_norm"])
    encodec_data, y = encodec_data.to(device), y.to(device)
    with torch.no_grad():
        x = encodec_model.quantizer.decode(encodec_data[None].transpose(0, 1))
        x = encodec_model.decoder(x)[0].detach()
    x = torchaudio.functional.resample(x, 24_000, 16_000)[0]
    # clip to max space we have for logits
    y = y[: int(round(x.shape[-1] / 16_000 * 10))]
    return x, y


def get_batch(split, dataset_idx=None):
    x_list = []
    y_list = []
    for _ in range(batch_size):
        x, y = get_sample(split, dataset_idx=dataset_idx)
        x_list.append(x)
        y_list.append(y)
    x, x_len = collate_fn(x_list, fixed_len=int(round(max_duration_s * 16_000)))
    y, y_len = collate_fn(y_list, fixed_len=int(round(max_duration_s * 16_000)))
    del x_list, y_list
    return x, x_len, y, y_len


# helps estimate an arbitrarily accurate loss over either split using many batches
@torch.no_grad()
def estimate_loss():
    n_loss_modalities = len(train_data) + len(val_data)
    effective_eval_iters = int(round(eval_iters / n_loss_modalities))
    out = {}
    model.eval()
    for split in ["train", "val"]:
        split_losses = []
        datatype_losses = []
        n_datasets = len(train_data) if split == "train" else len(val_data)
        for dataset_idx in range(n_datasets):
            losses = []
            for k in range(effective_eval_iters):
                X, X_len, Y, Y_len = get_batch(split, dataset_idx=dataset_idx)
                with ctx:
                    loss = model(X, X_len, targets=Y, targets_len=Y_len)
                loss = loss.item()
                losses.append(loss)
                datatype_losses.append(loss)
                split_losses.append(loss)
            out[f"{split}/loss_{dataset_names[dataset_idx]}"] = np.mean(losses)
        out[f"{split}/loss"] = np.mean(datatype_losses)
    model.train()
    return out


# learning rate decay scheduler (cosine with warmup)
def get_lr(it):
    # 1) linear warmup for warmup_iters steps
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    # 2) if it > lr_decay_iters, return min learning rate
    if it > lr_decay_iters:
        return min_lr
    # 3) in between, use cosine decay down to min learning rate
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # coeff ranges 0..1
    return min_lr + coeff * (learning_rate - min_lr)


# training loop
if master_process:
    print("Training...")
t0 = time.time()
t00 = time.time()
local_iter_num = 0  # number of iterations in the lifetime of this process
running_loss = []
while True:
    # determine and set the learning rate for this iteration
    lr = get_lr(iter_num) if decay_lr else learning_rate
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

    # evaluate the loss on train/val sets and write checkpoints
    if iter_num % eval_interval == 0 and master_process:
        time_since_last_loss = time.time() - t00
        t00 = time.time()
        print("estimating loss...")
        losses = estimate_loss()
        estimation_time = time.time() - t00
        eval_time_pct = np.clip(estimation_time / time_since_last_loss * 100, 0, 100)
        print(
            f"loss estimation took {estimation_time:.1f} seconds. ({eval_time_pct:.1f}% of loop)"
        )
        print(
            f"step {iter_num}: train loss {losses['train/loss']:.4f},"
            f" val loss {losses['val/loss']:.4f}"
        )
        if wandb_log:
            log_dict = {
                "iter": iter_num,
                "lr": lr,
            }
            for k, v in losses.items():
                log_dict[k] = v
            wandb.log(log_dict)
        if losses["val/loss"] < best_val_loss or always_save_checkpoint:
            if iter_num > 0:
                checkpoint = {
                    "model": raw_model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "model_args": model_args,
                    "iter_num": iter_num,
                    "best_val_loss": losses["val/loss"],
                    "config": config,
                }
                print(f"saving checkpoint to {out_dir}")
                if losses["val/loss"] < best_val_loss:
                    torch.save(checkpoint, os.path.join(out_dir, "best_ckpt.pt"))
                if always_save_checkpoint:
                    torch.save(checkpoint, os.path.join(out_dir, "last_ckpt.pt"))
            if losses["val/loss"] < best_val_loss:
                best_val_loss = losses["val/loss"]

    # end if eval test only
    if iter_num == 0 and eval_only:
        print("eval test done.")
        break

    # forward backward update, with optional gradient accumulation to simulate larger batch size
    # and using the GradScaler if data type is float16
    for micro_step in range(gradient_accumulation_steps):
        if ddp:
            # in DDP training we only need to sync gradients at the last micro step.
            # the official way to do this is with model.no_sync() context manager, but
            # I really dislike that this bloats the code and forces us to repeat code
            # looking at the source of that context manager, it just toggles this variable
            model.require_backward_grad_sync = (
                micro_step == gradient_accumulation_steps - 1
            )
        X, X_len, Y, Y_len = get_batch("train")
        with ctx:
            loss = model(X, X_len, targets=Y, targets_len=Y_len)
            loss_val = loss.item()  # loss as float. note: this is a CPU-GPU sync point
            loss = loss / gradient_accumulation_steps

        if debug_gradients and wandb_log and master_process:
            d = {
                "iter": iter_num,
                "debug_loss": loss_val,
            }
            grads = []
            for name, param in model.named_parameters():
                if param.grad is not None:
                    grads.append(param.grad.norm().item())
            if len(grads) > 0:
                d["debug_grads"] = np.mean(grads)
                wandb.log(d)

        # backward pass, with gradient scaling if training in fp16
        scaler.scale(loss).backward()
        running_loss.append(loss_val)
    # clip the gradient
    if grad_clip != 0.0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    # step the optimizer and scaler if training in fp16
    scaler.step(optimizer)
    scaler.update()
    # flush the gradients as soon as we can, no need for this memory anymore
    optimizer.zero_grad(set_to_none=True)

    # timing and logging
    t1 = time.time()
    dt = t1 - t0
    t0 = t1
    if iter_num % log_interval == 0 and master_process:
        avg_loss = np.mean(running_loss)
        running_loss = []
        print(f"iter {iter_num}: avg_loss {avg_loss:.3f}, step_time {dt*1000:.1f}ms")
    iter_num += 1
    local_iter_num += 1

    # termination conditions
    if iter_num > max_iters:
        break

if ddp:
    destroy_process_group()