# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/03_train_audio_algebra.ipynb (unless otherwise specified).

__all__ = ['init_jukebox_sample_rate', 'audio_for_jbx', 'load_audio_for_jbx', 'JUKEBOX_SAMPLE_RATE', 'IceBoxEncoder',
           'AudioAlgebra', 'main']

# Cell
from prefigure.prefigure import get_all_args, push_wandb_config
from copy import deepcopy
import math
import json

import accelerate
import sys
import torch
import torchaudio
from torch import optim, nn
from torch import multiprocessing as mp
from torch.nn import functional as F
#from torch.utils import data
from tqdm import tqdm, trange
from einops import rearrange, repeat

import wandb
from .viz import embeddings_table, pca_point_cloud, audio_spectrogram_image, tokens_spectrogram_image
import shazbot.blocks_utils as blocks_utils

# Cell
# TagBox Utils:
#  Utilities from Ethan Manilows's TagBox: https://github.com/ethman/tagbox
# slightly modified by Scott H. Hawley @drscotthawley

#JUKEBOX_SAMPLE_RATE = 44100  # ethan's original
JUKEBOX_SAMPLE_RATE = None

def init_jukebox_sample_rate(sr=44100): # will probably use 48000 in practice
    "SHH added this util to preserve rest of code minimall-modified"
    global JUKEBOX_SAMPLE_RATE
    JUKEBOX_SAMPLE_RATE = sr
    return

def audio_for_jbx(audio, trunc_sec=None, device=None):
    """Readies an audio TENSOR for Jukebox."""
    if audio.ndim == 1:
        audio = audio[None]
        audio = audio.mean(axis=0)

    # normalize audio
    norm_factor = torch.abs(audio).max()
    if norm_factor > 0:
        audio /= norm_factor

    if trunc_sec is not None:  # truncate sequence
        audio = audio[: int(JUKEBOX_SAMPLE_RATE * trunc_sec)]

    audio = audio[:, :, None]  # add one more dimension on the end?
    return audio


def load_audio_for_jbx(path, offset=0.0, dur=None, trunc_sec=None, device=None):
    """Loads a path for use with Jukebox."""
    audio, sr = librosa.load(path, sr=None, offset=offset, duration=dur)

    if JUKEBOX_SAMPLE_RATE is None: init_jukebox_sample_rate()

    if sr != JUKEBOX_SAMPLE_RATE:
        audio = librosa.resample(audio, sr, JUKEBOX_SAMPLE_RATE)

    return audio_for_jbx(audio, trunc_sec, device=device)

# Cell
# "Icebox": frozen Jukebox encoder for embeddings
class IceBoxEncoder(nn.Module):
    def __init__(self, global_args):
        super().__init__()

        n_io_channels = 2
        n_feature_channels = 8
        self.num_quantizers = global_args.num_quantizers
        self.ema_decay = global_args.ema_decay

        # for making Jukebox work with multi-GPU runs
        rank, local_rank, device = int(os.getenv('RANK')), int(os.getenv('LOCAL_RANK')), self.device
        dist_url = "tcp://127.0.0.1:9500"
        dist.init_process_group(backend="nccl")

        self.hps = Hyperparams()
        assert global_args.sample_rate == 44100, "Jukebox was pretrained at 44100 Hz."
        self.hps.sr = global_args.sample_rate #44100
        self.hps.levels = 3
        self.hps.hop_fraction = [.5,.5,.125]

        vqvae = "vqvae"
        self.vqvae = make_vqvae(setup_hparams(vqvae, dict(sample_length = 1048576)), self.device)
        for param in self.vqvae.parameters():  # FREEZE IT.  "IceBox"
            param.requires_grad = False

        self.encoder = self.vqvae.encode
        #self.encoder_ema = deepcopy(self.encoder)

        latent_dim = 64 # global_args.latent_dim. Jukebox is 64
        io_channels = 2#1 # 2.  Jukebox is mono but we decode in stereo

    def encode(self, *args, **kwargs):
        return self.encoder(*args, **kwargs)

# Cell
class AudioAlgebra(nn.Module):
    def __init__(self, global_args, device):
        super().__init__()

        self.device = device

        self.encoder = JukeBoxEncoder(global_args)


    def loss(self, reals):

        if self.pqmf_bands > 1:
            encoder_input = self.pqmf(reals)
        else:
            encoder_input = reals

        with torch.cuda.amp.autocast():
            latents = self.encoder(encoder_input).float()

        loss = 0# TODO: fill in mse_loss + mrstft_loss

        log_dict = {
            'train/mse_loss': mse_loss.detach(),
            'train/mb_distance': mb_distance.detach(),
        }

        return loss, log_dict

# Cell
def main():

    args = get_all_args()

    args.random_crop = False

    torch.manual_seed(args.seed)

    try:
        mp.set_start_method(args.start_method)
    except RuntimeError:
        pass

    accelerator = accelerate.Accelerator()
    device = accelerator.device
    print('Using device:', device, flush=True)

    aa_model = AudioAlgebra(args, device)

    accelerator.print('Parameters:', blocks_utils.n_params(diffusion_model))

    # If logging to wandb, initialize the run
    use_wandb = accelerator.is_main_process and args.name
    if use_wandb:
        import wandb
        config = vars(args)
        config['params'] = utils.n_params(diffusion_model)
        wandb.init(project=args.name, config=config, save_code=True)
