"""
Inference script for HowFar-Caarma: distance estimation from speech using HuBERT.

Usage:
    python inference.py --ckpt epoch18_val_acc7997.ckpt --audio path/to/audio.wav
"""
import argparse
import torch
import torch.nn as nn
import torchaudio
import pytorch_lightning as pl
from transformers import HubertModel, AutoFeatureExtractor


class Hubert_Model(nn.Module):
    def __init__(self, hubert_model_name="facebook/hubert-large-ls960-ft", cache_dir=""):
        super().__init__()
        self.encoder = HubertModel.from_pretrained(hubert_model_name, cache_dir=cache_dir)
        hidden_size = self.encoder.config.hidden_size
        self.layer_norm = nn.LayerNorm(hidden_size)
        self.dropout = nn.Dropout(0.1)
        self.bn = nn.BatchNorm1d(hidden_size)

    def forward(self, input_values, labels=None):
        outputs = self.encoder(input_values)
        hidden_states = outputs.last_hidden_state
        hidden_states = self.layer_norm(hidden_states)
        hidden_states = self.dropout(hidden_states)
        pooled = torch.mean(hidden_states, dim=1)
        pooled = self.bn(pooled)
        return pooled


class CassavaPLModule(pl.LightningModule):
    def __init__(self, hparams, model):
        super().__init__()
        self.model = model

    def forward(self, x):
        return self.model(x)


def load_model(ckpt_path, device="cuda"):
    base = Hubert_Model()
    model = CassavaPLModule.load_from_checkpoint(
        ckpt_path,
        hparams={"lr": 0.001, "batch_size": 1},
        model=base,
        strict=False,
        map_location=device,
    )
    model.eval()
    model.to(device)
    model.freeze()
    return model


def extract_embedding(model, audio_path, device="cuda"):
    waveform, sr = torchaudio.load(audio_path)
    if sr != 16000:
        waveform = torchaudio.functional.resample(waveform, sr, 16000)
    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)

    processor = AutoFeatureExtractor.from_pretrained("facebook/hubert-large-ls960-ft")
    inputs = processor(
        waveform.squeeze(0).numpy(),
        sampling_rate=16000,
        return_tensors="pt",
    ).input_values.to(device)

    with torch.no_grad():
        embedding = model(inputs)
    return embedding


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt", required=True)
    parser.add_argument("--audio", required=True)
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = load_model(args.ckpt, device=device)
    emb = extract_embedding(model, args.audio, device=device)
    print(f"Embedding shape: {emb.shape}")