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 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


import torch
from torchaudio.transforms import MelSpectrogram


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 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 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 AudioMetadataDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        s3_ids: 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.s3_ids = s3_ids
        self.base_s3_dir = base_s3_dir
        print(len(self.s3_ids))
        os.makedirs(self.tmp_dir, exist_ok=True)

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

    def __getitem__(self, idx):
        s3_id = self.s3_ids[idx]

        if s3_id in self.existing_ids:
            return None

        s3_filepath = f"{self.base_s3_dir}/{s3_id}.mp3"
        filepath = download_audio(s3_filepath, idx, self.tmp_dir)
        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, :]

        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():
            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
            )

        new_meta = {
            "id": s3_id,
            "features": {
                "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}",
            },
        }

        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_audio_production.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["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
    s3_ids = base_metas["id_x"].tolist()
    print(len(s3_ids))

    # create a meta dataset
    meta_dataset = AudioMetadataDataset(s3_ids, existing_ids, tmp_dir="/tmp/cjs")
    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)