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

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 torcheval.metrics as metrics
from torch.distributed import barrier, is_initialized
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from suno_utils.audio import Audio

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

# The hoot sample rate is 16kHz
SAMPLE_RATE = 16_000
enable_audio_augmentation = False


def dist_barrier():
    if is_initialized():
        barrier()


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


# Create a custom Dataset implementation
class HootDataset(Dataset):
    def __init__(
        self,
        data_metas,
        tokenizer,
        lyrics_key="text",
        max_duration=None,
        cache_dir=None,
    ):
        self.data_metas = data_metas
        self.tokenizer = tokenizer
        self.lyrics_key = lyrics_key
        self.max_duration = max_duration
        self.cache_dir = cache_dir

        # # Pre-compute tokens for all samples (only done once)
        # print_with_time("Pre-computing tokens for all samples...")
        # for idx, meta in tqdm.tqdm(
        #     enumerate(self.data_metas), total=len(self.data_metas)
        # ):
        #     if "tokens" not in meta:
        #         meta["tokens"] = tokenizer.encode(meta[lyrics_key])

    def __len__(self):
        return len(self.data_metas)

    def __getitem__(self, idx):
        meta = self.data_metas[idx]
        audio_file_path = meta["audio_filepath"]

        try:
            # Load audio using Audio class
            audio = Audio.from_file(audio_file_path)
            waveform = torch.from_numpy(audio.convert(SAMPLE_RATE, 2, 1).array_float)

            # The Audio class already handles:
            # - Converting to mono (last parameter = 1)
            # - Resampling to 16kHz (SAMPLE_RATE)
            # So we don't need the manual conversion and resampling anymore

            # Ensure waveform has proper shape for augmentation
            if waveform.ndim == 1:
                waveform = waveform.unsqueeze(0)  # Add channel dimension if needed

            if enable_audio_augmentation:
                # Augment with random silence or white noise padding (20% chance)
                random_value = random.random()
                if random_value < 0.2:
                    # Randomly choose between silence and white noise
                    noise_type = random.choice(["silence", "white_noise"])

                    # Randomly choose to pad at start, end, or both
                    pad_choice = random.choice(["start", "end", "both"])

                    def create_padding(duration):
                        """Create padding of specified duration"""
                        samples = int(duration * 16_000)
                        # Match the number of channels in the waveform
                        num_channels = waveform.shape[0]
                        if noise_type == "silence":
                            return torch.zeros(num_channels, samples)
                        else:  # white_noise
                            amplitude = random.uniform(0.01, 0.05)
                            return torch.randn(num_channels, samples) * amplitude

                    if pad_choice == "start":
                        padding = create_padding(random.uniform(0, 1.0))
                        waveform = torch.cat([padding, waveform], dim=-1)
                    elif pad_choice == "end":
                        padding = create_padding(random.uniform(0, 1.0))
                        waveform = torch.cat([waveform, padding], dim=-1)
                    else:  # both
                        start_padding = create_padding(random.uniform(0, 1.0))
                        end_padding = create_padding(random.uniform(0, 1.0))
                        waveform = torch.cat(
                            [start_padding, waveform, end_padding], dim=-1
                        )
                # Randomly crop 0 to 1 second from the beginning (20% chance)
                elif random_value < 0.4:
                    crop_duration = random.uniform(0, 1.0)
                    crop_samples = int(crop_duration * 16_000)
                    if waveform.shape[-1] > crop_samples:
                        waveform = waveform[..., crop_samples:]
                # randomly crop 0 to 1 second from the end (20% chance)
                elif random_value < 0.6:
                    crop_duration = random.uniform(0, 1.0)
                    crop_samples = int(crop_duration * 16_000)
                    if waveform.shape[-1] > crop_samples:
                        waveform = waveform[..., :-crop_samples]

            # Get tokens
            # if "tokens" in meta:
            #     y = meta["tokens"]  # Use pre-computed tokens
            # else:
            # Compute tokens on-the-fly if not already present
            cleaned_text = meta[self.lyrics_key].strip()
            encoded_array = self.tokenizer.encode(cleaned_text)
            repeat_ratio = max(np.bincount(encoded_array)) / (
                encoded_array.shape[0] + 1
            )
            # this is to prevent the loss from exploding
            if repeat_ratio > 0.35:
                raise ValueError(
                    f"Repeat ratio is too high: {repeat_ratio}, idx: {idx}"
                )
            y = encoded_array
            meta["tokens"] = y  # Update meta with computed tokens

            # Process audio
            x = waveform
            if x.shape[0] > 1:
                x = x.mean(axis=0)
            else:
                x = x.reshape(-1)

            # Clip to reasonable length
            y_np = y[: int(round(x.shape[-1] / 16_000 * 10))]
            y = torch.from_numpy(
                y_np.copy()
            )  # Make a copy to ensure memory is contiguous

            return x, y, idx
        except Exception as e:
            print(f"Error loading sample {idx}: {e}")
            # Return a small dummy sample to avoid batch failures
            return torch.zeros(16000), torch.zeros(1, dtype=torch.long), idx


