#!/usr/bin/env python3
"""
Train an audio embedding classifier on recorded samples.

Uses wav2vec2 to extract embeddings, then trains a simple classifier.

Usage:
    uv run python train_classifier.py
"""

import pickle
from pathlib import Path

import numpy as np
import torch
from scipy.io import wavfile
from torchaudio.pipelines import WAV2VEC2_BASE

RECORDINGS_DIR = Path(__file__).parent / "recordings"
MODEL_PATH = Path(__file__).parent / "classifier.pkl"
SAMPLE_RATE = 16000

# Categories (excluding wake_word which is handled separately)
CATEGORIES = [
    "pressure_below_15",
    "pressure_15_to_25",
    "pressure_25_to_35",
    "pressure_above_35",
    "drop_l",
    "drop_p",
    "drop_d",
    "both_eyes",
]


def load_audio(path: Path) -> torch.Tensor:
    """Load and preprocess audio file."""
    sr, data = wavfile.read(path)

    # Convert to float32
    if data.dtype == np.int16:
        data = data.astype(np.float32) / 32768.0
    elif data.dtype == np.int32:
        data = data.astype(np.float32) / 2147483648.0
    elif data.dtype != np.float32:
        data = data.astype(np.float32)

    # Convert to mono if stereo
    if len(data.shape) > 1:
        data = data.mean(axis=1)

    # Resample if needed
    if sr != SAMPLE_RATE:
        from scipy import signal
        num_samples = int(len(data) * SAMPLE_RATE / sr)
        data = signal.resample(data, num_samples)

    # Convert to torch tensor with batch dimension
    waveform = torch.from_numpy(data).unsqueeze(0)
    return waveform


def extract_embedding(waveform: torch.Tensor, model, device: str) -> np.ndarray:
    """Extract embedding from audio using wav2vec2."""
    with torch.no_grad():
        waveform = waveform.to(device)
        features, _ = model.extract_features(waveform)
        # Use the last layer's features, averaged over time
        embedding = features[-1].mean(dim=1).cpu().numpy()
    return embedding.flatten()


def load_samples():
    """Load all training samples."""
    samples = []
    labels = []

    for category in CATEGORIES:
        category_dir = RECORDINGS_DIR / category
        if not category_dir.exists():
            print(f"Warning: {category_dir} does not exist")
            continue

        wav_files = list(category_dir.glob("*.wav"))
        print(f"  {category}: {len(wav_files)} samples")

        for wav_file in wav_files:
            samples.append(wav_file)
            labels.append(category)

    return samples, labels


def main():
    print("=== Training Audio Classifier ===\n")

    # Check for samples
    if not RECORDINGS_DIR.exists():
        print(f"Error: {RECORDINGS_DIR} does not exist")
        print("Run ./collect_samples.sh first to record training samples")
        return

    print("Loading samples...")
    sample_paths, labels = load_samples()

    if len(sample_paths) == 0:
        print("\nNo samples found. Run ./collect_samples.sh first.")
        return

    if len(set(labels)) < 2:
        print(f"\nNeed samples from at least 2 categories. Found: {set(labels)}")
        return

    print(f"\nTotal: {len(sample_paths)} samples across {len(set(labels))} categories")

    # Check minimum samples per category
    from collections import Counter
    counts = Counter(labels)
    min_samples = min(counts.values())
    if min_samples < 3:
        print(f"\nWarning: Some categories have < 3 samples. More samples = better accuracy.")

    # Load model
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"\nLoading wav2vec2 on {device}...")
    bundle = WAV2VEC2_BASE
    model = bundle.get_model().to(device)
    model.eval()

    # Extract embeddings
    print("Extracting embeddings...")
    embeddings = []
    valid_labels = []

    for i, (path, label) in enumerate(zip(sample_paths, labels)):
        try:
            waveform = load_audio(path)
            embedding = extract_embedding(waveform, model, device)
            embeddings.append(embedding)
            valid_labels.append(label)
            print(f"  [{i+1}/{len(sample_paths)}] {path.name}")
        except Exception as e:
            print(f"  [{i+1}/{len(sample_paths)}] {path.name} - ERROR: {e}")

    if len(embeddings) < 2:
        print("\nNot enough valid samples to train.")
        return

    X = np.array(embeddings)
    y = np.array(valid_labels)

    print(f"\nEmbedding shape: {X.shape}")

    # Train classifier
    print("\nTraining classifier...")
    from sklearn.preprocessing import LabelEncoder
    from sklearn.svm import SVC
    from sklearn.model_selection import cross_val_score

    label_encoder = LabelEncoder()
    y_encoded = label_encoder.fit_transform(y)

    # Use SVM with probability estimates
    classifier = SVC(kernel='rbf', probability=True, C=1.0)

    # Cross-validation if enough samples
    if len(X) >= 4:
        n_splits = min(5, len(X))
        scores = cross_val_score(classifier, X, y_encoded, cv=n_splits)
        print(f"Cross-validation accuracy: {scores.mean():.2%} (+/- {scores.std()*2:.2%})")

    # Train on all data
    classifier.fit(X, y_encoded)

    # Save model
    model_data = {
        'classifier': classifier,
        'label_encoder': label_encoder,
        'categories': CATEGORIES,
    }

    with open(MODEL_PATH, 'wb') as f:
        pickle.dump(model_data, f)

    print(f"\nClassifier saved to: {MODEL_PATH}")
    print(f"Categories: {list(label_encoder.classes_)}")
    print("\nDone! The listener will now use this classifier.")


if __name__ == "__main__":
    main()