#!/usr/bin/env python3
"""Synthesize a small pop-style song with a KNOWN key, tempo, chords and melody,
so the pianoify pipeline can be scored objectively.

Writes <out>.mp3 (or .wav) plus <out>.truth.json.
"""
import json
import subprocess
import sys

import numpy as np
import soundfile as sf

SR = 44100
BPM = 96.0
BEAT = 60.0 / BPM

# 8-bar phrase in G major, played twice. One bar is split C|D (half-bar change).
PHRASE_CHORDS = [["G"], ["D"], ["Em"], ["C"], ["G"], ["D"], ["C", "D"], ["G"]]
CHORD_TONES = {  # root midi (octave 3) + quality
    "G": (55, [0, 4, 7]), "D": (50, [0, 4, 7]), "Em": (52, [0, 3, 7]),
    "C": (48, [0, 4, 7]), "Am": (57, [0, 3, 7]),
}
PHRASE_MELODY = [  # (midi, beats) per bar, 4 beats each
    [(67, 1), (71, 1), (74, 1.5), (72, 0.5)],
    [(71, 1), (69, 1), (66, 2)],
    [(67, 0.5), (69, 0.5), (71, 1), (76, 2)],
    [(74, 1), (72, 1), (71, 2)],
    [(67, 1), (67, 1), (71, 1), (74, 1)],      # repeated note: tests re-articulation
    [(78, 2), (76, 1), (74, 1)],
    [(72, 1), (76, 1), (74, 1), (69, 1)],
    [(67, 4)],
]
INTRO_BARS = 1   # drums only
REPEATS = 2


def env(n, attack=0.01, decay=1.5):
    t = np.arange(n) / SR
    a = np.clip(t / attack, 0, 1)
    return a * np.exp(-t / decay)


def tone(f, dur, harmonics=(1, .5, .3, .15), decay=1.5, vib=0.0):
    n = int(dur * SR)
    t = np.arange(n) / SR
    ph = 2 * np.pi * f * t
    if vib:
        ph += vib * np.sin(2 * np.pi * 5.5 * t) / 5.5
    sig = sum(a * np.sin(k * ph) for k, a in enumerate(harmonics, 1))
    rel = np.clip((n - np.arange(n)) / (0.03 * SR), 0, 1)
    return sig * env(n, decay=decay) * rel


def m2f(m):
    return 440.0 * 2 ** ((m - 69) / 12)


def add(buf, sig, t0, gain):
    i = int(t0 * SR)
    j = min(len(buf), i + len(sig))
    buf[i:j] += gain * sig[: j - i]


def main(out):
    rng = np.random.default_rng(0)
    total_bars = INTRO_BARS + REPEATS * len(PHRASE_CHORDS) + 1
    buf = np.zeros(int((total_bars * 4 + 2) * BEAT * SR))

    truth_chords, truth_notes = [], []
    # drums on every bar
    for b in range(total_bars):
        for beat in range(4):
            t = (b * 4 + beat) * BEAT
            if beat in (0, 2):  # kick
                n = int(0.25 * SR)
                tt = np.arange(n) / SR
                f = 50 + 70 * np.exp(-tt * 30)
                add(buf, np.sin(2 * np.pi * np.cumsum(f) / SR) * np.exp(-tt * 14), t, 0.9)
            else:  # snare
                n = int(0.18 * SR)
                add(buf, rng.standard_normal(n) * np.exp(-np.arange(n) / SR * 25), t, 0.25)
            for h in (0, 0.5):  # hats
                n = int(0.05 * SR)
                hat = np.diff(rng.standard_normal(n + 1)) * np.exp(-np.arange(n) / SR * 80)
                add(buf, hat, t + h * BEAT, 0.08)

    bar = INTRO_BARS
    for _ in range(REPEATS):
        for chords, mel in zip(PHRASE_CHORDS, PHRASE_MELODY):
            seg = 4 / len(chords)
            for k, name in enumerate(chords):
                t0 = (bar * 4 + k * seg) * BEAT
                root, ivs = CHORD_TONES[name]
                truth_chords.append({"start": t0, "end": t0 + seg * BEAT, "chord": name})
                # pad / comping: re-strike chord every beat
                for bt in range(int(seg)):
                    for iv in ivs:
                        add(buf, tone(m2f(root + iv), BEAT, decay=0.6), t0 + bt * BEAT, 0.12)
                # bass: root, octave down, eighth-note pulse
                for e in range(int(seg * 2)):
                    add(buf, tone(m2f(root - 12), BEAT / 2, (1, .6, .4, .2, .1), decay=0.4),
                        t0 + e * BEAT / 2, 0.30)
            pos = 0.0
            for m, d in mel:
                t0 = (bar * 4 + pos) * BEAT
                add(buf, tone(m2f(m), d * BEAT * 0.95, (1, .4, .25, .1), decay=3, vib=4), t0, 0.35)
                truth_notes.append({"start": t0, "end": t0 + d * BEAT, "midi": m})
                pos += d
            bar += 1
    # final G chord ring
    t0 = bar * 4 * BEAT
    for iv in CHORD_TONES["G"][1]:
        add(buf, tone(m2f(55 + iv), 4 * BEAT, decay=2), t0, 0.12)
    add(buf, tone(m2f(43), 4 * BEAT, decay=2), t0, 0.3)
    truth_chords.append({"start": t0, "end": t0 + 4 * BEAT, "chord": "G"})

    buf /= np.abs(buf).max() * 1.1
    wav = out + ".wav"
    sf.write(wav, buf, SR)
    if out.endswith(".mp3") or True:
        subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", wav, "-b:a", "192k", out + ".mp3"],
                       check=True)
    with open(out + ".truth.json", "w") as f:
        json.dump({"bpm": BPM, "key": "G major", "first_downbeat": 0.0,
                   "chords": truth_chords, "notes": truth_notes}, f, indent=1)
    print("wrote", out + ".mp3")


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "test_song")