import os
import json
import glob
import torch
import random
import torchaudio
import numpy as np

from tqdm import tqdm
from suno_ear.utils import load_audio
from suno_ear.system import EarSystem
from suno_boost.utils import apply_normalization


if __name__ == "__main__":
    num_compare = 5
    num_frames = 131072

    ckpt_path = "/app/suno/christian/ear-logs/ear/1l1qrmqw/checkpoints/epoch=210-step=1185398.ckpt"
    system = EarSystem.load_from_checkpoint(ckpt_path)
    system.eval()

    # search_dir = "/app/suno/christian/data/v3_generations_10k_24khz/"
    search_dir = "/app/suno/data/audio_2ch_48khz_lg/train/genius_hq/"
    # find files in search dir
    search_filepaths = glob.glob(os.path.join(search_dir, "*.wav"))
    # search_filepaths = np.random.choice(search_filepaths, 10000)

    ref_dir = "/app/suno/christian/data/codec_audio/reference-audio-wav_24khz/"
    ref_filepaths = glob.glob(os.path.join(ref_dir, "*.input.wav"))
    ref_filepaths = np.random.choice(ref_filepaths, num_compare)
    # ref_filepaths = [
    #    "/app/suno/christian/data/codec_audio/reference-audio-wav_24khz/02 Dreams.input.wav"
    # ]

    # load
    ref_audios = [
        load_audio(
            filepath,
            num_frames=num_frames,
            target_sample_rate=system.hparams.sample_rate,
        )
        for filepath in ref_filepaths
    ]

    best_results = {}
    worst_results = {}

    # score each file against the references
    for search_filepath in tqdm(search_filepaths):
        search_audio = load_audio(
            search_filepath,
            num_frames=num_frames,
            target_sample_rate=system.hparams.sample_rate,
        )

        if search_audio.shape[-1] < num_frames:
            continue

        overall_pref_preds = []
        overall_quant_preds = []
        for ref_audio in ref_audios:
            with torch.no_grad():  # run inference
                pref_preds, quant_preds = system.forward(
                    search_audio.unsqueeze(0),
                    ref_audio.unsqueeze(0),
                )

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

            overall_pref_preds.append(pref.item())
            overall_quant_preds.append(quant.item())

        overall_pref_pred = np.mean(overall_pref_preds)
        overall_quant_pred = np.mean(overall_quant_preds)

        if overall_pref_pred > 0.5:  # this means the audio is worse than ref
            worst_results[search_filepath] = {
                "pref": overall_pref_pred,
                "quant": overall_quant_pred,
                "score": overall_pref_pred * overall_quant_pred,
            }
        else:
            best_results[search_filepath] = {
                "pref": overall_pref_pred,
                "quant": overall_quant_pred,
                "score": (1 - overall_pref_pred) * overall_quant_pred,
            }

        # print top 5 best and worst sounding audios
        worst_quality = {
            k: v
            for k, v in sorted(
                worst_results.items(), key=lambda item: item[1]["score"], reverse=True
            )
        }
        print("worst quality")
        for idx, (k, v) in enumerate(worst_quality.items()):
            print(idx, k, v)
            if idx > 10:
                break

        best_quality = {
            k: v
            for k, v in sorted(
                best_results.items(), key=lambda item: item[1]["score"], reverse=True
            )
        }
        print()
        print("best quality")
        for idx, (k, v) in enumerate(best_quality.items()):
            print(idx, k, v)
            if idx > 10:
                break

        # save to disk
        with open("worst_quality.json", "w") as fp:
            json.dump(worst_quality, fp, indent=2)
        with open("best_quality.json", "w") as fp:
            json.dump(best_quality, fp, indent=2)