import argparse
from collections import defaultdict
import json
import numpy as np
import os
import pandas as pd
import pickle
import time
import torch
import torchaudio
from multiprocessing import Pool
from torchaudio.transforms import MelSpectrogram
from tqdm import tqdm
import pyloudnorm as pyln
import traceback
from suno_utils.tasks.ear import load_model as load_ear_model
from suno_utils.audio import Audio
import torch.nn as nn
import torch.nn.functional as F

import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)

## SHIMMER
# git+https://github.com/alphacsc/alphacsc.git@843990555914f2ebbcaef5fc40647d112ed3a51e
try:
    import warnings
    from numba.core.errors import NumbaPerformanceWarning

    warnings.filterwarnings("ignore", category=NumbaPerformanceWarning)
except ImportError:
    pass

## EAR
# Set CUDA device
device = "cuda"
torch.set_num_threads(64)
# neon/ear
# pip install -e .
# s3://suno-data/christian/ear/w5p4nhzn-epoch=27.ckpt
checkpoint_filepath = "/app/suno/data/dpo/models/ear_v2_s3080.pt"
NUM_FRAMES = 131072

# s3.download_file("suno-data", "m4burns/csc_model.pkl", csc_model_path)
csc_model_path = os.path.join("/home/tony/Data/Preference/local_model", "csc_model.pkl")
shimmer_model_path = os.path.join(
    "/home/tony/Data/Preference/local_model", "shimmer_cnn_2025-04-25.pt"
)


class AudioDefectCNN(nn.Module):
    @classmethod
    def from_pretrained(cls, model_name):
        model_path = os.path.join(os.path.dirname(__file__), model_name)
        if not os.path.exists(model_path):
            import boto3

            s3 = boto3.client("s3")
            s3.download_file("suno-data", f"m4burns/{model_name}", model_path)

        state = torch.load(model_path, map_location="cpu", weights_only=True)
        model = cls(
            in_channels=state["in_channels"],
            kernel_size=state["kernel_size"],
            n_filters=state["n_filters"],
            window_size=state["window_size"],
            hop_length=state["hop_length"],
            sample_rate=state["sample_rate"],
            high_cutoff=state["high_cutoff"],
            low_cutoff=state["low_cutoff"],
        )
        model.load_state_dict(state["state_dict"])
        return model

    def __init__(
        self,
        in_channels,
        kernel_size,
        n_filters,
        window_size,
        hop_length,
        sample_rate,
        high_cutoff,
        low_cutoff,
    ):
        super().__init__()

        self.window = np.hanning(window_size)
        self.window_size = window_size
        self.hop_length = hop_length
        self.sample_rate = sample_rate
        self.high_cutoff = high_cutoff
        self.low_cutoff = low_cutoff

        # First conv layer with kernel size matching defect length
        self.conv1 = nn.Conv1d(
            in_channels=in_channels,
            out_channels=n_filters,
            kernel_size=kernel_size,
            stride=1,  # Same length output
            padding=kernel_size // 2,  # Same padding
        )
        self.bn1 = nn.BatchNorm1d(n_filters)

        # Second conv layer for feature extraction
        self.conv2 = nn.Conv1d(
            in_channels=n_filters, out_channels=n_filters * 2, kernel_size=3, padding=1
        )
        self.bn2 = nn.BatchNorm1d(n_filters * 2)

        # Final conv layer to produce logits
        self.conv3 = nn.Conv1d(
            in_channels=n_filters * 2,
            out_channels=1,  # Binary classification
            kernel_size=1,
        )

        self.dropout = nn.Dropout(0.5)

    def spectrogram(self, signal):
        low_bin = int(self.low_cutoff * self.window_size / self.sample_rate)
        high_bin = int(self.high_cutoff * self.window_size / self.sample_rate)

        hops = signal.shape[0] // self.hop_length
        real_spec = np.zeros((hops, high_bin - low_bin + 1))
        for i in range(hops):
            start = i * self.hop_length
            if start + self.window_size > signal.shape[0]:
                break
            real_spec[i, :] = 10 * np.log10(
                np.abs(
                    np.fft.fft(signal[start : start + self.window_size] * self.window)[
                        low_bin : high_bin + 1
                    ]
                )
                + 1e-6
            )
        return real_spec

    def forward(self, x):
        """
        Forward pass

        Args:
            x (torch.Tensor): Input log spectrogram of shape (batch_size, channels, time)

        Returns:
            torch.Tensor: Logits for defect classification
        """
        input_length = x.size(-1)

        # First conv with ReLU activation and batch norm
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)
        assert (
            x.size(-1) == input_length
        ), f"Length mismatch after conv1: got {x.size(-1)}, expected {input_length}"

        # Second conv with ReLU activation and batch norm
        x = self.conv2(x)
        x = self.bn2(x)
        x = F.relu(x)

        x = self.dropout(x)

        # Final conv to produce logits
        x = self.conv3(x)

        # Remove channel dimension
        return x.squeeze(1)