# Collation function for batching
def hoot_collate_fn(batch):
    # Filter out any None values or failed loads
    batch = [b for b in batch if b is not None and b[0].numel() > 0]
    if not batch:
        return None

    x_list, y_list, idxs = zip(*batch)
    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)))

    return x, x_len, y, y_len, idxs


# Data prefetcher for asynchronous GPU transfer
class DataPrefetcher:
    def __init__(self, loader, device):
        self.loader = iter(loader)
        self.device = device
        self.stream = torch.cuda.Stream()
        self.next_batch = None
        self.preload()

    def preload(self):
        try:
            # Get next batch from loader
            self.next_batch = next(self.loader)
        except StopIteration:
            self.next_batch = None
            return

        # Transfer to GPU asynchronously
        if self.next_batch is not None:
            with torch.cuda.stream(self.stream):
                for i in range(
                    len(self.next_batch) - 1
                ):  # Skip idxs which is the last element
                    if isinstance(self.next_batch[i], torch.Tensor):
                        self.next_batch[i] = self.next_batch[i].to(
                            self.device, non_blocking=True
                        )

    def next(self):
        # Wait for the transfer to complete
        torch.cuda.current_stream().wait_stream(self.stream)

        # Get current batch and start preloading next one
        batch = self.next_batch

        # Important: Clear the reference before preloading to avoid holding two batches
        self.next_batch = None

        # Start preloading next batch
        self.preload()

        return batch

    def __del__(self):
        # Explicit cleanup when the prefetcher is deleted
        self.next_batch = None
        self.loader = None


master_addr = "localhost"
master_port = 12835

