#!/usr/bin/env python3
"""
Train on whispers, classify a single voiced word.

Run: python examples/single_word.py
"""
import numpy as np
from whisper_decoder import (
    MLP, train_mlp, extract_features, synth_word,
    generate_corpus, FEATURE_DIM, WORD_ORDER,
)


def main():
    print("generating whisper training set...")
    train = generate_corpus('whisper', 150, 20, seed=0,
                            reverb_train_frac=0.25)

    mu = train.X_train.mean(axis=0)
    sigma = train.X_train.std(axis=0) + 1e-9
    X_tr = (train.X_train - mu) / sigma

    model = MLP(FEATURE_DIM, 64, 32, 10, seed=0)
    print("training...")
    train_mlp(model, X_tr, train.y_train, epochs=200,
              batch=64, lr=3e-3, seed=0)

    word_index_inv = {i: w for w, i in train.word_index.items()}

    print()
    print("classify three voiced words with the whisper-trained model:")
    for word in ['zero', 'three', 'seven']:
        sig = synth_word(word, 'voiced', seed=hash(word) & 0xFFFF)
        feats = extract_features(sig)
        probs = model.predict_proba(((feats - mu) / sigma)[None, :])[0]
        top3 = np.argsort(-probs)[:3]
        pred = word_index_inv[int(top3[0])]
        print(f"  true: {word:<6s}  "
              f"predicted: {pred:<6s}  "
              f"p = {probs[top3[0]]:.3f}")
        print(f"    top-3: "
              + ", ".join(f"{word_index_inv[int(i)]}({probs[i]:.2f})"
                          for i in top3))


if __name__ == "__main__":
    main()