"""Minimal usage example for hv-contour.

Synthesizes three whistles with prescribed contours, encodes them, and
prints pairwise similarities.
"""
import math
import os
import wave
from array import array

from hv_contour import HVContour


OUTDIR = "./hv_contour_example"


def _write_wav(path: str, samples, sr: int = 16000) -> None:
    ints = array(
        "h",
        (max(-32768, min(32767, int(x * 32767))) for x in samples),
    )
    with wave.open(path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        w.writeframes(ints.tobytes())


def _whistle(contour: str, f0: float, f1: float,
             dur_ms: int = 600, sr: int = 16000):
    n = int(sr * dur_ms / 1000)
    dur_s = dur_ms / 1000.0
    out = []
    for i in range(n):
        t = i / n
        if contour == "R":
            f = f0 + (f1 - f0) * t
        elif contour == "F":
            f = f0 + (f1 - f0) * t
        elif contour == "W":
            peak = max(f0, f1) * 1.2
            f = (f0 + (peak - f0) * (t / 0.5) if t < 0.5
                 else peak - (peak - f1) * ((t - 0.5) / 0.5))
        elif contour == "U":
            trough = min(f0, f1) * 0.8
            f = (f0 - (f0 - trough) * (t / 0.5) if t < 0.5
                 else trough + (f1 - trough) * ((t - 0.5) / 0.5))
        else:
            f = f0
        phase = 2 * math.pi * f * (i / sr)
        env = math.sin(math.pi * i / n) ** 0.5
        out.append(0.5 * env * math.sin(phase))
    return out


def main() -> None:
    os.makedirs(OUTDIR, exist_ok=True)

    # Synthesize three contrasting contours
    samples = [
        ("up", _whistle("R", 800, 1600)),
        ("down", _whistle("F", 1600, 800)),
        ("flat", _whistle("S", 1200, 1200)),
    ]
    paths = []
    for name, sig in samples:
        p = os.path.join(OUTDIR, f"{name}.wav")
        _write_wav(p, sig)
        paths.append(p)

    model = HVContour(dim=50000, seed=42)

    print("=" * 60)
    print("hv-contour")
    print("=" * 60)
    print()

    encoded = []
    for path in paths:
        contours, hv = model.encode_file(path)
        tokens = [c.token for c in contours]
        print(f"  {os.path.basename(path):<12}  "
              f"tokens = {' '.join(tokens):<10}  "
              f"dim = {hv.dim}  ones = {sum(hv.bits)}")
        encoded.append((os.path.basename(path), hv))
    print()

    print("=" * 60)
    print("Pairwise similarity")
    print("=" * 60)
    print(f"  {'':<12}" + "".join(f"{n:>10}" for n, _ in encoded))
    for name_a, hv_a in encoded:
        row = f"  {name_a:<12}"
        for name_b, hv_b in encoded:
            row += f"{hv_a.similarity(hv_b):>10.3f}"
        print(row)
    print()

    print("=" * 60)
    print("Cross-check the audio round-trip")
    print("=" * 60)
    for name, _ in encoded:
        print(f"  synth {name:<6}  ->  extracted "
              f"{' '.join(c.token for c in model.encode_file(os.path.join(OUTDIR, f'{name}.wav'))[0])}")
    print()


if __name__ == "__main__":
    main()