out_dir = None
# train_input_path = "/home/tony/Data/Hoot/multi_long_filtered_train_manifest.json"
# val_input_path = "/home/tony/Data/Hoot/multi_long_filtered_test_manifest.json"
# train_input_path = "/home/tony/Data/Hoot/multi_balanced_train_manifest.json"
# val_input_path = "/home/tony/Data/Hoot/multi_balanced_test_manifest.json"
# train_input_path = "/home/tony/Data/Hoot/all_train_manifest.json"
# val_input_path = "/home/tony/Data/Hoot/all_test_manifest.json"
# train_input_path = "/home/tony/Data/Hoot/multi_filtered_train_manifest.json"
# val_input_path = "/home/tony/Data/Hoot/multi_filtered_test_manifest.json"
# train_input_path = "/home/tony/Data/Hoot/en_train_manifest_norm.json"
# val_input_path = "/home/tony/Data/Hoot/en_test_manifest_norm.json"
train_input_path = "/home/tony/Data/Hoot/v4_t1_cer_50_long_train.json"
train_suno_input_path = None
val_input_path = "/home/tony/Data/Hoot/v4_t1_cer_50_long_test.json"
# tokenizer_path = (
#     "/home/tony/Work/tony/hoot/tokenizers/multi/tokenizer_spe_bpe_v5120/tokenizer.model"
# )
# tokenizer_path = (
#     "/home/tony/Work/tony/hoot/tokenizers/multi_filtered/tokenizer_spe_bpe_v5120/tokenizer.model"
# )
# tokenizer_path = "/home/tony/Work/tony/hoot/tokenizer_spe_bpe_v20480/tokenizer.model"
# tokenizer_path = "/home/tony/Work/tony/hoot/tokenizers/v3/tokenizer_spe_bpe_v10240/tokenizer.model"
# Gpt tokenizer
# tokenizer_path = "/app/suno/data/dpo/models/tokenizer_60k.json"
# v5 tokenizer -- trained on all data
tokenizer_path = (
    "/home/tony/Work/tony/hoot/tokenizers/v5/tokenizer_spe_bpe_v20481/tokenizer.model"
)
max_duration_s = 4 * 60
debug_val_only = False
preload_checkpoint = None  # "/home/tony/Data/Hoot/stt_en_fastconformer_ctc_large.pt"
# preload_checkpoint = "/home/tony/Data/Hoot/stt_multilingual_fastconformer_hybrid_large_pc.pt"
preload_optimizer = False
preload_strict = False
preload_decoder = False  # use the same decoder as before; only if train from crash
suppress_compile_warnings = True
# eval items
custom_seed_offset = 1234
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 = True
wandb_project = "hoot-v2"
wandb_run_name = "suno_hoot_en_default_clean"
# data
gradient_accumulation_steps = 1  # used to simulate larger batch sizes
batch_size = 1  # if gradient_accumulation_steps > 1, this is the micro-batch size
max_tokens = 16000 * 40 * 60 * 3  # this is the max durations that we can fit in a batch
# model
n_layers = 18
n_embd = 512
n_augment_freq_masks = 2
decoder_type = "v1"
max_text_len = 2000
# 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.98
grad_clip = 0.1  # 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
lyrics_key = "text"  # text is the clean text, lyrics is the original lyrics
# -----------------------------------------------------------------------------
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)
eval_iters = min(eval_iters, max_iters)
if lr_decay_iters is None:
    lr_decay_iters = max_iters


# set up distributed variables
# os.environ["MASTER_ADDR"] = str(master_addr)
# os.environ["MASTER_PORT"] = str(master_port)
if "SLURM_PROCID" in os.environ:  # Running on SLURM
    if int(os.environ["SLURM_NTASKS_PER_NODE"]) != torch.cuda.device_count():
        raise ValueError(
            f"SLURM_NTASKS_PER_NODE ({os.environ['SLURM_NTASKS_PER_NODE']}) does not match"
            f" the number of CUDA devices ({torch.cuda.device_count()}) on node {os.environ['HOSTNAME']}"
        )
    ddp_rank = int(os.environ["SLURM_PROCID"])
    ddp_local_rank = int(os.environ["SLURM_LOCALID"])
    world_size = int(os.environ["SLURM_JOB_NUM_NODES"]) * int(
        os.environ["SLURM_NTASKS_PER_NODE"]
    )

else:  # Running locally
    ddp_rank = 0
    ddp_local_rank = 0
    world_size = 1
os.environ["RANK"] = str(ddp_rank)
os.environ["LOCAL_RANK"] = str(ddp_local_rank)

# various inits, derived attributes, I/O setup
ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
if ddp:
    print(
        f"initializing ddp with rank {ddp_rank}, local rank {ddp_local_rank}, world size {world_size}"
    )
    try:
        init_process_group(
            backend="nccl",
            # timeout=datetime.timedelta(seconds=24 * 60 * 60),
            rank=ddp_rank,
            world_size=world_size,
            device_id=torch.device(f"cuda:{ddp_local_rank}"),
        )
    except Exception as e:
        print(
            f"Distributed error on rank {ddp_rank} with host {os.environ['HOSTNAME']}"
        )
        raise e
    print(f"finished ddp with rank {ddp_rank}, world size {world_size}")
    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
