#!/usr/bin/env python
"""inference.py — Pi-Indic-Speaker-Embed-v1: Extract 512-dim speaker embeddings.

A clean, self-contained inference script for the Pi-Indic-Speaker-Embed-v1 model.
Downloads WavLM Large from HuggingFace at runtime; only the lightweight
trainable weights (LoRA + ECAPA-TDNN + pooling + projector, ~96 MB) are
stored in the checkpoint.

Usage as a library:
    from inference import PiIndicSpeakerEmbed

    model = PiIndicSpeakerEmbed.from_pretrained("Paytmlabs/Pi-indic-speaker-embed-v1")
    embedding = model.extract("audio.wav")          # → (512,) numpy array
    embeddings = model.extract_batch(["a.wav", "b.wav"])  # → (N, 512) numpy array

Usage from command line:
    python inference.py --audio audio1.wav audio2.wav --output embeddings.npy

Architecture:
    WavLM Large (frozen base + LoRA r=8/16 on layers 9-24)
        → multilayer weighted sum (1024-dim frame features)
        → utterance mean normalization
        → ECAPA-TDNN (3 Res2Net blocks, 1536-dim frame embeddings)
        → Channel-Attention Statistics Pooling (3072-dim utterance embedding)
        → RawNet3 Projector (BatchNorm + Linear → 512-dim speaker embedding)
        → L2 normalize

Requirements:
    pip install espnet s3prl transformers torch soundfile librosa numpy
"""
import os
import sys
import argparse
import numpy as np
import torch
import torch.nn.functional as F
import soundfile as sf
from pathlib import Path
from typing import Union, List, Optional


