#!/usr/bin/env python3
"""Synthesise a stereo pop-style track with known stems, so separation can be scored exactly.

Stems (all stereo, 16 kHz):
  voice   formant-synthesised singer (glottal pulses through vowel resonances) with vibrato,
          note-to-note glides, jitter and breath noise; panned dead centre
  chords  sustained pad, chord tones spread left/right
  bass    root notes, centre
  drums   kick (centre), snare (slightly left), hi-hats (right)

Writes <out>/stems/{voice,chords,bass,drums}.wav, <out>/mix.wav, <out>/mix.mp3 (128 kbps)
and <out>/truth.json (melody notes and chords).
"""
import json
import os
import subprocess
import sys

import numpy as np
import soundfile as sf
from scipy.signal import lfilter

SR = 16000
BPM = 96.0
BEAT = 60.0 / BPM
rng = np.random.default_rng(1)
TRANSPOSE = 0
PANS = [-0.7, 0.5, -0.2]
VOWEL_ORDER = ["ah", "oh", "ee", "eh", "oo"]

PHRASE_CHORDS = [["G"], ["D"], ["Em"], ["C"], ["G"], ["D"], ["C", "D"], ["G"]]
CHORD = {"G": (55, [0, 4, 7]), "D": (50, [0, 4, 7]), "Em": (52, [0, 3, 7]), "C": (48, [0, 4, 7])}
PHRASE_MELODY = [
    [(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)], [(78, 2), (76, 1), (74, 1)],
    [(72, 1), (76, 1), (74, 1), (69, 1)], [(67, 4)],
]
VOWELS = {"ah": [(730, 90), (1090, 110), (2440, 160), (3400, 250)],
          "oh": [(570, 80), (840, 100), (2410, 160), (3300, 250)],
          "ee": [(300, 60), (2290, 120), (3010, 180), (3700, 250)],
          "eh": [(530, 70), (1840, 110), (2480, 160), (3500, 250)],
          "oo": [(320, 60), (870, 90), (2240, 150), (3300, 250)]}
m2f = lambda m: 440.0 * 2 ** ((m - 69) / 12)


def pan(x, p):
    """p in [-1, 1]: -1 hard left, 0 centre (equal gain in both channels), 1 hard right."""
    return np.stack([x * np.sqrt((1 - p) / 2) * np.sqrt(2), x * np.sqrt((1 + p) / 2) * np.sqrt(2)]) / np.sqrt(2)


def voice_line(notes, total):
    """One continuous sung line: f0 contour with glides + vibrato, vowel per note, formant filtering."""
    f0 = np.zeros(total)
    amp = np.zeros(total)
    vowel_at = np.empty(total, dtype=object)
    vowel_cycle = VOWEL_ORDER
    for i, (t0, t1, m) in enumerate(notes):
        a, b = int(t0 * SR), min(total, int(t1 * SR))
        t = np.arange(b - a) / SR
        prev = notes[i - 1][2] if i and abs(notes[i - 1][1] - t0) < 0.02 else m
        glide = (prev - m) * np.exp(-t / 0.045)                       # ~45 ms portamento into the note
        vib = 0.3 * np.clip((t - 0.22) / 0.25, 0, 1) * np.sin(2 * np.pi * 5.5 * t)
        f0[a:b] = m2f(m + glide + vib)
        env = np.minimum(1, t / 0.03) * np.minimum(1, (t[-1] - t) / 0.06 + 0.02)
        amp[a:b] = env
        vowel_at[a:b] = vowel_cycle[i % len(vowel_cycle)]
    f0 *= 2 ** (0.08 * (rng.random(total) - 0.5) / 12)                 # jitter
    phase = np.cumsum(np.where(f0 > 0, f0, 0) / SR) % 1.0
    op, cl = 0.6, 0.15
    g = np.where(phase < op, 0.5 * (1 - np.cos(np.pi * phase / op)),
                 np.where(phase < op + cl, np.cos(0.5 * np.pi * (phase - op) / cl), 0.0))
    src = np.diff(g, prepend=0) * (f0 > 0) + 0.004 * rng.standard_normal(total) * (amp > 0)
    out = np.zeros(total)
    # filter each note with its vowel's formants (cascade of 2-pole resonators)
    for i, (t0, t1, m) in enumerate(notes):
        a, b = int(t0 * SR), min(total, int(t1 * SR + 0.08 * SR))
        seg = src[a:b].copy()
        for F, Bw in VOWELS[vowel_cycle[i % len(vowel_cycle)]]:
            r = np.exp(-np.pi * Bw / SR)
            c, d = 2 * r * np.cos(2 * np.pi * F / SR), -r * r
            seg = lfilter([1 - c - d], [1, -c, -d], seg)
        out[a:b] += seg * amp[a:b] if b - a <= len(amp[a:b]) else seg[:len(amp[a:b])] * amp[a:b]
    return out / (np.abs(out).max() + 1e-9), f0