def _load_csc_model():
    return pickle.load(open(csc_model_path, "rb"))


model_dict = _load_csc_model()
csc_model = model_dict["csc"]
shimmer_component_idx = model_dict["shimmer_component_idx"]
sr = model_dict["sr"]
low_cutoff = model_dict["low_cutoff"]
high_cutoff = model_dict["high_cutoff"]
window_size = model_dict["window_size"]
hop_length = model_dict["hop_length"]


def spectrogram(signal):
    window = np.hanning(window_size)
    low_bin = int(low_cutoff * window_size / sr)
    high_bin = int(high_cutoff * window_size / sr)

    hops = signal.shape[0] // hop_length
    real_spec = np.zeros((hops, high_bin - low_bin + 1))
    for i in range(hops):
        start = i * hop_length
        if start + window_size > signal.shape[0]:
            break
        real_spec[i, :] = 10 * np.log10(
            np.abs(
                np.fft.fft(signal[start : start + window_size] * window)[
                    low_bin : high_bin + 1
                ]
            )
            + 1e-6
        )
    return real_spec


def eval_csc(sol, real_specs):
    return sol.transform(real_specs.transpose(0, 2, 1))


def find_peaks(activation):
    threshold = np.max(activation) * 0.5
    peaks = []
    pos = 0
    while pos < len(activation):
        if activation[pos] > threshold:
            start = pos
            while pos < len(activation) and activation[pos] > threshold:
                pos += 1
            peaks.append(np.argmax(activation[start:pos]) + start)
        pos += 1
    return peaks


def _load_shimmer_model():
    shimmer_model = AudioDefectCNN.from_pretrained(shimmer_model_path)
    shimmer_model.to(device)
    shimmer_model.eval()
    return shimmer_model


def get_shimmer_score_from_audio(input_audio: Audio):
    audio_array = input_audio.convert(
        sample_rate=sr, byte_width=2, n_channels=1
    ).array_float
    # start_time = time.time()
    real_spec = spectrogram(audio_array)
    # print(f"Time taken spectrogram: {round(time.time() - start_time, 2)} seconds")
    activations = eval_csc(csc_model, real_spec[None, :, :])
    # print(f"Time taken eval_csc: {round(time.time() - start_time, 2)} seconds")
    shimmer_activation = activations[0][shimmer_component_idx]
    # print(
    #     f"Time taken shimmer_activation: {round(time.time() - start_time, 2)} seconds"
    # )
    peaks = find_peaks(shimmer_activation)
    # print(f"Time taken find_peaks: {round(time.time() - start_time, 2)} seconds")
    score = np.sum(shimmer_activation[peaks]) / (len(shimmer_activation) + 0.00001)
    #   print(f"Time taken score: {round(time.time() - start_time, 2)} seconds")
    return score


def get_shimmer_score_from_audio_v2(shimmer_model, input_audio: Audio):
    audio = input_audio.convert(
        sample_rate=shimmer_model.sample_rate, byte_width=2, n_channels=1
    )
    spec = shimmer_model.spectrogram(audio.array_float)

    spec_torch = torch.from_numpy(spec.T).unsqueeze(0).to(device, dtype=torch.float32)

    with torch.no_grad():
        preds = shimmer_model(spec_torch)
        probs = torch.sigmoid(preds).squeeze().cpu()

    return torch.sum(probs > 0.5).item() / audio.duration_s


