import os
import re
import glob
import json
import torch
import funcy
import faiss
import IPython
import numpy as np
import torchaudio
import itertools

from tqdm import tqdm
from typing import List
from dac.model.dac4 import DAC
from dac.model.discriminator2 import Discriminator as Discriminator_import
from dac.nn import loss as loss_import
from dac.utils.accelerator import Accelerator
from dac.utils import load_model


import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans, MiniBatchKMeans

from suno_utils.models.musicfm.modeling_MusicFM import MusicFM_MERTLong
from suno_utils.utils.s3 import read_from_s3

# from suno_utils.utils.text import (
#    write_jsonl,
#    read_jsonl,
#    write_json,
#    read_json,
#    normalize_whitespace,
# )

from transformers import WhisperProcessor, WhisperForConditionalGeneration


SAMPLE_RATE = 48_000
VAE_DIM = 128
VAE_T_MEMMAP = 1000
USE_MERT = False
MERT_DIM = 1024
MERT_T_MEMMAP = 250
NUM_FRAMES = int(10 * SAMPLE_RATE)
BATCH_SIZE = 32
MAX_N_CHUNKS = 3
MAX_FILES_PER_SUBSET = 500000
USE_WHISPHER = True


class SimpleAudioDataset(torch.utils.data.Dataset):
    def __init__(self, filepaths: List[str]):
        self.filepaths = filepaths

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

    def __getitem__(self, idx):
        audio_file = self.filepaths[idx]

        num_frames = torchaudio.info(audio_file).num_frames

        if num_frames > NUM_FRAMES:
            start_frame = np.random.randint(0, num_frames - NUM_FRAMES - 1)
            end_frame = start_frame + NUM_FRAMES
            # load all the audio
            audio, sr = torchaudio.load(
                audio_file, frame_offset=start_frame, num_frames=NUM_FRAMES
            )
        else:
            start_frame = 0
            end_frame = num_frames
            audio, sr = torchaudio.load(audio_file)

        if audio.shape[0] != 2:
            audio = audio.repeat(2, 1)

        if audio.shape[-1] < NUM_FRAMES:
            audio = torch.cat(
                [audio, torch.zeros(2, NUM_FRAMES - audio.shape[-1])], dim=-1
            )

        assert sr == SAMPLE_RATE

        return audio, start_frame, end_frame


def write_jsonl(data, filename):
    with open(filename, "w", encoding="utf-8") as f:
        for d in data:
            f.write(json.dumps(d, ensure_ascii=False) + "\n")


class AudioDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        filepaths: List[str],
    ):
        # check length of ecah file and create examples
        audio_chunks = []
        print("Checking audio files...")
        for audio_file in tqdm(filepaths):
            audio, sr = torchaudio.load(audio_file)
            audio /= audio.abs().max().clamp(1e-4)
            num_frames = audio.shape[-1]
            num_chunks = num_frames // NUM_FRAMES
            audio_file_chunks = []
            for i in range(num_chunks):
                start_frame = i * NUM_FRAMES
                end_frame = (i + 1) * NUM_FRAMES
                audio_chunk = audio[..., start_frame:end_frame]
                audio_chunk_energy = audio_chunk.pow(2).mean().item()
                if audio_chunk_energy > 1e-3 and np.random.rand() > 0.5:
                    audio_file_chunks.append((audio_file, start_frame, end_frame))

                if len(audio_file_chunks) >= MAX_N_CHUNKS:
                    break

            audio_chunks.extend(audio_file_chunks)

        print("Total", len(filepaths), "chunks", len(audio_chunks))
        print(np.random.choice(filepaths, 5))

        self.filepaths = filepaths
        self.audio_chunks = audio_chunks

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

    def __getitem__(self, idx):
        audio_file, start_frame, end_frame = self.audio_chunks[idx]

        # load all the audio
        audio, sr = torchaudio.load(
            audio_file, frame_offset=start_frame, num_frames=NUM_FRAMES
        )

        if audio.shape[0] != 2:
            audio = audio.repeat(2, 1)

        if audio.shape[-1] < NUM_FRAMES:
            audio = torch.cat(
                [audio, torch.zeros(2, NUM_FRAMES - audio.shape[-1])], dim=-1
            )

        assert sr == SAMPLE_RATE

        return audio


