# codec cycle raw audio and save to disk
import os
import glob
import torch
import argparse
import torchaudio
import numpy as np

from tqdm import tqdm
from typing import Optional
from suno_utils.tasks.dac_2c_12cb import DAC

EXTS = ["wav", "mp3", "flac", "ogg"]


def load_codec(codec_path: str):
    sd = torch.load(codec_path, map_location="cpu")
    model = DAC(**sd["metadata"]["kwargs"])
    model.load_state_dict(sd["state_dict"])
    model.eval()
    return model


def codec_cycle(
    audio: torch.Tensor,
    sample_rate: int,
    dac: torch.nn.Module,
    n_quantizers: Optional[int] = None,
    noise_level: Optional[float] = 0.0,
):
    length = audio.shape[-1]
    with torch.no_grad():
        audio = dac.preprocess(audio.unsqueeze(0), sample_rate)
        z, codes, latents, commitment_loss, codebook_loss = dac.encode(
            audio, n_quantizers
        )
        print(z)

        if noise_level > 0.0:
            p = 1.0  # Adjust this probability as needed
            z_noise = noise_level * torch.randn((z.shape[-2], z.shape[-1]))
            bernoulli_mask = torch.rand(z.shape[-1]) < p
            bernoulli_mask = bernoulli_mask.unsqueeze(0).expand(z.shape[-2], -1)
            print(bernoulli_mask.shape)
            z_noise_with_mask = z_noise * bernoulli_mask.float()
            print(z_noise_with_mask.shape)
            z += z_noise_with_mask
            print("noised", z)

        audio = dac.decode(z)

    return audio[..., :length].squeeze(0)


if __name__ == "__main__":
    codec_path = "checkpoints/dac/dac_2c_25x12.pt"
    model = load_codec(codec_path)
    codec_name = os.path.basename(codec_path).split(".")[0]
    print(f"Encoding audio with {codec_name}...")

    x, sr = torchaudio.load("/app/suno/christian/reference-audio-wav/02 Dreams.wav")

    if sr != 48000:
        x = torchaudio.functional.resample(x, sr, 48000)
        sr = 48000

    x = x[:, 524288 + 262144 : 2 * 524288]

    # peak normalize
    x /= x.abs().max()

    torchaudio.save(f"dreams-original.wav", x, sr)

    # create multiple encodings at different input levels
    for gain_db in [-0.0, -3.0, -6.0, -12.0, -18.0, -24.0, -32.0, -48.0]:
        gain_lin = 10 ** (gain_db / 20.0)
        y = codec_cycle(x * gain_lin, sr, model)
        y /= y.abs().max()
        torchaudio.save(f"dreams-encoded-gain_db={gain_db}.wav", y, sr)