#!/usr/bin/env python3
"""Generate a PROPS speaker-embedding GMM from a natural-language description."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch
from sentence_transformers import SentenceTransformer

from model import ComposedGMMMDN


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("description", help="Natural-language speaker-profile description")
    parser.add_argument("--model", default=str(Path(__file__).resolve().parent))
    parser.add_argument("--output", type=Path, default=Path("generated_gmm.npz"))
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()

    encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device=args.device)
    embedding = encoder.encode([args.description], convert_to_tensor=True).to(args.device)
    model = ComposedGMMMDN.from_pretrained(args.model, device=args.device)
    with torch.inference_mode():
        pi_logits, mu, sigma = model(embedding)
        pi = torch.softmax(pi_logits, dim=-1)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        args.output,
        pi_logits=pi_logits[0].cpu().numpy().astype(np.float32),
        pi=pi[0].cpu().numpy().astype(np.float32),
        mu=mu[0].cpu().numpy().astype(np.float32),
        sigma=sigma[0].cpu().numpy().astype(np.float32),
        description=np.asarray(args.description),
    )
    print(args.output)


if __name__ == "__main__":
    main()