"""Measure rendered-mix level, attacks, and texture changes in local FLACs.

Requires ffmpeg, numpy, matplotlib. This does not transcribe melodies or chords.
Positive spectral flux includes percussion; RMS is influenced by mastering.
"""
import os
os.environ['MPLCONFIGDIR'] = '/tmp/motif-matplotlib'
import subprocess, json, argparse
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
parser = argparse.ArgumentParser(description='Measure relative level and attack/texture changes in local FLAC references. Requires ffmpeg, numpy, and matplotlib. Does not transcribe notes or infer emotion.')
parser.add_argument('folder', nargs='?', default='references/devath')
parser.add_argument('--output', default='docs/rivet/research/measurements')
args = parser.parse_args()
root = Path(args.folder)
out = Path(args.output)
out.mkdir(parents=True, exist_ok=True)
files = sorted(root.glob('*.flac'))
if not files:
    parser.error('No FLAC files in the selected folder.')
fig, axes = plt.subplots(len(files), 1, squeeze=False, figsize=(15, 2.0 * len(files)), constrained_layout=True)
summary = []
axes = axes[:, 0]
for ax, path in zip(axes, files):
    raw = subprocess.check_output(['ffmpeg', '-v', 'error', '-i', str(path), '-f', 'f32le', '-ac', '1', '-ar', '11025', 'pipe:1'])
    y = np.frombuffer(raw, dtype='<f4')
    sr = 11025
    hop = 256
    size = 2048
    frames = np.lib.stride_tricks.sliding_window_view(y, size)[::hop]
    window = np.hanning(size)
    spectra = []
    rms = []
    for start in range(0, len(frames), 256):
        f = frames[start:start + 256]
        spectra.append(np.abs(np.fft.rfft(f * window, axis=1)).astype(np.float32))
        rms.extend(np.sqrt(np.mean(f * f, axis=1)))
    mag = np.concatenate(spectra)
    freq = np.fft.rfftfreq(size, 1 / sr)
    times = np.arange(len(mag)) * hop / sr
    log = np.log1p(mag)
    flux = np.maximum(np.diff(log, axis=0, prepend=log[:1]), 0).sum(axis=1)
    one = np.arange(0, len(y) / sr, 1.0)

    def seconds(values):
        return np.array([np.mean(values[(times >= t) & (times < t + 1)]) if np.any((times >= t) & (times < t + 1)) else 0 for t in one])
    energy = seconds(np.asarray(rms))
    attack = seconds(flux)
    bands = np.stack([seconds(mag[:, (freq >= lo) & (freq < hi)].sum(axis=1)) for lo, hi in [(40, 180), (180, 1000), (1000, 5000)]], axis=1)
    features = np.log1p(bands)
    features = (features - features.mean(0)) / (features.std(0) + 1e-08)
    novelty = np.zeros(len(one))
    for t in range(8, len(one) - 8):
        novelty[t] = np.linalg.norm(features[t - 8:t].mean(0) - features[t:t + 8].mean(0))
    peaks = []
    for idx in np.argsort(novelty)[::-1]:
        if idx < 16 or idx > len(one) - 16:
            continue
        if all((abs(int(idx) - p) >= 16 for p in peaks)):
            peaks.append(int(idx))
        if len(peaks) == 6:
            break
    peaks.sort()
    ax.plot(one, energy / (np.percentile(energy, 95) + 1e-08), label='Relative RMS', color='#8b77cf', linewidth=1)
    ax.plot(one, attack / (np.percentile(attack, 95) + 1e-08), label='Spectral attacks', color='#3f8d9e', alpha=0.65, linewidth=0.8)
    for p in peaks:
        ax.axvline(p, color='#b7823c', alpha=0.45, linewidth=0.7)
        ax.text(p, 1.65, f'{p // 60}:{p % 60:02}', fontsize=7, rotation=90)
    ax.set(title=path.stem, xlim=(0, len(y) / sr), ylim=(0, 2), ylabel='Relative level')
    ax.grid(alpha=0.15)
    item = {'track': path.stem, 'duration_seconds': round(len(y) / sr, 2), 'candidate_texture_changes_seconds': peaks, 'rms_20th_to_80th_ratio': round(float(np.percentile(energy, 80) / (np.percentile(energy, 20) + 1e-08)), 2), 'spectral_attack_20th_to_80th_ratio': round(float(np.percentile(attack, 80) / (np.percentile(attack, 20) + 1e-08)), 2)}
    summary.append(item)
    print(json.dumps(item), flush=True)
axes[0].legend(loc='upper right')
axes[-1].set_xlabel('Seconds')
fig.savefig(out / 'reference-timelines.png', dpi=130)
(out / 'measurements.json').write_text(json.dumps(summary, indent=2))