#!/usr/bin/env python3
"""
Predict the theme of Hindi songs (romance / sadness / celebration / motivational) with the models
of meet5568/meet_models.

For a NEW song, every one of the 5 cross-validation folds has its own lyrics encoder and its own
classifier; the script runs all 5 and averages their probabilities.

Examples
  # lyrics only (downloads the models from Hugging Face the first time; needs access to the repo)
  python predict.py --lyrics "tujhe dekha toh yeh jaana sanam, pyaar hota hai deewana sanam"
  python predict.py --lyrics-file song.txt
  # use an already downloaded copy of the repo
  python predict.py --model-dir /path/to/meet_models --lyrics-file song.txt
  # lyrics + audio: an .npz with arrays "instrument" and "mix" (and/or "vocal"), each (n_songs, 9984)
  python predict.py --lyrics-file song.txt --audio-npz song_audio.npz --combination lyrics+instrument+mix

Python API
  from predict import ThemePredictor
  p = ThemePredictor()                         # or ThemePredictor(model_dir="...")
  p.predict(["lyrics of song 1", "lyrics of song 2"])               # DataFrame of probabilities
  p.predict(lyrics, audio={"instrument": A, "mix": B}, combination="lyrics+instrument+mix")
"""

import argparse
import json
import os
import sys

import numpy as np
import pandas as pd
import torch

REPO = "meet5568/meet_models"
THEMES = ["romance", "sadness", "celebration", "motivational"]
AUDIO = ("vocal", "instrument", "mix")


def get_model_dir(model_dir=None, need_encoders=True, need_classifiers=True):
    """A local copy of the repo; downloads (only the needed parts) if no path is given."""
    if model_dir:
        return model_dir
    from huggingface_hub import snapshot_download
    patterns = ["code/*", "README.md"]
    if need_encoders:
        patterns.append("lyrics_encoder/theme_tuned/*")
    if need_classifiers:
        patterns.append("classification/*")
    # private repo: run `huggingface-cli login` once, or set the HF_TOKEN environment variable
    return snapshot_download(REPO, allow_patterns=patterns, token=os.environ.get("HF_TOKEN"))


class ThemePredictor:
    def __init__(self, model_dir=None, device=None):
        self.dir = get_model_dir(model_dir)
        sys.path.insert(0, os.path.join(self.dir, "code"))
        import classify  # model definitions (same code that trained the models)
        self.C = classify
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

    def encode_lyrics(self, texts, fold):
        from sentence_transformers import SentenceTransformer
        enc = SentenceTransformer(os.path.join(self.dir, "lyrics_encoder", "theme_tuned", f"fold_{fold}"),
                                  device=self.device)
        enc.max_seq_length = 400
        return enc.encode(texts, batch_size=16, normalize_embeddings=True, convert_to_numpy=True)

    def _fold_proba(self, combination, fold, inputs):
        d = os.path.join(self.dir, "classification", combination, f"fold_{fold}")
        info = json.load(open(os.path.join(d, "config.json")))
        if info["method"] == "late":                  # average the single-modality models
            return np.mean([self._fold_proba(m, fold, inputs) for m in info["modalities"]], axis=0)
        norm = np.load(os.path.join(d, "norm.npz"))
        model = self.C.build_model(info["run"], info["modalities"], info["in_dims"], info["model_config"])
        model.load_state_dict(torch.load(os.path.join(d, "model.pt"), map_location="cpu"))
        model = model.to(self.device).eval()
        xs = [torch.tensor((inputs[m] - norm[f"{m}_mean"]) / norm[f"{m}_std"], dtype=torch.float32,
                           device=self.device) for m in info["modalities"]]
        with torch.no_grad():
            return self.C.predict_proba(model, xs)

    def predict(self, lyrics=None, audio=None, combination=None):
        """lyrics: list of strings (or None); audio: dict modality -> (n, 9984) array (or None).
        combination: e.g. "lyrics", "lyrics+instrument+mix", "mix"; default = all given inputs."""
        audio = {k: np.asarray(v, dtype=np.float32) for k, v in (audio or {}).items()}
        if lyrics is not None and isinstance(lyrics, str):
            lyrics = [lyrics]
        given = (["lyrics"] if lyrics is not None else []) + [m for m in AUDIO if m in audio]
        combination = combination or "+".join(given)
        mods = combination.split("+")
        missing = [m for m in mods if m not in given]
        if missing:
            raise ValueError(f"combination {combination} needs inputs {missing}")
        n = len(lyrics) if lyrics is not None else len(next(iter(audio.values())))
        probs = []
        for k in range(5):
            inputs = dict(audio)
            if "lyrics" in mods:
                inputs["lyrics"] = self.encode_lyrics(lyrics, k)
            probs.append(self._fold_proba(combination, k, inputs))
        p = np.mean(probs, axis=0)
        out = pd.DataFrame(p, columns=THEMES)
        out.insert(0, "predicted_theme", [THEMES[i] for i in p.argmax(1)])
        assert len(out) == n
        return out


def main():
    ap = argparse.ArgumentParser(description="Predict the theme of Hindi songs")
    ap.add_argument("--lyrics", help="lyrics text of one song")
    ap.add_argument("--lyrics-file", help="text file with the lyrics of one song")
    ap.add_argument("--audio-npz", help=".npz with arrays vocal / instrument / mix, each (n_songs, 9984)")
    ap.add_argument("--combination", default=None, help="e.g. lyrics, lyrics+instrument+mix (default: all inputs given)")
    ap.add_argument("--model-dir", default=None, help="local copy of meet5568/meet_models (default: download)")
    args = ap.parse_args()
    lyrics = args.lyrics or (open(args.lyrics_file, encoding="utf-8").read() if args.lyrics_file else None)
    audio = dict(np.load(args.audio_npz)) if args.audio_npz else None
    if lyrics is None and audio is None:
        ap.error("give --lyrics / --lyrics-file and/or --audio-npz")
    p = ThemePredictor(args.model_dir).predict([lyrics] if lyrics else None, audio, args.combination)
    print(p.round(3).to_string(index=False))


if __name__ == "__main__":
    main()