n_gpus_per_node = torch.cuda.device_count()
dist_barrier()


def print_with_time(content):
    if master_process:
        print(f"[{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}]: {content}")


print_with_time(f"ddp init: world size {world_size} ddp_rank {ddp_rank}.")


gc.collect()  # Clear memory before starting
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_with_time(f"logging checkpoint here: {out_dir}")


# model init
print_with_time("Initializing a new model from scratch")
tokenizer = Tokenizer(tokenizer_path)
print_with_time(f"Total vocab size: {tokenizer.n_vocab}")
model_args = dict(
    n_layers=n_layers,
    n_embd=n_embd,
    n_classes=tokenizer.n_vocab,
    n_augment_freq_masks=n_augment_freq_masks,
    decoder_type=decoder_type,
    max_text_len=max_text_len,
)
hoot_configuration = HootConfig(**model_args)
model = Hoot(hoot_configuration)
model.to(device)

# load data & pre encode the labels...
print_with_time(f"loading data...")
train_data_metas = []

# load l is the old format...
try:
    with open(train_input_path, "r") as fp:
        for l in fp:
            l_dict = json.loads(l)
            train_data_metas.append(l_dict)
except Exception as e:
    print(f"Error loading {train_input_path}: {e}")
    with open(train_input_path, "r") as fp:
        train_data_metas = json.load(fp)
if train_suno_input_path is not None:
    with open(train_suno_input_path, "r") as fp:
        suno_data_metas = json.load(fp)
        train_data_metas.extend(suno_data_metas)
val_data_metas = []
with open(val_input_path, "r") as fp:
    for l in fp:
        l_dict = json.loads(l)
        val_data_metas.append(l_dict)
print_with_time(
    f"total train, {len(train_data_metas)}, {round(sum([m['duration'] for m in train_data_metas]) / 3600000, 2)} khrs, total valid, {len(val_data_metas)}, {round(sum([m['duration'] for m in val_data_metas]) / 3600000, 2)} khrs"
)

# 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.amp.GradScaler("cuda", 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_with_time(f"preloading checkpoint: {preload_checkpoint}")
    cur_state_dict = model.state_dict()
    checkpoint = torch.load(preload_checkpoint, map_location=device, weights_only=False)
    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_prefixs = ["_orig_mod.", "joint", "ctc_decoder"]
    if not preload_decoder:
        unwanted_prefixs.append("decoder")
    for k, v in list(state_dict.items()):
        for unwanted_prefix in unwanted_prefixs:
            if k.startswith(unwanted_prefix):
                state_dict[k[len(unwanted_prefix) :]] = state_dict.pop(k)
                break
    # 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 not preload_decoder else False
    )
    if preload_optimizer:
        print_with_time("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_with_time("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
dist_barrier()

# After loading data_metas, initialize datasets and dataloaders
train_dataset = HootDataset(
    train_data_metas, tokenizer, lyrics_key=lyrics_key, max_duration=max_duration_s
)

val_dataset = HootDataset(
    val_data_metas, tokenizer, lyrics_key=lyrics_key, max_duration=max_duration_s
)

# Set up samplers for distributed training
if ddp:
    train_sampler = DistributedSampler(
        train_dataset, num_replicas=world_size, rank=ddp_rank, shuffle=True
    )
    val_sampler = DistributedSampler(
        val_dataset, num_replicas=world_size, rank=ddp_rank, shuffle=False
    )
else:
    train_sampler = None
    val_sampler = None

# Create DataLoaders
num_workers = 4  # Adjust based on CPU cores
train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=(train_sampler is None),
    sampler=train_sampler,
    num_workers=num_workers,
    collate_fn=hoot_collate_fn,
    pin_memory=True,
    drop_last=True,
    persistent_workers=True,
    prefetch_factor=3,
)