def compute_stereo_width_simple(waveform, sample_rate):
    """
    Compute stereo width using a simple time-domain approach.
    Returns a value between 0 (mono) and 1 (maximum stereo spread).

    The calculation is based on comparing the difference signal (L-R)
    to the sum signal (L+R). A higher difference relative to the sum
    indicates more stereo content.

    """
    # Ensure stereo
    if waveform.size(0) == 1:
        raise ValueError("Audio file must be stereo (2 channels)")
    elif waveform.size(0) > 2:
        waveform = waveform[:2, :]  # Take first two channels if more exist

    # Get left and right channels
    left = waveform[0]
    right = waveform[1]

    # Compute difference and sum signals
    difference = left - right
    sum_signal = left + right

    # Compute RMS (Root Mean Square) energy of both signals
    diff_energy = torch.sqrt(torch.mean(difference**2))
    sum_energy = torch.sqrt(torch.mean(sum_signal**2))

    # Compute width as ratio of difference to total energy
    # Normalize to be between 0 and 1
    width = (diff_energy / (sum_energy + 1e-8)).item()
    width = min(width, 1.0)  # Clip to maximum of 1

    return width


def analyze_spectral_balance(waveform, sample_rate):
    """
    Analyze the spectral balance of an audio file and determine if it's
    bassy, mid-focused, or bright, while also computing the spectral centroid.

    Returns:
    - character: 'bassy', 'mid_focused', or 'bright'
    - bass_ratio, mid_ratio, high_ratio: Energy ratios
    - spectral_centroid: The average spectral centroid of the waveform
    """
    # Convert to mono if stereo
    if waveform.size(0) > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)

    # Create mel spectrogram
    mel_spec = MelSpectrogram(
        sample_rate=sample_rate,
        n_fft=2048,
        hop_length=512,
        n_mels=128,
        f_min=20,
        f_max=20000,
    )(waveform)

    # Convert to dB scale
    mel_spec_db = torch.log10(mel_spec + 1e-9)

    # Calculate average energy in each frequency band
    bass_energy = torch.mean(mel_spec_db[:, :40]).item()  # ~20-250 Hz
    mid_energy = torch.mean(mel_spec_db[:, 40:80]).item()  # ~250-4000 Hz
    high_energy = torch.mean(mel_spec_db[:, 80:]).item()  # ~4000-20000 Hz

    # Calculate relative ratios
    total_energy = bass_energy + mid_energy + high_energy
    bass_ratio = bass_energy / total_energy
    mid_ratio = mid_energy / total_energy
    high_ratio = high_energy / total_energy

    # Determine dominant characteristic
    if bass_ratio > max(mid_ratio, high_ratio):
        character = "bassy"
    elif mid_ratio > max(bass_ratio, high_ratio):
        character = "mid_focused"
    else:
        character = "bright"

    # Spectral Centroid computation
    # Frequency bins for the mel spectrogram
    mel_frequencies = torch.linspace(20, 20000, 128)
    spectral_centroid = torch.sum(
        mel_frequencies * torch.mean(mel_spec, dim=-1)
    ) / torch.sum(torch.mean(mel_spec, dim=-1))
    spectral_centroid = spectral_centroid.item()

    return character, bass_ratio, mid_ratio, high_ratio, spectral_centroid


def calculate_average_spectrum_db(waveform, n_fft=16384, hop_length=8192):
    # Keep as torch tensor or convert to torch tensor if it's numpy
    if not isinstance(waveform, torch.Tensor):
        waveform = torch.from_numpy(waveform)

    # Calculate the average spectrum using STFT for efficiency
    n_fft = 2048  # Choose an appropriate FFT size
    hop_length = n_fft // 4  # Standard hop length

    # Compute STFT using torch
    if waveform.dim() > 1:
        # For stereo, compute STFT for each channel
        stft_results = []
        for channel in range(waveform.shape[0]):
            stft = torch.stft(
                waveform[channel],
                n_fft=n_fft,
                hop_length=hop_length,
                window=torch.hann_window(n_fft),
                return_complex=True,
            )
            # Get magnitude
            stft_magnitude = torch.abs(stft)
            stft_results.append(stft_magnitude)

        # Average across time frames for each channel
        magnitude_spectrum = torch.stack(
            [torch.mean(stft, dim=1) for stft in stft_results]
        )
    else:
        # For mono
        stft = torch.stft(
            waveform,
            n_fft=n_fft,
            hop_length=hop_length,
            window=torch.hann_window(n_fft),
            return_complex=True,
        )
        # Get magnitude
        stft_magnitude = torch.abs(stft)
        magnitude_spectrum = torch.mean(stft_magnitude, dim=1)

    # Convert to dB scale
    spectrum_db = 20 * torch.log10(
        magnitude_spectrum + 1e-10
    )  # Adding small value to avoid log(0)

    # Convert to numpy for consistency with the rest of the code
    return spectrum_db.numpy()


