import os
import torch
import torchaudio
from torch.utils.data import Dataset, DataLoader
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm

import os
import gc
import time
import json
import glob
import random
import torch
import funcy
import queue
import threading
import torchaudio
import numpy as np

from suno_utils.utils.s3 import read_from_s3
from suno_utils.utils.text import read_jsonl
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from typing import List, Iterator, Dict, Any, Optional
from dataclasses import dataclass
from tqdm import tqdm
from dac.model.dac4 import DAC


# load the semantic model
from suno_utils.tasks.mert_25 import (
    preload_models as preload_semantic_models,
    encode as encode_semantic,
    EMBEDDING_RATE as SEMANTIC_HZ,
)

print("loading semantic model...")
semantic_model_filepath = "s3://suno-data/georg/models/semantic/mert_25.pt"
semantic_clusters_filepath = "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy"
_ = preload_semantic_models(
    checkpoint_filepath=semantic_model_filepath,
    centroids_filepath=semantic_clusters_filepath,
    device="cuda",
)


def corrupt_audio(audio, sr):
    audio_corrupt = audio.clone()
    # 10% chance to add highpass
    if np.random.uniform() < 0.25:
        hpf_cutoff = np.random.uniform(20, 1000)
        audio_corrupt = torchaudio.functional.highpass_biquad(
            audio_corrupt, sr, hpf_cutoff
        )
    # 10% chance to add lowpass
    if np.random.uniform() < 0.25:
        lpf_cutoff = np.random.uniform(1000, 20000)
        audio_corrupt = torchaudio.functional.lowpass_biquad(
            audio_corrupt, sr, lpf_cutoff
        )

    # 10% chance to make mono
    if np.random.uniform() < 0.7:
        audio_corrupt = audio_corrupt.mean(dim=0, keepdim=True)
        audio_corrupt = audio_corrupt.repeat(2, 1)

    # 10% chance to add noise
    if np.random.uniform() < 0.2:
        noise_level_db = -np.random.uniform(42, 96)
        noise_signal = torch.randn_like(audio_corrupt) * 10 ** (noise_level_db / 20)
        # apply filter to noise signal
        lpf_cutoff = np.random.uniform(1000, 20000)
        noise_signal = torchaudio.functional.lowpass_biquad(
            noise_signal, sr, lpf_cutoff
        )
        audio_corrupt = audio_corrupt + noise_signal

    # 5% chance to add contrast
    if np.random.uniform() < 0.05:
        contrast = np.random.uniform(0, 100.0)
        audio_corrupt = torchaudio.functional.contrast(
            audio_corrupt, enhancement_amount=contrast
        )

    # 10% add codec artifacts
    if np.random.uniform() < 1.0:
        compression = np.random.choice([32, 64, 128])
        audio_corrupt = torchaudio.functional.apply_codec(
            audio_corrupt, sr, "mp3", compression
        )

    # 10% chance to add distortion
    if np.random.uniform() < 0.3:
        drive_db = np.random.uniform(6, 24)
    else:
        drive_db = 0

    drive_lin = 10 ** (drive_db / 20)
    audio_corrupt *= drive_lin

    if np.random.uniform() < 0.5:
        audio_corrupt = torch.tanh(audio_corrupt)
    else:
        audio_corrupt = audio_corrupt.clamp(-1, 1)

    return audio_corrupt


def load_audio(meta):
    """Just load the audio file - I/O bound operation"""
    try:
        meta_id = meta["id"]

        # skip if npz file exists
        # npz_filepath = os.path.join(out_dir, f"{meta_id}.npz")
        # if os.path.exists(npz_filepath):
        #     return None

        start_time = time.time()
        audio, sr = read_from_s3(meta["audio_filepath"], read_f=torchaudio.load)
        print(f"loaded {meta['id']} in {time.time() - start_time:.2f} seconds")

        # crop to max 60s
        audio = audio[:, : int(48000 * 60.01)]

        # ensure stereo
        if audio.shape[0] == 1:
            audio = audio.repeat(2, 1)
        elif audio.shape[0] > 2:
            audio = audio[:2, :]

        # resample to 48k
        audio = torchaudio.functional.resample(audio, sr, 48000)

        # crop to first 30 seconds
        audio_48k = audio[:, : int(48000 * 30.01)]

        if audio_48k.shape[1] < int(48000 * 30.01):
            return None

        # Apply highpass filter
        audio_48k_corrupted = corrupt_audio(audio_48k, 48000)

        # Normalize for VAE
        audio_48k_normalized = audio_48k  # / audio_48k.abs().max().clamp(min=1e-5)
        audio_48k_corrupted_normalized = (
            audio_48k_corrupted  # / audio_48k_corrupted.abs().max().clamp(min=1e-5)
        )

        # Create 24k versions for semantic
        audio_24k = torchaudio.functional.resample(audio_48k, 48000, 24000)
        audio_24k_corrupted = torchaudio.functional.resample(
            audio_48k_corrupted, 48000, 24000
        )

        audio_24k_mono = audio_24k.mean(dim=0).unsqueeze(0)
        audio_24k_corrupted_mono = audio_24k_corrupted.mean(dim=0).unsqueeze(0)

        return {
            "id": meta["id"],
            "audio_48k": audio_48k_normalized,
            "audio_48k_corrupted": audio_48k_corrupted_normalized,
            "audio_24k_mono": audio_24k_mono,
            "audio_24k_corrupted_mono": audio_24k_corrupted_mono,
        }

    except Exception as e:
        print(f"Error loading {meta['id']}: {str(e)}")
        return None


