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

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

if __name__ == "__main__":
    # ckpt_path = (
    #    "/app/suno/christian/ear-logs/ear/6anrtc42/checkpoints/epoch=0-step=5618.ckpt"
    # )
    ckpt_path = "/app/suno/christian/ear-logs/ear/1l1qrmqw/checkpoints/epoch=63-step=359552.ckpt"
    # ckpt_path = "./checkpoints/epoch=3-step=22472.ckpt"

    pos_audio_ref_path = "/app/suno/christian/data/codec_audio/reference-audio-wav_24khz/02 Dreams.input.wav"
    # neg_audio_ref_path = "/app/suno/christian/data/v3_generations_10k_24khz/ff833686-e695-4ed3-91d9-b2beac45ad8c.wav"
    neg_audio_eval_paths = glob.glob(
        "/app/suno/christian/data/codec_audio/reference-audio-wav_24khz/*.input.wav"
    )
    neg_audio_eval_paths = neg_audio_eval_paths[:100]
    # neg_audio_eval_paths = [pos_audio_ref_path]

    mode = "test_generations"
    num_compare = 10

    audio_eval_paths = glob.glob(
        os.path.join(
            "/app/suno/christian/data/codec_audio/genius_hq/*.dac_2c_25x12.wav"
        )
    )

    # audio_eval_paths = glob.glob(
    #    os.path.join("/app/suno/christian/data/v3_generations_10k_24khz/", "*.wav")
    # )
    # pos_audio_eval_paths = pos_audio_eval_paths[:100]

    audio_eval_paths = glob.glob(
        os.path.join("/app/suno/christian/data/trending-05062024/", "*.mp3")
    )

    num_frames = 131072
    system = EarSystem.load_from_checkpoint(ckpt_path)
    system.eval()

    audio_ref = load_audio(
        pos_audio_ref_path,
        num_frames=num_frames,
        target_sample_rate=system.hparams.sample_rate,
    )

    # compare multiple things to reference
    correct = []
    for audio_eval_path in audio_eval_paths:

        audio_eval = load_audio(
            audio_eval_path,
            num_frames=num_frames,
            target_sample_rate=system.hparams.sample_rate,
        )

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

        # run inference
        with torch.no_grad():  # run inference
            # run forward
            pref_preds, quant_preds = system.forward(
                audio_eval.unsqueeze(0),
                audio_ref.unsqueeze(0),
            )
            f_pref_preds = pref_preds.mean(dim=1).squeeze(1)
            f_quant_preds = quant_preds.mean(dim=1).squeeze(1)
            f_pref = torch.sigmoid(f_pref_preds)
            f_quant = torch.argmax(f_quant_preds)

            # pref = f_pref
            # quant = f_quant

            # run reverse
            pref_preds, quant_preds = system.forward(
                audio_ref.unsqueeze(0),
                audio_eval.unsqueeze(0),
            )
            r_pref_preds = pref_preds.mean(dim=1).squeeze(1)
            r_quant_preds = quant_preds.mean(dim=1).squeeze(1)
            r_pref = torch.sigmoid(r_pref_preds)
            r_quant = torch.argmax(r_quant_preds)

            pref = (f_pref + (1 - r_pref)) / 2
            quant = (f_quant + r_quant) / 2
            print(f_pref.item(), r_pref.item(), f_quant.item())

        if pref > 0.5:
            correct.append(True)
        else:
            correct.append(False)

        # overall_pref_preds.append(torch.sigmoid(pref_preds).item())
        # overall_pref_preds.append(quant.item())
        print(f"{np.mean(correct)*100:0.2f} %")