def calculate_decay(
    audio_obj_first: Audio, audio_obj_last: Audio, n_fft=16384, hop_length=8192
):
    spectrum_first = calculate_average_spectrum_db(
        audio_obj_first.array_float, n_fft, hop_length
    )
    spectrum_last = calculate_average_spectrum_db(
        audio_obj_last.array_float, n_fft, hop_length
    )
    difference = spectrum_last - spectrum_first
    decay = np.sum(np.abs(difference))
    return decay


def analyze_loudness_factor(waveform, sample_rate):
    """
    Analyze loudness factor of an audio file and provide descriptive characteristics.
    Loudness factor is the LUFS measurement after peak normalization.

    """

    # peak normalize
    normalized_audio = waveform / torch.max(torch.abs(waveform)).clamp(min=1e-9)

    # Measure LUFS
    meter = pyln.Meter(sample_rate)
    loudness_factor = meter.integrated_loudness(normalized_audio.permute(1, 0).numpy())

    # Categorize and select descriptors
    if loudness_factor < -16:
        category = "dynamic"
    elif loudness_factor < -10:
        category = "moderate"
    else:
        category = "compressed"

    return category, loudness_factor


def analyze_clipping(audio_tensor, threshold=0.99, sample_rate=44100):
    """
    Vectorized analysis of audio clipping artifacts after peak normalization.

    Parameters:
        audio_tensor: torch.Tensor
            Audio samples (shape: [channels, samples] or [samples])
        threshold: float
            Threshold for considering a sample clipped (0.0 to 1.0)
        sample_rate: int
            Audio sample rate in Hz, used to normalize clip percentage

    Returns:
        tuple (total_clipped_samples, clips_per_second)
    """
    if not isinstance(audio_tensor, torch.Tensor):
        audio_tensor = torch.tensor(audio_tensor)

    if audio_tensor.dim() == 1:
        audio_tensor = audio_tensor.unsqueeze(0)

    # Peak normalize
    original_peak = audio_tensor.abs().max().clamp(min=1e-9)
    audio_tensor = audio_tensor / original_peak

    # Find clipped samples across all channels
    all_clips = audio_tensor.abs() >= threshold

    # Calculate total clipped samples (max across channels)
    total_clips = all_clips.sum(dim=-1).max().item()

    # Calculate clips per second
    duration_seconds = audio_tensor.shape[-1] / sample_rate
    clips_per_second = total_clips / duration_seconds

    return total_clips, clips_per_second


def load_ear_system():
    ear_system = load_ear_model(checkpoint_filepath)
    ear_system.to(device)
    ear_system.eval()
    return ear_system


def compare_quality(ear_system, audio_a: torch.Tensor, audio_b: torch.Tensor):
    """Compare the quality of two audio files using the trained model.

    Parameters
    ----------
    system : EarSystem
        the trained model
    audio_a : torch.Tensor
        audio tensor of shape (2, num_frames)
    audio_b : torch.Tensor
        audio tensor of shape (2, num_frames)

    Returns
    -------
    pref : torch.Tensor
        preference prediction

    quant : torch.Tensor
        quantification prediction

    """

    # move audio_a and audio_b to same device as system parameters
    audio_a = audio_a.to(ear_system.device)
    audio_b = audio_b.to(ear_system.device)

    # first, embed the audio that will be evaluated
    with torch.no_grad():
        embeds_a = ear_system.embed(audio_a)
        embeds_b = ear_system.embed(audio_b)

    # aggregate embeddings over time with a moving mean of frame size
    embeds_a = torch.nn.functional.adaptive_avg_pool1d(
        embeds_a.permute(0, 2, 1), 137
    ).permute(0, 2, 1)
    embeds_b = torch.nn.functional.adaptive_avg_pool1d(
        embeds_b.permute(0, 2, 1), 137
    ).permute(0, 2, 1)

    # concat embeds into singular tensors
    embeds = torch.cat((embeds_a, embeds_b), dim=-1)

    # no run through the projection to make predictions
    with torch.no_grad():
        pref_preds = ear_system.pref_classifier(embeds)
        quant_preds = ear_system.quant_classifier(embeds)

    # print(pref_preds.shape, quant_preds.shape)

    # get a final score by taking mean across seq of preds and chunks
    pref_preds = pref_preds.mean(dim=1).mean(dim=0)
    quant_preds = quant_preds.mean(dim=1).mean(dim=0)
    pref = torch.sigmoid(pref_preds)
    quant = torch.argmax(quant_preds, dim=0).float()

    return pref, quant