def process_batch(batch_data, codec_model, out_dir, num_vae_tokens: int):
    """Process a batch through models and save results"""
    if not batch_data:
        return

    # Stack for VAE processing
    original_batch = torch.stack([data["audio_48k"] for data in batch_data]).cuda()
    corrupted_batch = torch.stack(
        [data["audio_48k_corrupted"] for data in batch_data]
    ).cuda()

    # Process through VAE
    with torch.no_grad():
        vae_original = codec_model.encode(original_batch)["z"].cpu()
        vae_corrupted = codec_model.encode(corrupted_batch)["z"].cpu()

    # Clear GPU memory
    del original_batch, corrupted_batch
    torch.cuda.empty_cache()

    # Process semantic codes
    semantic_original = encode_semantic([d["audio_24k_mono"] for d in batch_data])[0]
    semantic_corrupted = encode_semantic(
        [d["audio_24k_corrupted_mono"] for d in batch_data]
    )[0]

    # Save results
    for idx, data in enumerate(batch_data):
        vae_orig = vae_original[idx, :, :num_vae_tokens]
        vae_corrupt = vae_corrupted[idx, :, :num_vae_tokens]

        np.savez(
            os.path.join(out_dir, f"{data['id']}.npz"),
            vae_latents_original=vae_orig,
            vae_latents_corrupted=vae_corrupt,
            semantic_codes=semantic_original.astype(np.uint16)[:, 0],
            semantic_codes_corrupted=semantic_corrupted.astype(np.uint16)[:, 0],
        )
        # print(f"saved {len(batch_data)} files")

    # Clear memory
    del vae_original, vae_corrupted, semantic_original, semantic_corrupted
    torch.cuda.empty_cache()


@dataclass
class AudioBatch:
    ids: List[str]
    audio_48k: torch.Tensor
    audio_48k_corrupted: torch.Tensor
    audio_24k_mono: torch.Tensor
    audio_24k_corrupted_mono: torch.Tensor


class AudioDataset(Dataset):
    def __init__(self, metas: List[dict], cache_dir: Optional[str] = None):
        """
        Args:
            metas: List of metadata dictionaries containing audio file information
            cache_dir: Optional directory to cache downloaded files
        """
        self.metas = metas
        self.cache_dir = cache_dir

    def __len__(self) -> int:
        return len(self.metas)

    def __getitem__(self, idx: int) -> Optional[dict]:
        """
        Process a single audio file and return the processed data
        """
        meta = self.metas[idx]

        try:
            result = load_audio(meta)
            if result is not None:
                return result
        except Exception as e:
            print(f"Error processing {meta['id']}: {str(e)}")
            return None


def collate_audio_batch(batch: List[dict]) -> Optional[AudioBatch]:
    """
    Collate function for DataLoader that combines individual results into tensors.
    Filters out None values from failed processing.
    """
    # Filter out None values from failed processing
    batch = [item for item in batch if item is not None]

    if not batch:
        return None

    return AudioBatch(
        ids=[item["id"] for item in batch],
        audio_48k=torch.stack([item["audio_48k"] for item in batch]),
        audio_48k_corrupted=torch.stack(
            [item["audio_48k_corrupted"] for item in batch]
        ),
        audio_24k_mono=torch.stack([item["audio_24k_mono"] for item in batch]),
        audio_24k_corrupted_mono=torch.stack(
            [item["audio_24k_corrupted_mono"] for item in batch]
        ),
    )


