"""Minimal inference for the converted RVC HuBERT Transformers model.

Dependencies:
  pip install torch transformers

Input waveform must already be mono 16 kHz float audio in [-1, 1].
This script intentionally avoids audio loading dependencies such as
soundfile, torchaudio, or librosa.
"""

from __future__ import annotations

import argparse

import torch
from transformers import HubertModel


@torch.inference_mode()
def extract_rvc_hubert_features(
    waveform_16k: torch.Tensor,
    model: HubertModel,
    *,
    output_layer: int = 12,
    repeat_factor: int = 2,
) -> torch.Tensor:
    """Return RVC-compatible HuBERT features.

    Args:
        waveform_16k: Tensor shaped [T] or [B, T], mono 16 kHz float audio.
        model: Converted Transformers HubertModel.
        output_layer: 12 matches RVC v2's HuBERT content features.
        repeat_factor: 2 converts HuBERT's ~50 Hz features to RVC's ~100 Hz conditioning rate.

    Returns:
        Tensor shaped [B, frames, 768].
    """

    if waveform_16k.ndim == 1:
        waveform_16k = waveform_16k.unsqueeze(0)
    if waveform_16k.ndim != 2:
        raise ValueError(f"Expected waveform shape [T] or [B, T], got {tuple(waveform_16k.shape)}")

    device = next(model.parameters()).device
    output = model(
        input_values=waveform_16k.to(device=device, dtype=torch.float32),
        output_hidden_states=True,
    )
    if output_layer >= 0:
        features = output.hidden_states[output_layer]
    else:
        features = output.last_hidden_state

    features = features.float()
    if repeat_factor > 1:
        features = features.repeat_interleave(int(repeat_factor), dim=1)
    return features


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="assets/hubert/hubert_base_transformers")
    parser.add_argument("--seconds", type=float, default=2.0)
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = HubertModel.from_pretrained(args.model).to(device).eval()

    # Replace this dummy input with your own mono 16 kHz float waveform tensor.
    waveform = torch.zeros(int(16000 * float(args.seconds)), dtype=torch.float32)
    features = extract_rvc_hubert_features(waveform, model)
    print(features.shape)


if __name__ == "__main__":
    main()