def prepare_audio(
    input_audio: Audio,
    num_frames: int = NUM_FRAMES,
    start_s: float | None = None,
    end_s: float | None = None,
):
    sample_rate = input_audio.sample_rate
    audio = torch.from_numpy(input_audio.array_float)

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

    # crop audio based on metadata example
    if start_s is not None and end_s is not None:
        start_frame = int(start_s * sample_rate)
        end_frame = int(end_s * sample_rate)
        audio = audio[:, start_frame:end_frame]

    # if the file is long, only take part of it
    if audio.shape[-1] > (sample_rate * 120):
        audio = audio[:, : sample_rate * 120]

    # downmix and resample decoded audio to 24khz
    audio = torchaudio.functional.resample(audio, sample_rate, 24_000)

    # pad by repeating the signal if shorter than window
    if audio.shape[-1] < num_frames:
        pad_size = num_frames - audio.shape[-1]
        audio = torch.nn.functional.pad(audio, (1, pad_size), mode="replicate")

    # chunk into non-overlapping blocks of num_frames
    audio_chunks = []
    num_chunks = audio.shape[-1] // num_frames
    for n in range(num_chunks):
        start_idx = n * num_frames
        end_idx = start_idx + num_frames
        audio_chunks.append(audio[:, start_idx:end_idx])

    # loudness norm
    meter = pyln.Meter(24_000)

    for audio_chunk_idx in range(len(audio_chunks)):
        x_lufs_db = meter.integrated_loudness(audio.T.numpy())
        if x_lufs_db == -float("inf"):
            gain_lin = 1.0
        else:
            delta_lufs_db = -20.0 - x_lufs_db
            gain_lin = 10.0 ** (np.clip(delta_lufs_db, a_min=-120, a_max=48.0) / 20.0)
        audio_chunks[audio_chunk_idx] *= gain_lin

    # take the last chunk keeping the list
    # audio_chunks = audio_chunks[-1:]

    return torch.stack(audio_chunks)


def get_audio_quality_scores(input_audio: Audio):
    sample_rate = input_audio.sample_rate
    audio = torch.from_numpy(input_audio.array_float)
    # 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, :]

    # do the analysis here
    with torch.no_grad():
        category, loudness_factor = analyze_loudness_factor(audio, sample_rate)
        spectral_character, bass_ratio, mid_ratio, high_ratio, spectral_centroid = (
            analyze_spectral_balance(audio, sample_rate)
        )
        stereo_width = compute_stereo_width_simple(audio, sample_rate)
        total_clips, clips_per_second = analyze_clipping(audio, sample_rate=sample_rate)

    return {
        "loudness_factor": f"{loudness_factor:0.3f}",
        "spectral_character": spectral_character,
        "spectral_centroid": f"{spectral_centroid:0.3f}",
        "bass_ratio": f"{bass_ratio:0.3f}",
        "mid_ratio": f"{mid_ratio:0.3f}",
        "high_ratio": f"{high_ratio:0.3f}",
        "stereo_width": f"{stereo_width:0.3f}",
        "total_clips": total_clips,
        "clips_per_second": f"{clips_per_second:0.3f}",
    }