val_loader = DataLoader(
    val_dataset,
    batch_size=batch_size,
    shuffle=False,
    sampler=val_sampler,
    num_workers=num_workers,
    collate_fn=hoot_collate_fn,
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=3,
)


def get_sample(split):
    if split == "train":
        data_metas = train_data_metas
    else:
        data_metas = val_data_metas
    idx = random.randint(0, len(data_metas) - 1)
    meta = data_metas[idx]
    audio_file_path = meta["audio_filepath"]
    waveform = Audio.from_file(audio_file_path)
    waveform = torch.from_numpy(waveform.convert(SAMPLE_RATE, 2, 1).array_float)
    # try to load the tokens directly
    y = meta.get("tokens", tokenizer.encode(meta[lyrics_key]))
    # update the data cache
    if "tokens" not in meta:
        data_metas[idx]["tokens"] = y
    # print(meta[lyrics_key], y, y.shape)
    x = waveform
    # clip to max space we have for logits
    y = torch.from_numpy(y[: int(round(x.shape[-1] / 16_000 * 10))])
    if x.shape[0] > 1:
        x = x.mean(axis=0)
    else:
        x = x.reshape(-1)
    # print(x.shape, y.shape)
    return x.to(device), y.to(device), idx


def get_batch(split):
    x_list = []
    y_list = []
    idxs = []
    # curr_num_of_seq = 0
    for _ in range(batch_size):
        x, y, idx = get_sample(split)
        # make sure that we will not exceeed the max length of the info we can pack in
        # if curr_num_of_seq + x.shape[0] < max_tokens:
        # let's check label frequency -- this can be really screwed up and blows up the loss
        # we have done data cleaning, don't need to do this again. Hopefuly that speeds up a bit
        # max_label_freq = torch.max(torch.bincount(y)).item() / (y.shape[0] + 1)
        # while max_label_freq >= 0.5:
        #     print("WTF", idx, max_label_freq)
        #     x, y, idx = get_sample(split)
        #     max_label_freq = torch.max(torch.bincount(y)).item() / (y.shape[0] + 1)
        x_list.append(x)
        y_list.append(y)
        idxs.append(idx)
    x, x_len = collate_fn(x_list, fixed_len=None)
    y, y_len = collate_fn(y_list, fixed_len=None)
    del x_list, y_list
    return x, x_len, y, y_len, idxs


# helps estimate an arbitrarily accurate loss over either split using many batches
@torch.no_grad()
def estimate_loss():
    effective_eval_iters = min(max(eval_iters, 1), 100)
    out = {}
    model.eval()

    for split in ["train", "val"]:
        losses = []
        default_wer = metrics.WordErrorRate(device=device)
        truth_lyrics = []
        pred_lyrics = []

        # Use the appropriate loader
        loader = train_loader if split == "train" else val_loader
        # Create new prefetcher for evaluation
        prefetcher = DataPrefetcher(loader, device)

        for k in range(effective_eval_iters):
            batch = prefetcher.next()
            if batch is None:
                # Explicitly clean up before creating new prefetcher
                del prefetcher
                torch.cuda.empty_cache()
                prefetcher = DataPrefetcher(loader, device)
                batch = prefetcher.next()
                if batch is None:  # If still None, skip this iteration
                    continue

            # Extract tensors and remove batch reference
            X, X_len, Y, Y_len, _ = batch
            del batch

            with ctx:
                curr_loss, decoded = model(
                    X, X_len, targets=Y, targets_len=Y_len, return_decoded=True
                )

            # Process batch
            for i in range(decoded.shape[0]):
                logits = decoded[i].detach().cpu().numpy()
                pred_label = tokenizer.decode_logits(logits)
                valid_y = Y[i, :][Y[i, :] > 0]
                true_label = tokenizer.decode(valid_y.detach().cpu().tolist())
                pred_lyrics.append(pred_label)
                truth_lyrics.append(true_label)

            losses.append(curr_loss.item())

            # Explicit cleanup after each evaluation batch
            del X, X_len, Y, Y_len, decoded, curr_loss
            torch.cuda.empty_cache()

        # Calculate metrics
        default_wer.update(pred_lyrics, truth_lyrics)
        wer_value = default_wer.compute().item()

        out[f"{split}/loss"] = np.mean(losses)
        out[f"{split}/wer"] = wer_value

        # Complete cleanup after evaluating this split
        del default_wer, pred_lyrics, truth_lyrics, losses
        del prefetcher
        torch.cuda.empty_cache()

    model.train()
    return out


