"""Generate synthetic appliance-sound samples so judges can test without hardware.

Run from the project root:  python assets/generate_samples.py
Produces mono 16 kHz WAVs in assets/.
"""
import os

import numpy as np
import soundfile as sf

SR = 16000
DUR = 8.0
OUT = os.path.dirname(__file__)


def _norm(y: np.ndarray) -> np.ndarray:
    y = y - np.mean(y)
    peak = np.max(np.abs(y)) + 1e-9
    return (0.9 * y / peak).astype(np.float32)


def washer_bearing() -> np.ndarray:
    t = np.linspace(0, DUR, int(SR * DUR), endpoint=False)
    base = 0.04 * np.random.randn(len(t))           # low broadband noise floor
    base += 0.08 * np.sin(2 * np.pi * 55 * t)       # drum hum
    clicks = np.zeros_like(t)                         # sharp rhythmic 4 Hz impacts
    for k in np.arange(0, DUR, 0.25):
        i = int(k * SR)
        env = np.exp(-np.linspace(0, 1, 300) * 40)
        clicks[i:i+300] += 1.0 * env * np.sin(2 * np.pi * 2000 * np.linspace(0, .0188, 300))
    return _norm(base + clicks)


def fan_imbalanced() -> np.ndarray:
    t = np.linspace(0, DUR, int(SR * DUR), endpoint=False)
    hum = 0.4 * np.sin(2 * np.pi * 50 * t)
    wobble = 1 + 0.5 * np.sin(2 * np.pi * 3.3 * t)  # amplitude imbalance
    return _norm(hum * wobble + 0.05 * np.random.randn(len(t)))


def motor_squeal() -> np.ndarray:
    t = np.linspace(0, DUR, int(SR * DUR), endpoint=False)
    am = 1 + 0.3 * np.sin(2 * np.pi * 6 * t)
    tone = 0.5 * np.sin(2 * np.pi * 2500 * t) * am
    return _norm(tone + 0.04 * np.random.randn(len(t)))


def washer_good() -> np.ndarray:
    t = np.linspace(0, DUR, int(SR * DUR), endpoint=False)
    hum = 0.3 * np.sin(2 * np.pi * 50 * t) + 0.1 * np.sin(2 * np.pi * 100 * t)
    return _norm(hum + 0.02 * np.random.randn(len(t)))


SAMPLES = {
    "sample_washer_bearing.wav": washer_bearing,
    "sample_fan_imbalanced.wav": fan_imbalanced,
    "sample_motor_squeal.wav": motor_squeal,
    "sample_washer_good.wav": washer_good,
}

if __name__ == "__main__":
    np.random.seed(0)
    for name, fn in SAMPLES.items():
        sf.write(os.path.join(OUT, name), fn(), SR)
        print("wrote", name)