def tone(f, dur, harm, decay):
    n = int(dur * SR)
    t = np.arange(n) / SR
    s = sum(a * np.sin(2 * np.pi * f * k * t) for k, a in enumerate(harm, 1))
    return s * np.minimum(1, t / 0.01) * np.exp(-t / decay) * np.clip((n - np.arange(n)) / (0.03 * SR), 0, 1)


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


def main(out):
    os.makedirs(os.path.join(out, "stems"), exist_ok=True)
    intro, reps = 1, 2
    bars = intro + reps * len(PHRASE_CHORDS) + 1
    total = int((bars * 4 + 1) * BEAT * SR)
    chords = np.zeros((2, total)); bass = np.zeros((2, total)); drums = np.zeros((2, total))
    notes, chord_truth = [], []
    for b in range(bars):
        for beat in range(4):
            t = (b * 4 + beat) * BEAT
            if beat in (0, 2):
                n = int(0.25 * SR); tt = np.arange(n) / SR
                kick = np.sin(2 * np.pi * np.cumsum(50 + 70 * np.exp(-tt * 30)) / SR) * np.exp(-tt * 14)
                add(drums, pan(kick, 0.0), t, 0.8)
            else:
                n = int(0.18 * SR)
                snare = rng.standard_normal(n) * np.exp(-np.arange(n) / SR * 25)
                add(drums, pan(snare, -0.3), t, 0.25)
            for h in (0, 0.5):
                n = int(0.05 * SR)
                hat = np.diff(rng.standard_normal(n + 1)) * np.exp(-np.arange(n) / SR * 80)
                add(drums, pan(hat, 0.6), t + h * BEAT, 0.07)
    bar = intro
    for _ in range(reps):
        for chs, mel in zip(PHRASE_CHORDS, PHRASE_MELODY):
            seg = 4 / len(chs)
            for k, name in enumerate(chs):
                t0 = (bar * 4 + k * seg) * BEAT
                root, ivs = CHORD[name]
                root += TRANSPOSE
                chord_truth.append({"start": t0, "end": t0 + seg * BEAT, "chord": name})
                for j, iv in enumerate(ivs):                               # pad: tones spread L/R, held
                    add(chords, pan(tone(m2f(root + iv), seg * BEAT, (1, .5, .3, .15), 1.5), PANS[j]), t0, 0.16)
                for e in range(int(seg * 2)):
                    add(bass, pan(tone(m2f(root - 12), BEAT / 2, (1, .6, .4, .2, .1), 0.4), 0.0), t0 + e * BEAT / 2, 0.3)
            pos = 0.0
            for m, d in mel:
                t0 = (bar * 4 + pos) * BEAT
                notes.append((t0, t0 + d * BEAT * 0.97, m + TRANSPOSE))
                pos += d
            bar += 1
    v, f0 = voice_line(notes, total)
    voice = pan(v, 0.0) * 0.6
    mix = voice + chords + bass + drums
    peak = np.abs(mix).max() * 1.05
    stems = {"voice": voice, "chords": chords, "bass": bass, "drums": drums}
    for k, s in stems.items():
        sf.write(os.path.join(out, "stems", k + ".wav"), (s / peak).T, SR, subtype="FLOAT")
    sf.write(os.path.join(out, "mix.wav"), (mix / peak).T, SR, subtype="FLOAT")
    np.save(os.path.join(out, "f0_true.npy"), f0)
    subprocess.run(["ffmpeg", "-v", "error", "-y", "-i", os.path.join(out, "mix.wav"), "-b:a", "128k",
                    os.path.join(out, "mix.mp3")], check=True)
    json.dump({"sr": SR, "bpm": BPM, "notes": [{"start": a, "end": b, "midi": m} for a, b, m in notes],
               "chords": chord_truth}, open(os.path.join(out, "truth.json"), "w"), indent=1)
    print(f"wrote {out}: {total / SR:.1f} s, stems {list(stems)}")


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("out", nargs="?", default="data")
    ap.add_argument("--transpose", type=int, default=0, help="semitones (moves voice and chords)")
    ap.add_argument("--seed", type=int, default=1)
    ap.add_argument("--pans", default="-0.7,0.5,-0.2", help="pan of the three chord tones")
    ap.add_argument("--vowels", default="ah,oh,ee,eh,oo")
    a = ap.parse_args()
    rng = np.random.default_rng(a.seed)
    TRANSPOSE, PANS, VOWEL_ORDER = a.transpose, [float(x) for x in a.pans.split(",")], a.vowels.split(",")
    main(a.out)