import os
import sys

sys.path.insert(0, "/home/christian/code/christian/scripts")
import json
import glob
import math
import torch
import torchaudio
import numpy as np
import pandas as pd
import pyloudnorm as pyln

from tqdm import tqdm

from train_ear import AudioQualityModel, create_label_encoder


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 AudioDataset(torch.utils.data.Dataset):
    def __init__(self, filepaths, start_s=0.0, duration_s=60.0, target_sr=48000):
        self.filepaths = filepaths
        self.start_s = start_s
        self.duration_s = duration_s
        self.target_sr = target_sr

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

    def __getitem__(self, idx):
        filepath = self.filepaths[idx]
        try:
            audio, sr = torchaudio.load(filepath)
            if sr != self.target_sr:
                audio = torchaudio.functional.resample(audio, sr, self.target_sr)

            # if less than 120 seconds, repeat the audio until we reach target length
            target_length = int(self.target_sr * self.duration_s)
            if audio.shape[1] < target_length:
                # Calculate how many full repeats we need
                num_repeats = math.ceil(target_length / audio.shape[1])
                # Repeat the audio that many times
                audio = audio.repeat(1, num_repeats)
                # Then crop to exact target length
                audio = audio[:, :target_length]
            else:
                # Crop to specified duration
                start_idx = 0
                end_idx = target_length
                audio = audio[:, start_idx:end_idx]

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

            # Normalize
            audio = audio / audio.abs().max()

            return {"audio": audio, "filepath": filepath}
        except Exception as e:
            print(f"Error loading {filepath}: {str(e)}")
            return None


def collate_fn(batch):
    # Filter out None values from failed loads
    batch = [b for b in batch if b is not None]
    if not batch:
        return None

    return {
        "audio": torch.stack([item["audio"] for item in batch]),
        "filepath": [item["filepath"] for item in batch],
    }


if __name__ == "__main__":

    batch_size = 4
    num_workers = 32

    # load manifest
    manifest_filepath = "/app/suno/data/audio_2ch_48khz_lg/ear_train.csv"
    with open(manifest_filepath, "r") as f:
        filepaths = [line.strip() for line in f.readlines()]
    print(len(filepaths))

    # load model
    # model_filepath = (
    #    "/app/suno/christian/checkpoints/ear-v2/2024-12-20_11-26-16_s1551/last_ckpt.pt"
    # )
    model_filepath = (
        "/app/suno/christian/checkpoints/ear-v2/2025-02-14_17-09-21_s8910/last_ckpt.pt"
    )

    ckpt = torch.load(model_filepath)
    model = AudioQualityModel(**ckpt["run_config"]["model"])
    state_dict = ckpt["model"]
    new_state_dict = {}
    for key, value in state_dict.items():
        new_key = key.replace("module.", "")
        new_state_dict[new_key] = value
    model.load_state_dict(new_state_dict)
    model.eval()
    model.cuda()

    # also load corruptions config
    # corruptions = ckpt["corruptions"]
    # label_encoder = create_label_encoder(corruptions)
    # idx_to_label = {idx: label for label, idx in label_encoder.items()}

    # Create dataset and dataloader
    dataset = AudioDataset(filepaths)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        collate_fn=collate_fn,
        shuffle=False,
        pin_memory=True,
    )

    results = []
    results_filepath = (
        "/home/christian/code/christian/metadata/ear_train_corruptions_v2.csv"
    )

    # Create CSV file with headers
    with open(results_filepath, "w") as f:
        f.write("filepath,score\n")

    write_every = 32
    batch_results = []
    with torch.no_grad():
        for batch in tqdm(dataloader):
            if batch is None:
                continue

            # Move batch to GPU
            audio = batch["audio"].cuda()
            filepaths = batch["filepath"]

            # Get predictions
            scores = model.get_score_batch(audio)

            # Collect results
            for i in range(len(filepaths)):
                batch_results.append((filepaths[i], scores[i].item()))

            # Write results every 100 batches
            if len(batch_results) >= write_every * batch_size:
                with open(results_filepath, "a") as f:
                    for filepath, score in batch_results:
                        f.write(f"{filepath},{score}\n")
                batch_results = []

    # Write any remaining results
    if batch_results:
        with open(results_filepath, "a") as f:
            for filepath, score in batch_results:
                f.write(f"{filepath},{score}\n")