def process_request_batch(request_batch):
    try:
        ear_system = load_ear_system()
        shimmer_model = _load_shimmer_model()
        results = defaultdict(dict)
        for request_id, negative_id, positive_id in tqdm(request_batch):
            try:
                full_negative_audio = Audio.from_s3(
                    f"s3://suno-data-uploads/studio/uploads/{negative_id}.mp3",
                    n_channels=2,
                )
                full_positive_audio = Audio.from_s3(
                    f"s3://suno-data-uploads/studio/uploads/{positive_id}.mp3",
                    n_channels=2,
                )
                # n_total_chunks = min(
                #     max(
                #         int(
                #             min(
                #                 full_negative_audio.duration_s,
                #                 full_positive_audio.duration_s,
                #             )
                #             // 30
                #         ),
                #         1,
                #     ),
                #     2,  # take maximum of two chunks
                # )
                n_total_chunks = max(
                    int(
                        min(
                            full_negative_audio.duration_s,
                            full_positive_audio.duration_s,
                        )
                        // 30
                    ),
                    1,
                )
                negative_average_spectrum = 0
                positive_average_spectrum = 0
                for local_chunk_idx in range(n_total_chunks):
                    negative_audio = full_negative_audio.get_segment(
                        local_chunk_idx * 30, (local_chunk_idx + 1) * 30
                    )
                    positive_audio = full_positive_audio.get_segment(
                        local_chunk_idx * 30, (local_chunk_idx + 1) * 30
                    )
                    # print(f"Processing {local_chunk_idx} of {n_total_chunks}")
                    negative_ear_quality_scores, _ = ear_system.get_score(
                        negative_audio, return_scores=True
                    )
                    positive_ear_quality_scores, _ = ear_system.get_score(
                        positive_audio, return_scores=True
                    )
                    negative_shimmer_score = get_shimmer_score_from_audio_v2(
                        shimmer_model, negative_audio
                    )
                    positive_shimmer_score = get_shimmer_score_from_audio_v2(
                        shimmer_model, positive_audio
                    )
                    negative_quality_scores = get_audio_quality_scores(negative_audio)
                    positive_quality_scores = get_audio_quality_scores(positive_audio)
                    negative_spectrum_decay = 0
                    positive_spectrum_decay = 0
                    if local_chunk_idx == 0:
                        negative_average_spectrum = calculate_average_spectrum_db(
                            negative_audio.array_float
                        )
                        positive_average_spectrum = calculate_average_spectrum_db(
                            positive_audio.array_float
                        )
                    else:
                        current_negative_average_spectrum = (
                            calculate_average_spectrum_db(negative_audio.array_float)
                        )
                        negative_spectrum_decay = np.sum(
                            np.abs(
                                current_negative_average_spectrum
                                - negative_average_spectrum
                            )
                        )

                        current_positive_average_spectrum = (
                            calculate_average_spectrum_db(positive_audio.array_float)
                        )
                        positive_spectrum_decay = np.sum(
                            np.abs(
                                current_positive_average_spectrum
                                - positive_average_spectrum
                            )
                        )

                    # keep the first index for easy filtering
                    results[request_id][
                        negative_id
                        + ("" if local_chunk_idx == 0 else f"_{local_chunk_idx}")
                    ] = {
                        "ear_v2_quality_scores": [
                            round(current_score, 4)
                            for current_score in negative_ear_quality_scores
                        ],
                        "shimmer_score": negative_shimmer_score,
                        "abs_loudness_factor": round(
                            float(
                                get_audio_quality_scores(negative_audio)[
                                    "loudness_factor"
                                ]
                            ),
                            4,
                        ),
                        "spectrum_decay": float(negative_spectrum_decay),
                        **negative_quality_scores,
                    }
                    results[request_id][
                        positive_id
                        + ("" if local_chunk_idx == 0 else f"_{local_chunk_idx}")
                    ] = {
                        "ear_v2_quality_scores": [
                            round(current_score, 4)
                            for current_score in positive_ear_quality_scores
                        ],
                        "shimmer_score": positive_shimmer_score,
                        "abs_loudness_factor": round(
                            float(
                                get_audio_quality_scores(positive_audio)[
                                    "loudness_factor"
                                ]
                            ),
                            4,
                        ),
                        "spectrum_decay": float(positive_spectrum_decay),
                        **positive_quality_scores,
                    }
            except Exception as e:
                print(f"Error processing request {request_id}: {e}")
                traceback.print_exc()
                results[request_id][negative_id] = None
                results[request_id][positive_id] = None
        return results
    except Exception as e:
        print(f"Error in process_request_batch: {e}")
        traceback.print_exc()
        return {}


# test case
def main_test():
    ## Example usage
    test_batch = [
        (
            "0002c166-7355-4335-bc47-414ed0914f82",  # request_id
            "22bf283f-982b-40c3-aecc-3878c5583062",  # negative_id
            "30c532ed-9422-4db3-8f4c-2b56c3aed34a",  # positive_id
        )
    ]
    test_result = process_request_batch(test_batch)
    print(test_result)
    for request_id, results in test_result.items():
        for negative_id, result in results.items():
            print(f"{request_id} - {negative_id}")
            print(result["shimmer_score"])
            print("-" * 100)
        for positive_id, result in results.items():
            print(f"{request_id} - {positive_id}")
            print(result["shimmer_score"])
            print("-" * 100)
    return