def process_dataset(
    codec_model: torch.nn.Module,
    dataset: AudioDataset,
    out_dir: str,
    num_workers: int = 16,
    batch_size: int = 8,
    device: str = "cuda",
) -> None:
    """
    Process the entire dataset using DataLoader for parallel processing.
    """
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        collate_fn=collate_audio_batch,
        shuffle=False,
        pin_memory=True,
    )

    codec_model = codec_model.to(device)
    codec_model.eval()

    for batch in tqdm(loader, desc="Processing audio files"):
        if batch is None:
            continue

        try:
            # Move batch to GPU and process
            audio_batch = batch.audio_48k.to(device)
            audio_corrupted_batch = batch.audio_48k_corrupted.to(device)

            # Process through VAE
            with torch.no_grad():
                vae_original = codec_model.encode(audio_batch)["z"].cpu()
                vae_corrupted = codec_model.encode(audio_corrupted_batch)["z"].cpu()

            # Process semantic codes
            semantic_original = encode_semantic([d for d in batch.audio_24k_mono])[0]
            semantic_corrupted = encode_semantic(
                [d for d in batch.audio_24k_corrupted_mono]
            )[0]

            # Save results for each item in batch
            for idx, id in enumerate(batch.ids):
                vae_orig = vae_original[idx, :, :750]  # num_vae_tokens=750
                vae_corrupt = vae_corrupted[idx, :, :750]

                np.savez(
                    os.path.join(out_dir, f"{id}.npz"),
                    vae_latents_original=vae_orig,
                    vae_latents_corrupted=vae_corrupt,
                    semantic_codes=semantic_original.astype(np.uint16)[:, 0],
                    semantic_codes_corrupted=semantic_corrupted.astype(np.uint16)[:, 0],
                )

        except Exception as e:
            print(f"Error processing batch: {str(e)}")
            continue


def main(
    out_dir: str,
    num_workers: int = 16,
    chunk_size: int = 100,
    num_samples: int = 10000,
):
    VAE_RATE_HZ = 25

    # load the vae model
    device = "cuda:0"
    # device = "cpu"
    if VAE_RATE_HZ == 100:
        checkpoint_filepath = "s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth"
        # checkpoint_filepath = "/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/dac/weights.pth"
        num_vae_tokens = 3000
    elif VAE_RATE_HZ == 25:
        checkpoint_filepath = "s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth"
        num_vae_tokens = 750
    else:
        raise ValueError(f"VAE_RATE_HZ must be 100 or 25, got {VAE_RATE_HZ}")
    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
    }
    codec_model = DAC(**sd["metadata"]["kwargs"])
    codec_model.load_state_dict(sd["state_dict"])
    codec_model.eval()
    codec_model.to(device)

    base_metas = (
        "/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl"
    )
    metas = read_jsonl(base_metas)
    print(f"loaded {len(metas):,} metas")

    # also load the lyric alignments
    genius_alignments_filepath = (
        "/home/tony/Work/tony/hoot/tmp/genius_hq_alignments_t30_v1.jsonl"
    )
    # if os.path.exists(genius_alignments_filepath):
    #    print("loading genius alignments")
    #    genius_alignments = read_jsonl(genius_alignments_filepath, progress=False)
    #    genius_alignments_map = {a[0]: a[1] for a in genius_alignments}
    # else:
    #    genius_alignments_map = {}
    # merge alignments into single map
    # alignments_map = {**genius_alignments_map}
    # print(f"loaded {len(alignments_map):,} alignments")

    # create metas map
    metas_map = {meta["id"]: meta for meta in metas}
    print(f"loaded {len(metas_map):,} metas")

    # Setup
    os.makedirs(out_dir, exist_ok=True)
    random.seed(11)
    random.shuffle(metas)
    subset_metas = metas[:num_samples]

    # Create dataset
    dataset = AudioDataset(metas=subset_metas, cache_dir=os.path.join(out_dir, "cache"))

    # Process dataset
    process_dataset(
        codec_model=codec_model,
        dataset=dataset,
        out_dir=out_dir,
        num_workers=num_workers,
        batch_size=batch_size,
    )


if __name__ == "__main__":
    out_dir = "/home/christian/data/dpo/genius_hq_corrupt_dpo_25hz_30_v2_npz"
    num_samples = 10000
    num_workers = 64
    batch_size = 8
    main(out_dir, num_workers, batch_size, num_samples)