class PiIndicSpeakerEmbed:
    """Pi-Indic-Speaker-Embed-v1: 512-dim speaker embedding extractor for Indic languages.

    The model is warm-started from espnet/voxcelebs12_ecapa_wavlm_joint and
    fine-tuned on Kathbath (12 Indic languages) + VoxCeleb with LoRA on WavLM
    layers 9-24 and full ECAPA-TDNN fine-tuning.

    Achieves 1.586% EER on Kathbath validation (unseen speakers) and 4.003% EER
    on IndicVoices-R test set (Hindi, Marathi, Bengali, Telugu).
    """

    def __init__(self, model_path: str, config_path: str, device: str = "auto",
                 target_duration: float = 6.0, num_eval: int = 3, batch_size: int = 10):
        """Initialize the model.

        Args:
            model_path: Path to the stripped checkpoint (.pth file)
            config_path: Path to the config YAML file
            device: "cuda", "cpu", or "auto" (uses CUDA if available)
            target_duration: Audio chunk length in seconds (default 6.0)
            num_eval: Number of random crops per utterance for score averaging (default 3)
            batch_size: Batch size for forward pass (default 10)
        """
        self.target_duration = target_duration
        self.num_eval = num_eval
        self.batch_size = batch_size

        if device == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        self._build_model(config_path)
        self._load_checkpoint(model_path)
        self.model.eval()

    def _build_model(self, config_path: str):
        """Build the ESPnet speaker model from config, with LoRA injection."""
        # Register custom frontend + preprocessor
        from espnet2.tasks.spk import SpeakerTask, frontend_choices, preprocessor_choices
        _register_custom_components()

        import yaml
        with open(config_path) as f:
            config = yaml.safe_load(f)

        self.spk_num = config.get("_spk_num", 3790)  # doesn't affect embedding extraction

        model_args = SpeakerTask.get_parser().parse_args([
            "--config", config_path,
            "--output_dir", "/tmp/pi_indic_dummy",
            "--ngpu", "1" if "cuda" in self.device else "0",
            "--train_data_path_and_name_and_type", "data/train/wav.scp,speech,sound",
            "--train_data_path_and_name_and_type", "data/train/utt2spk,spk_labels,text",
            "--valid_data_path_and_name_and_type", "data/eval/trial.scp,speech,sound",
            "--valid_data_path_and_name_and_type", "data/eval/trial2.scp,speech2,sound",
            "--valid_data_path_and_name_and_type", "data/eval/trial_label,spk_labels,text",
            "--spk2utt", "data/train/spk2utt",
            "--spk_num", str(self.spk_num),
            "--train_shape_file", "exp/spk_train_ecapa_wavlm_lora_512/spk_stats/train/speech_shape",
            "--valid_shape_file", "exp/spk_train_ecapa_wavlm_lora_512/spk_stats/valid/speech_shape",
        ])

        print("Building model (WavLM Large + LoRA + ECAPA-TDNN)...")
        self.model = SpeakerTask.build_model(model_args).to(self.device)

        # Build preprocessor
        from spk_preprocessor_lora import SpkLoRAPreprocessor
        pp_conf = config.get("preprocessor_conf", {})
        self.preprocessor = SpkLoRAPreprocessor(
            train=False,
            target_duration=pp_conf.get("target_duration", self.target_duration),
            sample_rate=pp_conf.get("sample_rate", 16000),
            num_eval=pp_conf.get("num_eval", self.num_eval),
            noise_apply_prob=0.0, rir_apply_prob=0.0, noise_info=[], rir_scp=None,
        )

    def _load_checkpoint(self, model_path: str):
        """Load the stripped checkpoint (LoRA + ECAPA + pooling + projector only).

        WavLM base weights are NOT in the checkpoint — they're downloaded from
        HuggingFace (microsoft/wavlm-large) during model building. Only the
        LoRA adapters and downstream components are loaded here.
        """
        print(f"Loading checkpoint: {os.path.basename(model_path)}")
        state_dict = torch.load(model_path, map_location="cpu", weights_only=False)
        if isinstance(state_dict, dict) and "model" in state_dict:
            state_dict = state_dict["model"]

        missing, unexpected = self.model.load_state_dict(state_dict, strict=False)
        n_missing = len(missing)
        n_unexpected = len(unexpected)
        if n_missing > 0:
            # Missing keys are expected: WavLM base weights (loaded from HF) + loss weights
            non_wavlm_missing = [k for k in missing if not k.startswith("frontend.")]
            if non_wavlm_missing:
                print(f"  Warning: {len(non_wavlm_missing)} unexpected missing keys: {non_wavlm_missing[:3]}")
            else:
                print(f"  {n_missing} missing keys (WavLM base + loss — loaded from HF / not needed)")
        if n_unexpected > 0:
            print(f"  Warning: {n_unexpected} unexpected keys: {unexpected[:3]}")
        print("  Checkpoint loaded successfully.")

    @torch.no_grad()
    def extract(self, audio_path: str) -> np.ndarray:
        """Extract a 512-dim L2-normalized speaker embedding from an audio file.

        Args:
            audio_path: Path to a WAV file (any sample rate, auto-resampled to 16kHz).

        Returns:
            (512,) numpy float32 array, L2-normalized (unit length on hypersphere).
        """
        audio = self._load_audio(audio_path)
        speech = self._preprocess(audio)
        speech = speech.unsqueeze(0).to(self.device)  # (1, num_eval, T)

        # Flatten num_eval into batch for forward pass
        org_shape = speech.shape[:2]
        speech = speech.flatten(0, 1)  # (num_eval, T)

        embds = self.model(speech=speech, spk_labels=None,
                           extract_embd=True, task_tokens=None)
        embds = F.normalize(embds, p=2, dim=1)  # L2 normalize

        # Average across num_eval crops
        embds = embds.view(org_shape[0], org_shape[1], -1)  # (1, num_eval, 512)
        embd = embds.mean(dim=1)  # (1, 512)
        embd = F.normalize(embd, p=2, dim=1)  # re-normalize after averaging

        return embd.squeeze(0).cpu().numpy()

    @torch.no_grad()
    def extract_batch(self, audio_paths: List[str]) -> np.ndarray:
        """Extract 512-dim L2-normalized embeddings for multiple audio files.

        Args:
            audio_paths: List of paths to WAV files.

        Returns:
            (N, 512) numpy float32 array, L2-normalized.
        """
        embeddings = []
        for i in range(0, len(audio_paths), self.batch_size):
            batch_paths = audio_paths[i:i + self.batch_size]
            batch_speech = []
            for path in batch_paths:
                audio = self._load_audio(path)
                speech = self._preprocess(audio)
                batch_speech.append(speech)
            speech_tensor = torch.stack(batch_speech, dim=0).to(self.device)
            org_shape = speech_tensor.shape[:2]
            speech_tensor = speech_tensor.flatten(0, 1)  # (B*num_eval, T)

            embds = self.model(speech=speech_tensor, spk_labels=None,
                               extract_embd=True, task_tokens=None)
            embds = F.normalize(embds, p=2, dim=1)
            embds = embds.view(org_shape[0], org_shape[1], -1)
            embds = embds.mean(dim=1)
            embds = F.normalize(embds, p=2, dim=1)

            embeddings.append(embds.cpu().numpy())

        return np.concatenate(embeddings, axis=0)

    def _load_audio(self, path: str, sample_rate: int = 16000) -> np.ndarray:
        """Load audio file as float32 numpy array at target sample rate."""
        audio, sr = sf.read(path, dtype="float32")
        if len(audio.shape) > 1:
            audio = audio[:, 0]
        if sr != sample_rate:
            import librosa
            audio = librosa.resample(audio, orig_sr=sr, target_sr=sample_rate)
        return audio

    def _preprocess(self, audio: np.ndarray) -> torch.Tensor:
        """Preprocess audio: pad/crop to target_duration, num_eval random crops."""
        data = {"speech": audio, "speech2": audio}
        data = self.preprocessor._speech_process(data)
        speech = data["speech"]
        if not isinstance(speech, torch.Tensor):
            speech = torch.from_numpy(np.asarray(speech))
        if speech.dim() == 1:
            speech = speech.unsqueeze(0)
        return speech  # (num_eval, T)

    @classmethod
    def from_pretrained(cls, repo_id: str = "Paytmlabs/Pi-indic-speaker-embed-v1",
                        device: str = "auto", **kwargs) -> "PiIndicSpeakerEmbed":
        """Download model files from HuggingFace and initialize.

        Args:
            repo_id: HuggingFace repo ID (default: "Paytmlabs/Pi-indic-speaker-embed-v1")
            device: "cuda", "cpu", or "auto"

        Returns:
            PiIndicSpeakerEmbed instance ready for inference.
        """
        from huggingface_hub import hf_hub_download

        # Download files from HuggingFace
        model_path = hf_hub_download(repo_id=repo_id, filename="valid.eer.ave_5best_stripped.pth")
        config_path = hf_hub_download(repo_id=repo_id, filename="config.yaml")

        # Also need the custom code files
        code_dir = os.path.dirname(model_path)
        for fname in ["lora_s3prl_frontend.py", "lora_wavlm.py", "spk_preprocessor_lora.py"]:
            hf_hub_download(repo_id=repo_id, filename=fname)

        # Add code dir to path
        if code_dir not in sys.path:
            sys.path.insert(0, code_dir)

        return cls(model_path=model_path, config_path=config_path, device=device, **kwargs)


