from contextlib import contextmanager, nullcontext
import datetime
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

from modules.audio_classifier import Model, HootConfig


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


INT16_MAX = np.iinfo(np.int16).max

data_dir = None
out_dir = None
train_file_name = "audio_16khz_tr.bin"
val_file_name = "audio_16khz_val.bin"
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 = 100
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 = 64  # if gradient_accumulation_steps > 1, this is the micro-batch size
# model
n_layers = 18
n_embd = 512
n_classes = 2
highfreq = None
window = 16_000 * 5
n_codebooks = 4
# adamw optimizer
learning_rate = 1e-4  # 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 = "float32"  # "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
# -----------------------------------------------------------------------------

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 True  # 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("loading data...")
train_data = np.memmap(os.path.join(data_dir, train_file_name), dtype=np.int16, mode="r")
val_data = np.memmap(os.path.join(data_dir, val_file_name), dtype=np.int16, mode="r")

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

bandwidth = round(n_codebooks / 2 * 1.5, 1)
encodec_model = EncodecModel.encodec_model_24khz()
encodec_model.set_target_bandwidth(bandwidth)
encodec_model.eval()
encodec_model.to(device)

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

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


@torch.no_grad()
def cycle_through_codec(array_batch):
    frames = encodec_model.encode(array_batch)[0][0]
    decoded_frames = encodec_model.decode([(frames, None)]).detach()
    return decoded_frames


def get_batch(split):
    # ix = np.arange(batch_size)
    data = train_data if split == "train" else val_data
    ix = np.random.randint(0, high=data.shape[0] - window, size=(batch_size,))
    x = torch.cat(
        [torch.from_numpy(data[i : i + window][None, None].astype(np.float32) / INT16_MAX) for i in ix],
        dim=0,
    ).to(device)
    # y = torch.from_numpy(np.array([True] *3 + [False] *4))
    y = torch.rand(batch_size, device=device) > 0.5
    x[y] = cycle_through_codec(x[y])
    x = x.squeeze(1)
    # y = (torch.rand(batch_size, device=device) > 0.5).long()
    y = y.long()
    signals_len = torch.ones(x.shape[0], dtype=torch.int64) * x.shape[-1]
    targets_len = torch.ones(x.shape[0], dtype=torch.int64)

    x, y, signals_len, targets_len = (
        x.to(device),
        y.to(device),
        signals_len.to(device),
        targets_len.to(device),
    )
    assert not torch.isnan(x).any().item()
    return x, signals_len, y, targets_len


@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ["train", "val"]:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, sl, Y, tl = get_batch(split)
            with ctx:
                _, loss = model(X, sl, Y, tl)
            losses[k] = loss.item()
        out[f"{split}/loss"] = losses.mean()
    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
        loss.backward()
        running_loss.append(loss_val)
    # clip the gradient
    if grad_clip != 0.0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    # step the optimizer and scaler if training in fp16
    optimizer.step()
    # 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}, "
            f"curr loss: {loss_val:.3f}, "
            f"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()