"""Example: Run DiariZen segmentation on a WAV file."""

import argparse
import sys
import numpy as np

from diarizen_sdk import DiarizenSegmenter
from diarizen_sdk.postprocess import log_probs_to_probs, top_speakers_at_frame


def main():
    parser = argparse.ArgumentParser(description="DiariZen speaker segmentation")
    parser.add_argument("audio", help="Path to 16kHz mono WAV file")
    parser.add_argument("--cnn-model", default="cnn_features.axmodel",
                        help="Path to CNN NPU model")
    parser.add_argument("--backend-model", default="backend.onnx",
                        help="Path to backend ONNX model")
    args = parser.parse_args()

    # Load audio
    try:
        import soundfile as sf
        audio, sr = sf.read(args.audio, dtype="float32")
    except ImportError:
        print("soundfile not available, using scipy.io.wavfile")
        from scipy.io import wavfile
        sr, audio = wavfile.read(args.audio)
        audio = audio.astype(np.float32) / 32768.0

    if audio.ndim > 1:
        audio = audio[:, 0]  # Use first channel

    print(f"Audio: {len(audio)} samples @ {sr} Hz")

    # Create segmenter and run
    segmenter = DiarizenSegmenter(args.cnn_model, args.backend_model)
    log_probs = segmenter(audio, sr)

    print(f"Output shape: {log_probs.shape}")
    print(f"  Frames: {log_probs.shape[1]}, Classes: {log_probs.shape[2]}")

    # Show results for key frames
    probs = log_probs_to_probs(log_probs)
    check_frames = [0, 50, 100, 150, 198]
    print("\nTop-3 speaker classes per selected frame:")
    for f in check_frames:
        top = top_speakers_at_frame(log_probs, f, top_k=3)
        items = ", ".join(f"cls {c}: {lp:.2f}" for c, lp in top)
        print(f"  Frame {f:3d}: {items}")

    # Overall most active class
    mean_probs = probs[0].mean(axis=0)
    top_class = int(np.argmax(mean_probs))
    print(f"\nMost active class overall: {top_class} "
          f"(avg prob={mean_probs[top_class]:.4f})")


if __name__ == "__main__":
    main()