"""LAION-CLAP's music checkpoint on the same benchmark as song_describer.py, for comparison.

    uv run clap_eval.py
"""
import argparse
from pathlib import Path

import laion_clap
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from song_describer import clips, load_captions, load_tracks, print_results

SR = 48000
CHECKPOINT = "music_audioset_epoch_15_esc_90.14.pt"


def normalize(vectors):
    return vectors / np.linalg.norm(vectors, axis=-1, keepdims=True)


@torch.inference_mode()
def embed(model, wave):
    size = 10 * SR
    chunks = np.stack([wave[start:start + size] for start in range(0, len(wave) - size + 1, size)])
    # CLAP was trained on audio that went through int16.
    chunks = (np.clip(chunks, -1, 1) * 32767).astype(np.int16).astype(np.float32) / 32767
    return normalize(model.get_audio_embedding_from_data(chunks).mean(0))


@torch.inference_mode()
def embed_text(model, captions):
    return normalize(np.concatenate([model.get_text_embedding(captions[start:start + 64]) for start in range(0, len(captions), 64)]))


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--data", type=Path, default=Path("data/song_describer"))
    parser.add_argument("--device", default="cuda", help="cuda, mps or cpu")
    args = parser.parse_args()

    captions = load_captions(args.data)
    model = laion_clap.CLAP_Module(enable_fusion=False, amodel="HTSAT-base", device=args.device)
    model.load_ckpt(hf_hub_download("lukewys/laion_clap", CHECKPOINT), verbose=False)

    tracks, middle, average = [], [], []
    for track, wave in load_tracks(args.data, SR):
        three = [embed(model, clip) for clip in clips(wave, SR)]
        tracks.append(track)
        middle.append(three[1])
        average.append(normalize(np.mean(three, axis=0)))

    order = np.argsort(tracks)
    views = {"one 10 s clip": middle, "three 10 s clips": average}
    views = {name: np.stack(audio)[order] for name, audio in views.items()}
    text = embed_text(model, captions.text.tolist())
    print_results("CLAP music", text, views, captions["corpus-id"].to_numpy(), np.array(tracks)[order])


if __name__ == "__main__":
    main()