import os
import json
import torch
import torchaudio
import pandas as pd
import pyloudnorm as pyln

from tqdm import tqdm
from suno_utils.utils.s3 import read_from_s3
from suno_utils.utils.text import read_jsonl, write_jsonl
from torchaudio.transforms import MelSpectrogram


def measure_similarity(audio_a, audio_b, n_fft=2048, hop_length=512, n_mels=64):
    a_mfccs = torchaudio.transforms.MFCC(
        sample_rate=48000,
        log_mels=True,
        melkwargs={
            "n_fft": n_fft,
            "hop_length": hop_length,
            "n_mels": n_mels,
        },
    )(audio_a)
    b_mfccs = torchaudio.transforms.MFCC(
        sample_rate=48000,
        log_mels=True,
        melkwargs={
            "n_fft": n_fft,
            "hop_length": hop_length,
            "n_mels": n_mels,
        },
    )(audio_b)
    return torch.nn.functional.mse_loss(a_mfccs, b_mfccs)


def download_audio(s3_filepath: str, example_id: str, tmp_dir: str):
    filename = os.path.basename(s3_filepath)
    out_filepath = os.path.join(tmp_dir, f"{example_id}-{filename}")
    # only download the file if its not already downloaded
    if not os.path.isfile(out_filepath):
        os.system(f"aws s3 cp {s3_filepath} {out_filepath} > /dev/null 2>&1")
    return out_filepath


class AudioPairsMetadataDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        pairs: list,
        existing_ids: list = None,
        tmp_dir: str = "/mnt/localdisk/tmp/cjs",
        base_s3_dir: str = "s3://suno-data-uploads/studio/uploads",
    ):
        self.tmp_dir = tmp_dir
        self.existing_ids = existing_ids
        self.pairs = pairs
        self.base_s3_dir = base_s3_dir
        print(len(self.pairs))
        os.makedirs(self.tmp_dir, exist_ok=True)

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

    def __getitem__(self, idx):
        pair = self.pairs[idx]

        if pair[0]["request_id"] in self.existing_ids:
            return None

        negative_s3_filepath = f"{self.base_s3_dir}/{pair[0]['id_x']}.mp3"
        negative_filepath = download_audio(negative_s3_filepath, idx, self.tmp_dir)

        positive_s3_filepath = f"{self.base_s3_dir}/{pair[1]['id_x']}.mp3"
        positive_filepath = download_audio(positive_s3_filepath, idx, self.tmp_dir)

        audios = []
        for filepath in [negative_filepath, positive_filepath]:
            try:
                audio, sample_rate = torchaudio.load(filepath)
                # crop audio to max of 30sec
                audio = audio[:, : int(30 * sample_rate)]

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

                audios.append(audio)

            except Exception as e:
                print(f"Error: {e}")
                print(f"Failed to process: {filepath}")
                return None

            # delete audio file
            # os.remove(filepath)

        # do the analysis here
        with torch.no_grad():
            mfcc_similarity = measure_similarity(audios[0], audios[1])

        new_meta = {
            "request_id": pair[0]["request_id"],
            "mfcc_similarity": mfcc_similarity.item(),
        }

        return new_meta


def collate_fn(batch):
    return batch


if __name__ == "__main__":

    batch_size = 128
    num_workers = 200

    base_metas_filepath = (
        "/home/tony/Data/Preference/up_v1/interesting_clips_up_u_1_20241201_full.pkl"
    )
    out_metas_filepath = "/home/christian/code/christian/metadata/interesting_clips_up_u_1_20241201_full_mfcc_similarity.jsonl"

    if os.path.isfile(out_metas_filepath):
        print(f"Output file already exists: {out_metas_filepath}")
        # load the existing file
        existing_metas = read_jsonl(out_metas_filepath)
        print(f"Loaded {len(existing_metas)} existing metas")
        existing_ids = [meta["request_id"] for meta in existing_metas]
    else:
        existing_ids = []

    base_metas = pd.read_pickle(base_metas_filepath)

    print(len(base_metas))

    # get the s3 filepaths
    # collect the positive and negative pairs from the base_meta
    # they are stored in order with the negative first, followed by the positive
    # iterate over the base_meta and collect the pairs (skip by 2)
    pairs = []
    for i in range(0, len(base_metas), 2):
        pairs.append((base_metas.iloc[i], base_metas.iloc[i + 1]))

    # create a meta dataset
    meta_dataset = AudioPairsMetadataDataset(
        pairs, existing_ids, tmp_dir="/app/suno/christian/dpo/mp3"
    )
    meta_dataloader = torch.utils.data.DataLoader(
        meta_dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    for new_metas in tqdm(meta_dataloader):
        new_metas = [meta for meta in new_metas if meta is not None]
        write_jsonl(new_metas, out_metas_filepath, do_append=True)