"""Heli-WAV spektrogramm PNG-ks (ilma matplotlibita). Kasutus: python specgram.py heli.wav [välja.png]
Kõne paistab horisontaalsete formant-triipude ja silbirütmina; müra on ühtlane udu."""
import struct, sys, wave, zlib
import numpy as np


def write_png(path, img):
    """img: uint8 (h, w) hallskaala -> PNG."""
    h, w = img.shape
    raw = b"".join(b"\x00" + img[y].tobytes() for y in range(h))
    def chunk(tag, data):
        c = tag + data
        return struct.pack(">I", len(data)) + c + struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
    png = b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 0, 0, 0, 0))
    png += chunk(b"IDAT", zlib.compress(raw, 6)) + chunk(b"IEND", b"")
    open(path, "wb").write(png)


def specgram(path, out=None, nfft=512, hop=160, fmax=4000, px_per_sec=100):
    w = wave.open(path); fs = w.getframerate()
    x = np.frombuffer(w.readframes(w.getnframes()), np.int16).astype(np.float64) / 32768
    n = (len(x) - nfft) // hop
    fr = np.stack([x[i * hop:i * hop + nfft] * np.hanning(nfft) for i in range(n)])
    S = np.abs(np.fft.rfft(fr, axis=1)) ** 2
    f = np.fft.rfftfreq(nfft, 1 / fs)
    S = S[:, f <= fmax]
    db = 10 * np.log10(S + 1e-12)
    band = db[:, (f[f <= fmax] >= 300) & (f[f <= fmax] <= 3000)]      # skaala kõneribast, mitte filtri nullidest
    lo, hi = np.percentile(band, 10), np.percentile(band, 99.5)
    img = np.clip((db - lo) / (hi - lo), 0, 1)
    img = (255 * (1 - img)).astype(np.uint8).T[::-1]          # tume = tugev, sagedus üles
    # laius ajas: px_per_sec
    dur = len(x) / fs
    tw = max(int(dur * px_per_sec), 50)
    idx = np.linspace(0, img.shape[1] - 1, tw).astype(int)
    img = img[:, idx]
    out = out or path.rsplit(".", 1)[0] + "_spec.png"
    write_png(out, np.ascontiguousarray(img))
    print(f"{out}  ({img.shape[1]}x{img.shape[0]}, {dur:.1f} s, 0-{fmax} Hz)")
    return out


if __name__ == "__main__":
    specgram(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else None)