#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright      2026  AXERA-TECH  (authors: Magnetar)
#
# CAMPPlus speaker embedding demo (AXera NPU).
#
# Mirrors the Python flows of 3D-Speaker-MT.axera:
#   1) 1:1 speaker verification (export_campplus_onnx.py demo):
#        python3 example.py --models-dir models --wav1 a.wav --wav2 b.wav
#      same speaker ~0.67, different speakers ~0.06
#   2) chunked embedding extraction + speaker clustering (ax_cam_bin.py):
#        python3 example.py --models-dir models --audio wav/vad_example.wav \
#            --diarize --speaker-num 0

import argparse
import os
import time

import numpy as np

from campplus_sdk import CampplusModel, chunk, cosine_similarity
from campplus_sdk.inference import load_wav


def parse_args():
    parser = argparse.ArgumentParser(description="CAMPPlus speaker embedding demo")
    parser.add_argument("--models-dir", type=str, default="models",
                        help="Directory containing campplus.axmodel")
    parser.add_argument("--wav1", type=str, default=None,
                        help="First wav for 1:1 speaker verification")
    parser.add_argument("--wav2", type=str, default=None,
                        help="Second wav for 1:1 speaker verification")
    parser.add_argument("--audio", type=str, default=None,
                        help="Audio for chunked embedding extraction / diarization")
    parser.add_argument("--diarize", action="store_true",
                        help="Run speaker clustering (requires clustering deps)")
    parser.add_argument("--speaker-num", type=int, default=0,
                        help="Known speaker count (0 = auto, requires clustering deps)")
    return parser.parse_args()


def load_wav_16k(wav_file):
    """Read wav, resample if needed, return mono [T] tensor."""
    return load_wav(wav_file, target_sr=16000)


def main():
    args = parse_args()
    if not args.wav1 and not args.audio:
        raise SystemExit("Specify --wav1/--wav2 or --audio (see --help)")

    t0 = time.time()
    model = CampplusModel(args.models_dir)
    print(f"Model loaded in {time.time() - t0:.2f}s")

    if args.audio:
        speech = load_wav_16k(args.audio)
        speech = speech.numpy()
        if speech.ndim > 1:
            speech = speech[0]
        fs = 16000
        audio_duration = len(speech) / fs
        print(f"Audio duration: {audio_duration:.2f}s")

        # VAD-free chunking identical to ax_meeting_transc_demo.py
        chunks = chunk(0, audio_duration)
        t1 = time.time()
        embeddings = model(speech, fs, chunks=chunks)
        cost = time.time() - t1
        print(f"Embedding extraction: {cost:.2f}s for {len(chunks)} chunks "
              f"({cost / max(len(chunks), 1) * 1000:.2f} ms/chunk)")
        print(f"Embeddings shape: {embeddings.shape}")

        if args.diarize:
            try:
                from campplus_sdk.clustering import do_clustering
            except ImportError as e:
                raise SystemExit(
                    f"clustering deps missing ({e}); run: "
                    "pip install scipy scikit-learn fastcluster umap-learn hdbscan")
            t1 = time.time()
            speaker_num, diar_results = do_clustering(
                chunks, embeddings,
                speaker_num=args.speaker_num if args.speaker_num else None)
            print(f"Clustering: {time.time() - t1:.2f}s, "
                  f"speakers: {speaker_num}")
            for seg_st, seg_ed, spk in diar_results:
                print(f"  Speaker_{spk}: [{seg_st:.2f} {seg_ed:.2f}]")
        return

    # 1:1 speaker verification
    wav1 = load_wav_16k(args.wav1)
    wav2 = load_wav_16k(args.wav2)
    emb1 = model.extract(wav1)
    emb2 = model.extract(wav2)
    sim = cosine_similarity(emb1, emb2)
    print(f"\n{args.wav1} vs {args.wav2}")
    print(f"cosine similarity: {sim:.4f}")
    print("note: the higher the similarity, the more likely the two audios "
          "belong to the same speaker")


if __name__ == "__main__":
    main()