def _register_custom_components():
    """Register custom LoRA frontend and preprocessor into ESPnet."""
    from espnet2.tasks.spk import frontend_choices, preprocessor_choices
    from lora_s3prl_frontend import LoRAS3prlFrontend
    from spk_preprocessor_lora import SpkLoRAPreprocessor
    frontend_choices.classes["lora_s3prl"] = LoRAS3prlFrontend
    preprocessor_choices.classes["spk_lora"] = SpkLoRAPreprocessor


def main():
    parser = argparse.ArgumentParser(description="Extract speaker embeddings with Pi-Indic-Speaker-Embed-v1")
    parser.add_argument("--audio", nargs="+", required=True, help="Path(s) to audio file(s)")
    parser.add_argument("--output", default=None, help="Output .npy file (default: print to stdout)")
    parser.add_argument("--model_path", default=None, help="Path to checkpoint (if not downloading from HF)")
    parser.add_argument("--config_path", default=None, help="Path to config YAML (if not downloading from HF)")
    parser.add_argument("--repo_id", default="Paytmlabs/Pi-indic-speaker-embed-v1", help="HF repo ID")
    parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
    args = parser.parse_args()

    if args.model_path and args.config_path:
        model = PiIndicSpeakerEmbed(args.model_path, args.config_path, device=args.device)
    else:
        model = PiIndicSpeakerEmbed.from_pretrained(args.repo_id, device=args.device)

    if len(args.audio) == 1:
        emb = model.extract(args.audio[0])
        embeddings = emb.reshape(1, -1)
    else:
        embeddings = model.extract_batch(args.audio)

    print(f"Embedding shape: {embeddings.shape}")
    print(f"L2 norm: {np.linalg.norm(embeddings, axis=1)}")

    if args.output:
        np.save(args.output, embeddings)
        print(f"Saved to {args.output}")
    else:
        np.set_printoptions(precision=4, suppress=True)
        for i, (path, emb) in enumerate(zip(args.audio, embeddings)):
            print(f"\n{path}:")
            print(f"  {emb}")


if __name__ == "__main__":
    main()