def get_model_test_transcribe():
    model.eval()
    idx = random.randint(0, len(train_data_metas) - 1)
    meta = train_data_metas[idx]
    audio_file_path = meta["audio_filepath"]
    tokenized_label = tokenizer._tokenizer.decode(
        tokenizer._tokenizer.encode(meta[lyrics_key])
    )

    # Load audio using Audio class
    audio = Audio.from_file(audio_file_path)
    arr = torch.from_numpy(audio.convert(SAMPLE_RATE, 2, 1).array_float)

    # The Audio class already handles mono conversion and resampling
    # So we don't need these lines anymore:
    # if arr.shape[0] > 1:
    #     arr = arr.mean(dim=0, keepdim=True)
    # arr = torchaudio.functional.resample(arr, sr, 16_000)

    arr = arr.to(device)
    arrays, arrays_len = collate_fn([arr])
    with torch.no_grad():
        decoded, _ = model.forward(arrays, arrays_len)
    # no need to unpad cause batch 1
    logits = decoded[0].detach().cpu().numpy()
    pred_label = tokenizer.decode_logits(logits)
    print_with_time(f"origin: {tokenized_label}")
    print_with_time(f"pred: {pred_label}")
    model.train()


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


# Initialize the prefetcher before training
train_prefetcher = DataPrefetcher(train_loader, device)

# Get first batch
batch = train_prefetcher.next()
if batch is None:
    train_prefetcher = DataPrefetcher(train_loader, device)
    batch = train_prefetcher.next()
X, X_len, Y, Y_len, idxs = batch

# In the training loop, replace the manual batch loading:
# Delete this line within the micro_step loop:
# X, X_len, Y, Y_len, idxs = get_batch("train")

