"""
Curry speaker verification: pair scoring and VoxCeleb1-O evaluation.

  python inference.py a.wav b.wav                                   # cosine score for one pair
  python inference.py --trials veri_test.txt --root /path/to/VoxCeleb1/wav/   # EER / minDCF
  python inference.py --repo /local/dir ...                         # use a local copy instead of the Hub
"""
import argparse
import numpy as np
import torch
from tqdm import tqdm
from scipy.interpolate import interp1d
from scipy.optimize import brentq
from sklearn.metrics import roc_curve
from transformers import AutoModel

REPO = "cmu-mlsp/Curry"


def compute_eer(labels, scores):
    fpr, tpr, thr = roc_curve(labels, scores, pos_label=1)
    eer = brentq(lambda x: 1.0 - x - interp1d(fpr, tpr)(x), 0.0, 1.0)
    return eer, float(interp1d(fpr, thr)(eer))


def compute_min_dcf(labels, scores, p_target=0.01, c_miss=1, c_fa=1):
    fpr, tpr, thr = roc_curve(labels, scores, pos_label=1)
    fnr = 1.0 - tpr
    c_det = c_miss * fnr * p_target + c_fa * fpr * (1 - p_target)
    i = int(np.argmin(c_det))
    return c_det[i] / min(c_miss * p_target, c_fa * (1 - p_target)), thr[i]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default=REPO, help="Hub repo id or local directory")
    ap.add_argument("--trials", help="trial file: label enroll_path test_path")
    ap.add_argument("--root", default="", help="prefix prepended to trial paths")
    ap.add_argument("--no-mean-norm", action="store_true", help="disable eval-set embedding mean subtraction")
    ap.add_argument("wavs", nargs="*")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained(args.repo, trust_remote_code=True).eval().to(device)

    if args.trials:
        trials = np.loadtxt(args.trials, str)
        paths = sorted(set(trials[:, 1]) | set(trials[:, 2]))
        embs = {p: model.embed_file(args.root + p).cpu().numpy() for p in tqdm(paths, desc="embedding")}
        if not args.no_mean_norm:
            mu = np.mean(list(embs.values()), axis=0)
            embs = {k: v - mu for k, v in embs.items()}
        labels, scores = [], []
        for lab, a, b in trials:
            ea, eb = embs[a], embs[b]
            scores.append(float(ea @ eb / (np.linalg.norm(ea) * np.linalg.norm(eb))))
            labels.append(int(lab))
        eer, thr = compute_eer(labels, scores)
        print(f"cosine EER: {eer*100:.2f}% (threshold {thr:.3f})")
        for p in (0.01, 0.001):
            dcf, _ = compute_min_dcf(labels, scores, p_target=p)
            print(f"minDCF(p={p}): {dcf:.4f}")
    elif len(args.wavs) == 2:
        print(f"cosine similarity: {model.verify(*args.wavs):.4f}")
    else:
        for w in args.wavs:
            print(w, tuple(model.embed_file(w).shape))


if __name__ == "__main__":
    main()