#!/usr/bin/env python3
"""
FIGMA inference: encode music and text into a shared space and retrieve.

Examples
--------
# Text -> Audio: rank a folder of .wav clips by a text query
python inference.py --checkpoint figma.ckpt \
    --audio_dir ./clips --query "a song in F minor at 120 BPM in 4/4 time" --topk 5

# Dump audio embeddings for a folder of clips
python inference.py --checkpoint figma.ckpt --audio_dir ./clips --save_embeddings emb.pt
"""

import os
import glob
import argparse

import torch
import librosa

from figma_model import Figma, get_tokenizer, SAMPLE_RATE, CLIP_SECONDS


def load_audio_batch(paths, clip_seconds, sr=SAMPLE_RATE):
    fixed = int(sr * clip_seconds)
    out = torch.zeros(len(paths), 1, fixed)
    for i, p in enumerate(paths):
        wav, _ = librosa.load(p, sr=sr)
        wav = torch.tensor(wav, dtype=torch.float32)
        take = min(wav.numel(), fixed)
        out[i, 0, :take] = wav[:take]
    return out


@torch.no_grad()
def embed_audios(model, paths, clip_seconds, device, batch_size=32):
    embs = []
    for i in range(0, len(paths), batch_size):
        batch = load_audio_batch(paths[i:i + batch_size], clip_seconds).to(device)
        embs.append(model.encode_audio(batch).cpu())
    return torch.cat(embs) if embs else torch.empty(0)


@torch.no_grad()
def embed_texts(model, tokenizer, texts, device, max_len=128, batch_size=64):
    embs = []
    for i in range(0, len(texts), batch_size):
        tok = tokenizer(texts[i:i + batch_size], return_tensors="pt",
                        padding="max_length", truncation=True, max_length=max_len)
        tok = {k: v.to(device) for k, v in tok.items()}
        embs.append(model.encode_text(tok).cpu())
    return torch.cat(embs) if embs else torch.empty(0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", required=True, help="Path to the FIGMA checkpoint (.ckpt)")
    ap.add_argument("--audio_dir", required=True, help="Folder of audio files")
    ap.add_argument("--audio_glob", default="*.wav")
    ap.add_argument("--query", help="Text query for text->audio retrieval")
    ap.add_argument("--topk", type=int, default=5)
    ap.add_argument("--save_embeddings", help="Path to save {paths, audio_emb} .pt")
    ap.add_argument("--clip_seconds", type=float, default=CLIP_SECONDS)
    ap.add_argument("--gpu", type=int, default=0)
    args = ap.parse_args()

    device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu")
    model = Figma.from_checkpoint(args.checkpoint, device=device)
    tokenizer = get_tokenizer()

    paths = sorted(glob.glob(os.path.join(args.audio_dir, args.audio_glob)))
    if not paths:
        raise SystemExit(f"No audio matching {args.audio_glob} in {args.audio_dir}")
    print(f"Encoding {len(paths)} audio clips at {args.clip_seconds}s ...")
    audio_emb = embed_audios(model, paths, args.clip_seconds, device)

    if args.save_embeddings:
        torch.save({"paths": paths, "audio_emb": audio_emb}, args.save_embeddings)
        print(f"Saved embeddings -> {args.save_embeddings}")

    if args.query:
        q = embed_texts(model, tokenizer, [args.query], device)[0]
        scores = audio_emb @ q
        topk = scores.topk(min(args.topk, len(paths)))
        print(f"\nTop {topk.indices.numel()} matches for: {args.query!r}")
        for rank, (idx, sc) in enumerate(zip(topk.indices.tolist(), topk.values.tolist()), 1):
            print(f"  {rank:>2}. {sc:.4f}  {os.path.basename(paths[idx])}")


if __name__ == "__main__":
    main()