def main():
    total_job_n_gpus = 8
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file_name", type=str)
    parser.add_argument("--job_idx", type=int, default=0)
    args = parser.parse_args()
    print(f"CUDA_VISIBLE_DEVICES: {os.environ['CUDA_VISIBLE_DEVICES']}")
    start_time = time.time()
    # Load pretrained ear modelimport pandas as pd
    # input_folder_path = "/home/tony/Data/Preference/up_v1"
    # input_file_name = "interesting_clips_up_u_1_20250125_full.pkl"
    # input_folder_path = "/home/tony/Data/Preference/up_diff2_v1"
    # input_folder_path = "/home/tony/Data/Preference/up_v2_d5/"
    # input_folder_path = "/home/tony/Data/Preference/carp_t1/"
    input_folder_path = "/home/tony/Data/Preference/dorado_t1/"
    # input_file_name = "interesting_clips_up_u_4_20250215_full.pkl"
    input_file_name = args.input_file_name
    output_file_name = "full_pair_quality"
    df = pd.read_pickle(f"{input_folder_path}/{input_file_name}")
    df = df.sort_values(by=["request_id", "preference"])
    print("Preference data shape", df.shape)
    request_jobs = []
    assert df["preference"].nunique() == 2
    current_requests = []
    df["request_id"] = df["request_id"].astype(str)
    df["s3_id"] = df["s3_id"].astype(str)
    for row_id, row in df.iterrows():
        if row["preference"] == 0 and row_id % 2 == 0:
            current_requests.append(row["request_id"])
            current_requests.append(row["s3_id"])
        elif row["preference"] == 1 and row_id % 2 == 1:
            current_requests.append(row["s3_id"])
            request_jobs.append(current_requests)
            current_requests = []
        else:
            raise ValueError(f"Invalid preference: {row['preference']}")

    request_jobs = sorted(request_jobs)
    print(f"Total jobs (requests): {len(request_jobs)}")
    # find the specific chunk
    request_jobs = request_jobs[
        (len(request_jobs) // total_job_n_gpus) * args.job_idx : (
            len(request_jobs) // total_job_n_gpus
        )
        * (args.job_idx + 1)
    ]
    print(f"Chunked jobs (requests): {len(request_jobs)}")

    if os.path.exists(f"{input_folder_path}/{output_file_name}.json"):
        with open(f"{input_folder_path}/{output_file_name}.json", "r") as f:
            known_results = json.load(f)
    else:
        known_results = {}

    print(
        f"Pre-filtered jobs: {len(request_jobs)}, known results: {len(known_results)}"
    )
    request_jobs = [job for job in request_jobs if job[0] not in known_results.keys()]

    print(f"Total jobs: {len(request_jobs)}")
    # only when you debug
    # request_jobs = request_jobs[:20]

    # Using process_map from tqdm.contrib.concurrent for better parallelization
    # Split request_jobs into batches
    n_processes = 16
    batch_size = len(request_jobs) // n_processes + 1
    batches = [
        request_jobs[i : i + batch_size]
        for i in range(0, len(request_jobs), batch_size)
        # for i in range(0, 10, batch_size)
    ]
    print(f"Total batches: {len(batches)}, batch size: {batch_size}")

    # Use ThreadPool instead of Pool for multi-threading rather than multi-processing
    from concurrent.futures import ThreadPoolExecutor

    with ThreadPoolExecutor(max_workers=n_processes) as executor:
        results = list(executor.map(process_request_batch, batches))

    # Combine results from all batches
    combined_results = {}
    for batch_result in results:
        combined_results.update(batch_result)

    known_results.update(combined_results)
    # with open(f"{input_folder_path}/pair_quality_{args.job_idx}.json", "w") as f:
    if total_job_n_gpus > 1:
        with open(
            f"{input_folder_path}/{output_file_name}_{args.job_idx}.json", "w"
        ) as f:
            json.dump(known_results, f, indent=4)
    else:
        with open(f"{input_folder_path}/{output_file_name}.json", "w") as f:
            json.dump(known_results, f, indent=4)
    print(
        f"DONE!! Processed {len(combined_results)} -- to total {len(known_results)}. Total time: {round(time.time() - start_time, 2)}s"
    )


if __name__ == "__main__":
    # Make sure export is setup: export CUDA_VISIBLE_DEVICES=5
    main()
    # main_test()