# And add this after optimizer.zero_grad(set_to_none=True):
next_batch = train_prefetcher.next()
if next_batch is None:
    # Reset for new epoch
    if ddp:
        train_loader.sampler.set_epoch(iter_num // len(train_loader) + 1)
    train_prefetcher = DataPrefetcher(train_loader, device)
    next_batch = train_prefetcher.next()
X, X_len, Y, Y_len, idxs = next_batch

# training loop
if master_process:
    print_with_time("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_with_time("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_with_time(
            f"loss estimation took {estimation_time:.1f} seconds. ({eval_time_pct:.1f}% of loop)"
        )
        print_with_time(
            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": float(losses["val/loss"]),
                    "config": config,
                }
                print_with_time(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"))
                    torch.save(
                        checkpoint,
                        os.path.join(out_dir, f"{iter_num // 1000}k_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_with_time("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
            )
        with ctx:
            # Create text_input with random masking (50% chance)
            if decoder_type == "v1":
                loss = model(X, X_len, targets=Y, targets_len=Y_len)
            elif decoder_type == "v2":
                text_input = Y.clone()
                mask = torch.rand(text_input.shape[0]) < 0.8
                # Zero out masked text inputs
                text_input[mask] = torch.zeros_like(text_input[mask])

                # Apply random shuffling to 10% of the batch
                shuffle_mask = torch.rand(text_input.shape[0]) < 0.1
                if shuffle_mask.any():
                    shuffle_indices = shuffle_mask.nonzero(as_tuple=True)[0]
                    for i in shuffle_indices:
                        # Get non-zero token positions only
                        non_zero_mask = text_input[i] > 0
                        non_zero_positions = non_zero_mask.nonzero(as_tuple=True)[0]

                        if (
                            len(non_zero_positions) > 1
                        ):  # Only shuffle if there's content to shuffle
                            # Get the actual non-zero tokens
                            non_zero_tokens = text_input[i, non_zero_positions]
                            # Shuffle only the non-zero tokens
                            shuffled_tokens = non_zero_tokens[
                                torch.randperm(len(non_zero_tokens))
                            ]
                            # Put shuffled tokens back in their original positions
                            text_input[i, non_zero_positions] = shuffled_tokens

                # crop the text input to the max length (ALWAYS apply, not just when shuffling)
                text_input = text_input[:, : hoot_configuration.max_text_len]
                loss = model(
                    X, X_len, targets=Y, targets_len=Y_len, text_input=text_input
                )
            else:
                raise ValueError(f"Unknown decoder type: {decoder_type}")
            loss_val = loss.item()  # loss as float. note: this is a CPU-GPU sync point
            # print_with_time(f"iter: {iter_num}, loss_val: {loss_val}")
            loss = loss / gradient_accumulation_steps
        # Get next batch asynchronously
        next_batch = train_prefetcher.next()
        if next_batch is None:
            # Reset for new epoch
            if ddp:
                train_loader.sampler.set_epoch(iter_num // len(train_loader) + 1)
            # Clean up old prefetcher first
            del train_prefetcher
            torch.cuda.empty_cache()
            train_prefetcher = DataPrefetcher(train_loader, device)
            next_batch = train_prefetcher.next()

        # Extract tensors and immediately remove batch reference
        next_X, next_X_len, next_Y, next_Y_len, next_idxs = next_batch
        del next_batch

        # After backward pass and optimizer step, update tensors
        X, X_len, Y, Y_len, idxs = next_X, next_X_len, next_Y, next_Y_len, next_idxs
        del next_X, next_X_len, next_Y, next_Y_len, next_idxs
        if debug_gradients and wandb_log and master_process:
            d = {
                "iter": iter_num,
                "debug_loss": float(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"] = float(np.mean(grads))
                wandb.log(d)

        # backward pass, with gradient scaling if training in fp16
        if torch.isnan(loss) or torch.isinf(loss):
            print_with_time(
                f"Loss is BAD at iter {iter_num}, {loss}, idxs are {idxs}, skip by setting loss to 0"
            )
            loss = torch.tensor([0.0], requires_grad=True).to(device)
        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_with_time(
            f"iter {iter_num}: avg_loss {avg_loss:.3f}, idxs are {idxs}, step_time {dt*1000:.1f}ms"
        )
        d = {
            "iter": iter_num,
            "train_running_loss": avg_loss,
        }
        wandb.log(d)
    iter_num += 1
    local_iter_num += 1

    # Memory management and termination
    if iter_num % 1000 == 0:
        # Regular memory cleanup
        gc.collect()
        torch.cuda.empty_cache()

        # More aggressive cleanup every 500 iterations if memory usage is high
        # if (
        #     iter_num % 500 == 0
        #     and torch.cuda.memory_allocated() > 0.8 * torch.cuda.max_memory_allocated()
        # ):
        #     # Reset prefetcher completely
        #     del train_prefetcher
        #     torch.cuda.empty_cache()
        #     train_prefetcher = DataPrefetcher(train_loader, device)

        #     # Get fresh batch
        #     batch = train_prefetcher.next()
        #     if batch is not None:
        #         X, X_len, Y, Y_len, idxs = batch
        #         del batch

    # Termination condition
    if iter_num > max_iters:
        break

if ddp:
    destroy_process_group()