if __name__ == "__main__":

    # load VAE
    device = "cuda:0"
    # checkpoint_filepath = "/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/dac/weights.pth"
    checkpoint_filepath = "s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth"
    load_f = funcy.partial(torch.load, map_location="cpu")

    if checkpoint_filepath.startswith("s3://"):
        sd = read_from_s3(checkpoint_filepath, read_f=load_f)
    else:
        sd = load_f(checkpoint_filepath)

    sd["metadata"]["kwargs"] = {
        k: v
        for k, v in sd["metadata"]["kwargs"].items()
        if k in DAC.__init__.__code__.co_varnames
    }
    model_100hz = DAC(**sd["metadata"]["kwargs"])
    model_100hz.load_state_dict(sd["state_dict"])
    model_100hz.eval()
    model_100hz.to(device)

    # load MERT
    if USE_MERT:
        model_filepath = "s3://suno-data/minz/models/musicfm_concat_epoch=51.pt"
        centroids_filepath = "s3://suno-data/minz/models/musicfm_concat_centroids.npy"
        mert_model = MusicFM_MERTLong(
            is_flash=False,
            stat_path="s3://suno-data/minz/models/mertlong_stats.json",
            model_path=model_filepath,
        )
        mert_model.cuda()
        mert_model.eval()
        for name, param in mert_model.named_parameters():
            param.requires_grad = False

    # load whisper model for transcription
    if USE_WHISPHER:
        whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-base")
        whisper_model = WhisperForConditionalGeneration.from_pretrained(
            "openai/whisper-base"
        )
        whisper_model.cuda()

    # root_dir = "/app/suno/christian/data/tiktok_covers_lyrics_48khz"
    root_dir = "/app/suno/data/audio_2ch_48khz_lg"
    dataset_name = "genius_hq"
    # set up train and test splits

    embed_dir = "/app/suno/christian/data/suno_diffusion_genius_hq_lyrics"
    os.makedirs(embed_dir, exist_ok=True)

    for subset_name in ["train"]:

        subset_filepaths = glob.glob(
            os.path.join(root_dir, subset_name, dataset_name, "*.wav")
        )
        subset_filepaths = subset_filepaths

        dataset = SimpleAudioDataset(subset_filepaths)
        dataloader = torch.utils.data.DataLoader(
            dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=32
        )

        # output memmap
        vae_example_size = VAE_DIM * VAE_T_MEMMAP
        mert_example_size = MERT_DIM * MERT_T_MEMMAP

        vae_mm_write_idx = 0
        mert_mm_write_idx = 0

        print("Creating memmap...")
        out_vae_mm = np.memmap(
            f"{embed_dir}/vae_{subset_name}.bin",
            dtype=np.float32,
            mode="w+",
            shape=(len(dataset) * vae_example_size),
        )

        if USE_MERT:
            out_mert_mm = np.memmap(
                f"{embed_dir}/mert_{subset_name}.bin",
                dtype=np.float32,
                mode="w+",
                shape=(len(dataset) * mert_example_size),
            )

        if USE_COARSE:
            out_codec_mm = np.memmap(
                f"{embed_dir}/codec_{subset_name}.bin",
                dtype=np.float32,
                mode="w+",
                shape=(len(dataset) * N_COARSE_CODEBOOKDS * CODEC_T_MEMMAP),
            )

        audios = []
        metas_out = []
        global_idx = 0
        print("Encoding...")
        for batch in tqdm(dataloader):
            audio_chunks, start_fames, end_frames = batch

            # transcribe
            if USE_WHISPHER:
                audios_16k = []
                audio_chunks_16k = torchaudio.functional.resample(
                    audio_chunks.mean(dim=1, keepdim=False), 48000, 16000
                )
                # convert audio to list of mono numpy arrays at 16 khz
                for audio in audio_chunks_16k:
                    audios_16k.append(audio.numpy())

                input_features = whisper_processor(
                    audios_16k,
                    sampling_rate=16000,
                    return_tensors="pt",
                ).input_features

                input_features = input_features.cuda()

                # Generate token ids
                predicted_ids = whisper_model.generate(input_features)

                # Decode token ids to text
                transcriptions = whisper_processor.batch_decode(
                    predicted_ids, skip_special_tokens=True
                )

            audios = audio_chunks.to(device)

            if USE_MERT:
                with torch.no_grad():
                    audios_24k_mono = torchaudio.functional.resample(
                        audios.mean(dim=1, keepdim=False), 48_000, 24_000
                    )

                    out = mert_model.get_latent(audios_24k_mono)
                    mert_z = out

            with torch.no_grad():
                out = model_100hz.encode(audios)
                vae_z = out["z"]

            # store metas
            for bidx in range(vae_z.shape[0]):
                meta = {}
                meta["audio_file"] = subset_filepaths[global_idx]
                meta["dataset"] = dataset_name
                if USE_WHISPHER:
                    meta["lyrics"] = transcriptions[bidx]
                    print(transcriptions[bidx])
                meta["start_frame"] = start_fames[bidx].item()
                meta["end_frame"] = end_frames[bidx].item()
                meta["start_sec"] = meta["start_frame"] / SAMPLE_RATE
                meta["end_sec"] = meta["end_frame"] / SAMPLE_RATE

                metas_out.append(meta)
                global_idx += 1

            # store results into memmap
            for bidx in range(vae_z.shape[0]):
                # write vae embed to memmap
                vae_embed = vae_z[bidx : bidx + 1, ...].reshape(-1).cpu().numpy()
                # check for nan in embed
                if np.isnan(vae_embed).any():
                    print("Found NaN in VAE embed")
                    continue

                out_vae_mm[vae_mm_write_idx : vae_mm_write_idx + vae_example_size] = (
                    vae_embed
                )
                vae_mm_write_idx += vae_example_size

                # write mert embed to memmap
                if USE_MERT:
                    mert_embed = mert_z[bidx : bidx + 1, ...].reshape(-1).cpu().numpy()
                    # check for nan in embed
                    if np.isnan(mert_embed).any():
                        print("Found NaN in MERT embed")
                        continue

                    out_mert_mm[
                        mert_mm_write_idx : mert_mm_write_idx + mert_example_size
                    ] = mert_embed
                    mert_mm_write_idx += mert_example_size

        # flush
        out_vae_mm.flush()
        del out_vae_mm

        # save out metas
        write_jsonl(metas_out, f"{embed_dir}/{subset_name}_metas.jsonl")

        if USE_MERT:
            out_mert_mm.flush()
            del out_mert_mm