# #!/bin/zsh

import os
import subprocess

import cv2
import numpy as np

for songname in [
        '01-guest',
        '02_dancing_happy',
        '03-bear-trim2',
        '03-bear',
        '04-disneyswing',
        '05-encanto',
        '06-frozen',
        '07-onejump',
        '08-lionking',
        '09-pianoman-mix',
        '10-disneytap',
        '11-club',
        '12-sunnyside2',
        '13-supercali',
        '14-groove',
        '15-mermaid',
        '16-all',
        '17-parade-mix',
]:

    WAV_DIR = f"/tmp/htdemucs/{songname}"
    instruments = [
        ("vocals", "3k", (1, .3, 0)),
        ("other", "3k", (1, .8, .5)),
        ("drums", "1k", (1, .9, 1)),
        ("bass", "400", (0, .4, 1)),
    ]

    def generate_spectrogram(input_file, output_file, rate):
        command = ["sox", input_file, "-n", "remix", "1", "rate", rate, "spectrogram", "-X", "50", "-y", "100", "-z", "80", "-m", "-r", "-o", output_file]
        subprocess.check_call(command)

    def tint_image(img, tint):
        img = img.astype(np.float32) / 255
        tint_img = np.full((img.shape[0], img.shape[1], 3), tint[::-1], dtype=np.float32)
        return cv2.multiply(img, tint_img)

    stack = []
    for name, rate, tint in instruments:
        input_file = f"{WAV_DIR}/{name}.wav"
        spectro_file = f"/tmp/spectro_{name}.png"
        generate_spectrogram(input_file, spectro_file, rate)
        img = cv2.imread(spectro_file)
        tinted_img = tint_image(img, tint)
        stack.append(tinted_img)

    out = np.concatenate(stack, axis=0)
    cv2.imwrite(os.path.expandvars(f"$LIGHT9_SHOW/spectrogram/{songname}.png"), out * 255)