#!/usr/bin/env python3
# This is free and unencumbered software released into the public domain. For more detail,
# see the LICENCE file at https://github.com/adefossez/seewav
# Original author: adefossez
"""
Generates a nice waveform visualization from an audio file, save it as a mp4 file.
"""
import json
import math
import subprocess as sp
import sys
import tempfile
from pathlib import Path
import shutil
import os
import time
from tqdm import tqdm

# from functools import partial


import numpy as np
import PIL
from PIL import Image, ImageDraw, ImageFont, ImageFilter


from suno_utils.audio import Audio
from suno_utils.utils.text import clean_and_wrap_text

_is_main = False

ASSETS_PATH = Path(os.environ.get("SUNO_ASSETS_PATH", str((Path(__file__).parent / "assets"))))


def colorize(text, color):
    """
    Wrap `text` with ANSI `color` code. See
    https://stackoverflow.com/questions/4842424/list-of-ansi-color-escape-sequences
    """
    code = f"\033[{color}m"
    restore = "\033[0m"
    return "".join([code, text, restore])


def fatal(msg):
    """
    Something bad happened. Does nothing if this module is not __main__.
    Display an error message and abort.
    """
    if _is_main:
        head = "error: "
        if sys.stderr.isatty():
            head = colorize("error: ", 1)
        print(head + str(msg), file=sys.stderr)
        sys.exit(1)


def read_info(media):
    """
    Return some info on the media file.
    """
    proc = sp.run(
        [
            "ffprobe",
            "-loglevel",
            "panic",
            str(media),
            "-print_format",
            "json",
            "-show_format",
            "-show_streams",
        ],
        capture_output=True,
    )
    if proc.returncode:
        raise IOError(f"{media} does not exist or is of a wrong type.")
    return json.loads(proc.stdout.decode("utf-8"))


def read_audio(audio, seek=None, duration=None):
    """
    Read the `audio` file, starting at `seek` (or 0) seconds for `duration` (or all)  seconds.
    Returns `float[channels, samples]`.
    """

    info = read_info(audio)
    channels = None
    stream = info["streams"][0]
    if stream["codec_type"] != "audio":
        raise ValueError(f"{audio} should contain only audio.")
    channels = stream["channels"]
    samplerate = float(stream["sample_rate"])

    # Good old ffmpeg
    command = ["ffmpeg", "-y"]
    command += ["-loglevel", "panic"]
    if seek is not None:
        command += ["-ss", str(seek)]
    command += ["-i", audio]
    if duration is not None:
        command += ["-t", str(duration)]
    command += ["-f", "f32le"]
    command += ["-"]

    proc = sp.run(command, check=True, capture_output=True)
    wav = np.frombuffer(proc.stdout, dtype=np.float32)
    return wav.reshape(-1, channels).T, samplerate


def sigmoid(x):
    return 1 / (1 + np.exp(-x))


def envelope(wav, window, stride):
    """
    Extract the envelope of the waveform `wav` (float[samples]), using average pooling
    with `window` samples and the given `stride`.
    """
    # pos = np.pad(np.maximum(wav, 0), window // 2)
    wav = np.pad(wav, window // 2)
    out = []
    for off in range(0, len(wav) - window, stride):
        frame = wav[off : off + window]
        out.append(np.maximum(frame, 0).mean())
    out = np.array(out)
    # Some form of audio compressor based on the sigmoid.
    out = 1.9 * (sigmoid(2.5 * out) - 0.5)
    return out


def draw_env(envs, out, fg_colors, bg_color, size, text_image):
    """
    Internal function, draw a single frame (two frames for stereo) using Pillow and save
    it to the `out` file as png. envs is a list of envelopes over channels, each env
    is a float[bars] representing the height of the envelope to draw. Each entry will
    be represented by a bar.
    """

    im = PIL.Image.new(mode="RGB", size=size)

    draw = ImageDraw.Draw(im)

    draw.rectangle((0, 0, size[0], size[1]), fill=bg_color)

    K = len(envs)  # Number of waves to draw (waves are stacked vertically)
    T = len(envs[0])  # Numbert of time steps
    pad_ratio = 0.1  # spacing ratio between 2 bars
    width = 1.0 / (T * (1 + 2 * pad_ratio))
    pad = pad_ratio * width
    delta = 2 * pad + width
    # ctx.set_line_width(width)
    for i in range(K):
        m_scale = min(envs[i])
        scale = max(envs[i]) - min(envs[i])
        for step in range(T):
            rel_scale = (envs[i][step] - m_scale) / scale if scale != 0 else 1
            half = 0.5 * envs[i][step]  # (semi-)height of the bar
            half /= K  # as we stack K waves vertically
            midrule = ((2 if i == 1 else 0) + 2 * i) / (2 * K)  # midrule of i-th wave
            draw.line(
                (
                    (pad + step * delta) * size[0],
                    midrule * size[1],
                    (pad + step * delta) * size[0],
                    (midrule + (-half if i == 1 else 0.9 * half)) * size[1],
                ),
                fill=tuple(
                    int(bg_color[j] + rel_scale * (x - bg_color[j])) for j, x in enumerate(fg_colors[i])
                ),
                width=int(width * size[0]),
            )

    im.paste(text_image, (0, 0), text_image)
    return im


def interpole(x1, y1, x2, y2, x):
    return y1 + (y2 - y1) * (x - x1) / (x2 - x1)


def generate_frame(
    idx,
    rate,
    sr,
    stride,
    bars,
    envs,
    speed,
    smooth,
    tmp,
    fg_color,
    fg_color2,
    bg_color,
    size,
    text_image,
    offset: int,
):
    pos = ((idx / rate) * sr) / stride / bars
    off = int(pos)
    loc = pos - off
    denvs = []
    for env in envs:
        env1 = env[off * bars : (off + 1) * bars]
        env2 = env[(off + 1) * bars : (off + 2) * bars]

        # we want loud parts to be updated faster
        maxvol = math.log10(1e-4 + env2.max()) * 10
        speedup = np.clip(interpole(-6, 0.5, 0, 2, maxvol), 0.5, 2)
        w = sigmoid(speed * speedup * (loc - 0.5))
        denv = (1 - w) * env1 + w * env2
        denv *= smooth
        denvs.append(denv)

    # blurred_image = blurred_image.rotate(int(idx / frames_count) * 360)
    # blurred_image.paste(image, mask=image)
    # blurred_image.show()
    # text_image.paste(blurred_image, (0, 0))

    return draw_env(
        denvs,
        tmp / f"{idx+offset:06d}.png",
        (fg_color, fg_color2),
        bg_color,
        size,
        text_image,
    )


def contains_chinese(text):
    for char in text:
        if "\u4e00" <= char <= "\u9fff":
            return True
    return False


def contains_hebrew(text):
    for char in text:
        if "\u0590" <= char <= "\u05FF":
            return True
    return False


def contains_arabic(text):
    for char in text:
        if "\u0600" <= char <= "\u06FF":
            return True
    return False


def contains_japanese(text):
    for char in text:
        if (
            "\u3000" <= char <= "\u303f"
            or "\u3040" <= char <= "\u309f"  # Japanese-style punctuation
            or "\u30a0" <= char <= "\u30ff"  # Hiragana
            or "\u4e00" <= char <= "\u9faf"  # Katakana
            or "\u3400" <= char <= "\u4dbf"  # Kanji  # Rare Kanji
        ):
            return True
    return False


def contains_korean(text):
    for char in text:
        if "\uac00" <= char <= "\ud7af":  # Hangul Syllables
            return True
    return False


def get_font(text: str, font_size: int, is_square: bool):
    is_english = text.isascii()

    is_cjk = not is_english and (
        contains_chinese(text) or contains_japanese(text) or contains_korean(text)
    )

    is_hebrew = (not is_english) and contains_hebrew(text)
    is_arabic = (not is_english) and contains_arabic(text)
    # print("text", text, is_english, is_cjk, is_hebrew, is_arabic)
    font = (
        (
            ImageFont.truetype(
                str((ASSETS_PATH / "Roobert-SemiBold.ttf").resolve()),
                font_size,
            )
            if is_square
            else ImageFont.truetype(str((ASSETS_PATH / "Manrope-SemiBold.ttf").resolve()), font_size)
        )
        if is_english
        else (
            ImageFont.truetype(
                str((ASSETS_PATH / "NotoSansHebrew-Medium.ttf").resolve()),
                font_size,
            )
            if is_hebrew
            else ImageFont.truetype(str((ASSETS_PATH / "NotoSansCJK-Medium.ttc").resolve()), font_size)
            if (is_cjk)
            else ImageFont.truetype(
                str((ASSETS_PATH / "NotoSansArabic-VariableFont_wdth,wght.ttf").resolve()),
                font_size,
            )
            if (is_arabic)
            else ImageFont.truetype(str((ASSETS_PATH / "NotoSans-Medium.ttf").resolve()), font_size)
        )
    )
    return font


def visualize_section(
    audio: Path,
    tmp: Path,
    ffmpeg_pipe: sp.Popen,
    png_path: Path | None = None,
    is_square: bool | None = False,
    seek=None,
    duration=None,
    rate=60,
    bars=50,
    speed=4,
    audio_time=0.4,
    oversample=3,
    fg_color=(62, 71, 175),
    fg_color2=(62, 71, 175),
    bg_color=(1, 1, 1),
    size=(400, 400),
    stereo=False,
    text="",
    image_transparency=0,
    offset: int = 0,
    tags=None,
    title=None,
    deactivate_animation=True,
    aligned_text=[],
):
    print("Start visualization...")
    start_time = time.time()
    try:
        wav, sr = read_audio(audio, seek=seek, duration=duration)
    except (IOError, ValueError) as err:
        fatal(err)
        raise
    # wavs is a list of wav over channels
    wavs = []
    if stereo:
        if wav.shape[0] == 2:
            # "stereo requires stereo audio file, do demuc conversion"
            # take < 1 sec typically
            # TODO: could also look into
            # https://github.com/facebookresearch/demucs
            # out_audio = demucs.encode(torch.from_numpy(wav.copy()))
            # arr_bg = out_audio[:3].sum(0)
            # arr_vocals = out_audio[3]
            wavs.append(wav[0])
            wavs.append(wav[1])
        else:
            raise ValueError(f"Don't accept the input size {wav.shape[0]}")
    else:
        wav = wav.mean(0)
        wavs.append(wav)

    for i, wav in enumerate(wavs):
        wavs[i] = wav / wav.std()

    window = int(sr * audio_time / bars)
    stride = int(window / oversample)
    # envs is a list of env over channels
    envs = []
    for wav in wavs:
        env = envelope(wav, window, stride)
        env = np.pad(env, (bars // 2, 2 * bars))
        envs.append(env)

    duration = len(wavs[0]) / sr
    frames = int(rate * duration)
    smooth = np.hanning(bars)

    if text is None:
        text = ""

    margin: int = 0 if is_square else 40
    padding: int = 40 if is_square else 5
    stroke_width: int = 4 if is_square else 0
    image_size = min(size[0], size[1]) - (2 * margin) if png_path is not None else 0
    # text_images = []

    logo_src = PIL.Image.open(ASSETS_PATH / "Suno-Logo-Lockup-White.png")
    logo_width, logo_height = logo_src.size
    logo_scale_ratio = 0.08
    logo_src = logo_src.resize(
        (math.ceil(logo_width * logo_scale_ratio), math.ceil(logo_height * logo_scale_ratio)),
        resample=PIL.Image.BICUBIC,
    ).convert("RGBA")

    if png_path is None or png_path == "":
        png_path = ASSETS_PATH / "CoverPlaceholder.png"
    image_src = PIL.Image.open(png_path)
    image_src = image_src.convert("RGBA")
    image_src_resized = image_src.copy()
    image_src_resized = image_src_resized.resize(
        (360, 360),
        # (math.ceil(width * scale_ratio), math.ceil(height * scale_ratio)),
        resample=PIL.Image.BICUBIC,
    )
    blurred_image_src = image_src.copy()
    blurred_image_src = blurred_image_src.filter(ImageFilter.GaussianBlur(radius=50))
    blurred_image_src = blurred_image_src.resize(
        # (100, 100),
        # (math.ceil(width * scale_ratio), math.ceil(height * scale_ratio)),
        (math.ceil(size[0] * 2.25), math.ceil(size[1] * 2.25)),
        resample=PIL.Image.BICUBIC,
    )
    # text_draw_src = ImageDraw.Draw(text_image_src)
    text_image_src = PIL.Image.new(mode="RGBA", size=size, color=(255, 255, 255, 0))
    empty_lyrics = len(aligned_text) == 0
    has_any_aligned_time = any(t for t in aligned_text if "start_s" in t)
    if not empty_lyrics and not has_any_aligned_time:
        # assuming evenly spaced
        for i, t in enumerate(aligned_text):
            t["start_s"] = (i / len(aligned_text)) * duration
    try:
        text = clean_and_wrap_text(text)
    except Exception as e:
        print("Error cleaning text", e)
        pass

    print(f"Before calculating fonts... {time.time() - start_time:.2f} seconds.")

    font_spacing = 7
    font_size = 18
    lyrics_font = get_font(text, font_size, is_square)

    if tags is None:
        tag_text = ""
    else:
        tag_text = f"{tags[0:30]}..." if len(tags) > 30 else tags
    tag_font = get_font(tag_text, font_size=18, is_square=is_square)

    if title is None:
        title_text = ""
    else:
        title_text = f"{title[0:30]}..." if len(title) > 30 else title
    title_font = get_font(title_text, font_size=32, is_square=is_square)

    text_image = text_image_src.copy()

    blurred_image_src = blurred_image_src.crop(
        (0, 0, size[0] * 2.25, size[1] * 2.25),
    )
    lyrics_img_src = Image.new("RGBA", size=(500, 210), color=(255, 255, 255, 0))

    roobert_font_1 = ImageFont.truetype(
        str((ASSETS_PATH / "Roobert-SemiBold.ttf").resolve()),
        32,
    )
    roobert_font_2 = ImageFont.truetype(
        str((ASSETS_PATH / "Roobert-SemiBold.ttf").resolve()),
        18,
    )
    text_draw_test = ImageDraw.Draw(text_image_src)
    tags_bbox = text_draw_test.textbbox(
        (0, 0),
        tag_text,
        font=roobert_font_2,
    )
    title_bbox = text_draw_test.textbbox(
        (0, 0),
        title_text,
        font=roobert_font_1,
    )

    # these are some shared images
    rect_image = Image.new("RGBA", image_src_resized.size, (255, 255, 255, 0))
    rect_draw = ImageDraw.Draw(rect_image)
    rect_draw.rounded_rectangle(
        (0, 0, image_src_resized.size[0], image_src_resized.size[1]),
        fill=(255, 255, 255, 255),
        radius=16,
    )
    rect_image2 = Image.new(
        "RGBA", (image_src_resized.size[0] + 20, image_src_resized.size[1] + 20), (255, 255, 255, 0)
    )
    rect_draw2 = ImageDraw.Draw(rect_image2)
    # rect_draw2.rounded_rectangle(
    #     (0, 0, rect_image2.size[0], rect_image2.size[1]), fill=(116, 116, 116, 255), radius=20
    # )
    rect_image3 = Image.new(
        "RGBA", (image_src_resized.size[0] + 20, image_src_resized.size[1] + 20), (116, 116, 116, 255)
    )
    rect_draw3 = ImageDraw.Draw(rect_image3)
    text_width = 500
    offset_from_image_total = ((size[0] - image_size) - (margin)) - text_width
    offset_from_image = offset_from_image_total * 0.5
    text_start_pos_x = (
        (size[0] - text_width) * 0.5 if is_square else margin + image_size + offset_from_image
    )
    text_start_pos_y = ((size[1] - 300) * 0.5) + 400

    def process_frame(idx):
        assert png_path is not None
        text_image = text_image_src.copy()
        # apply a rotation based on the index
        degrees = idx
        blurred_image = blurred_image_src.copy()
        blurred_image = blurred_image.rotate(degrees)
        # PIE SLICE
        rect_draw3.pieslice(
            (0, 0, rect_image3.size[0], rect_image3.size[1]),
            225 + degrees,
            270 + degrees,
            fill=(255, 255, 255, 255),
        )
        rect_image3_pieslice = rect_image3.filter(ImageFilter.GaussianBlur(radius=30))

        text_image.paste(
            blurred_image, (-1 * int(0.5 * size[0] * 1.125), -1 * int(0.5 * size[1] * 1.125))
        )
        text_image_start_x = int(size[0] / 2) - int(image_src_resized.width / 2) - 10
        text_image_start_y = 80 if not empty_lyrics else 230
        # text_image.paste(
        #     rect_image2,
        #     (
        #         text_image_start_x,
        #         text_image_start_y,
        #     ),
        #     rect_image2,
        # )
        # text_image.paste(
        #     rect_image3_pieslice,
        #     (
        #         text_image_start_x,
        #         text_image_start_y,
        #     ),
        #     mask=rect_image2,
        # )
        text_image.paste(
            image_src_resized,
            (text_image_start_x + 10, text_image_start_y - 40),
            rect_image,
        )
        logo = logo_src.copy()
        text_image.paste(logo, (text_image.size[0] - 190, text_image.size[1] - 75), logo)

        text_draw2 = ImageDraw.Draw(text_image)
        text_offset = (
            (max(0, text_start_pos_x) if is_square else max(image_size + margin, text_start_pos_x)),
            max(0, text_start_pos_y),
        )
        if tag_text:
            text_draw2.text(
                (
                    int(size[0] / 2) - int(tags_bbox[2] / 2),
                    int(text_offset[1]) - (240 if not empty_lyrics else 110) + 55,
                ),
                tag_text,
                fill=(255, 255, 255, 166),
                font=tag_font,
            )
        if title_text:
            text_draw2.text(
                (
                    int(size[0] / 2) - int(title_bbox[2] / 2),
                    int(text_offset[1]) - (280 if not empty_lyrics else 150) + 53,
                ),
                title_text,
                fill=(255, 255, 255, 166),
                font=title_font,
            )
        return (text_image, idx)

    print(f"Before generating the frames... {time.time() - start_time:.2f} seconds.")
    # text_imgs = [process_frame(idx) for idx in range(frames)]

    # we don't need to generate more frames than the requested frames, or max 360
    degree_imgs = [process_frame(idx) for idx in range(min(360, frames))]
    print(f"After generating the degree images... {time.time() - start_time:.2f} seconds.")

    len_degree = len(degree_imgs)
    text_imgs = [degree_imgs[idx % len_degree] for idx in range(frames)]

    # text_imgs.sort(key=lambda i: i[1])
    # text_images = [i[0] for i in text_imgs]
    text_images = []
    last_index = 0
    end_index = 0

    lyrics_img_cache = None
    lyrics_bbox_cache = None
    last_index_cache = None

    # clean up aligned_text to all the expected dicts, as there is also full hoot_lyrics
    aligned_text = [w for w in aligned_text if "word" in w]
    print(f"aligned_text len: {len(aligned_text)}")

    print(f"Before pasting the text... {time.time() - start_time:.2f} seconds.")
    for idx, image in enumerate(text_imgs):
        text_width = 500
        image_copy = image[0].copy()

        last_index_copy = last_index

        offset_from_image_total = ((size[0] - image_size) - (margin)) - text_width
        offset_from_image = (
            # offset_from_image_total - margin
            # if offset_from_image_total < 100
            # else
            offset_from_image_total * 0.5
        )
        text_start_pos_x = (
            (size[0] - text_width) * 0.5
            if is_square
            else (
                margin + image_size + offset_from_image
                if png_path is not None
                else ((size[0] - text_width) * 0.5)
            )
        )
        text_start_pos_y = ((size[1] - 300) * 0.5) + 400
        text_offset = (
            (max(0, text_start_pos_x) if is_square else max(image_size + margin, text_start_pos_x)),
            max(0, text_start_pos_y),
        )

        elapsed = (idx / frames) * duration
        previous_word = None

        for i, word in enumerate(aligned_text[last_index:]):
            if (
                "start_s" in word
                and word["start_s"] < max(0, elapsed - 3)  # backset 3 sec to get prev line
                and (previous_word is None or previous_word.endswith("\n"))
            ):
                last_index_copy = i + last_index

            previous_word = word["word"]

            if last_index_copy + 40 > len(aligned_text):
                end_index = len(aligned_text)

            if "start_s" in word and word["start_s"] > elapsed and last_index_copy == last_index_cache:
                break

            if (
                "word" in word
                and word["word"].endswith("\n")
                and (i + last_index > last_index_copy + 40 and i + last_index < last_index_copy + 60)
            ):
                end_index = i + last_index
                break

        last_index = last_index_copy
        lyrics_text = clean_and_wrap_text(
            "".join([t["word"] for t in aligned_text[last_index : end_index + 1]])
        )

        if last_index == last_index_cache:
            lyrics_img = lyrics_img_cache.copy()
            lyrics_bbox = lyrics_bbox_cache
        else:
            lyrics_img = lyrics_img_src.copy()
            text_draw = ImageDraw.Draw(lyrics_img)

            lyrics_bbox = text_draw.multiline_textbbox(
                # text_offset,
                (0, 0),
                lyrics_text,
                font=lyrics_font,
                spacing=font_spacing,
                align="left",
            )
            text_draw.multiline_text(
                # text_offset,
                (0, 0),
                lyrics_text,
                fill=(255, 255, 255, 255),
                font=lyrics_font,
                spacing=font_spacing,
                align="left",
            )

            lyrics_bbox_cache = lyrics_bbox
            lyrics_img_cache = lyrics_img.copy()
            last_index_cache = last_index

        # if lyrics_bbox[3] > 300:
        #    delta = (lyrics_bbox[3] - 200) / frames
        # else:
        #    delta = 0
        # if delta > 0:
        #    lyrics_img = lyrics_img.crop(
        #        (0, min(idx * delta, lyrics_bbox[3]), lyrics_img.size[0], lyrics_bbox[3] + 300)
        #    )
        # pixels = lyrics_img.load()
        # for y in range(0, 5):
        #    max_alpha = 166 - (int(166 / 5) * (5 - y))
        #    for x in range(lyrics_img.size[0]):
        #        pixels[x, y] = pixels[x, y][:3] + (min(pixels[x, y][3], max_alpha),)

        # lyrics_bottom = size[1] - (int(text_offset[1]) - 180)
        # for y in range(lyrics_bottom - 8, lyrics_bottom):
        #    max_alpha = 166 - (int(166 / 8) * (8 - (lyrics_bottom - y)))
        #    for x in range(lyrics_img.size[0]):
        #        pixels[x, y] = pixels[x, y][:3] + (min(pixels[x, y][3], max_alpha),)

        image_copy.paste(
            lyrics_img,
            (
                int(size[0] / 2) - int(lyrics_bbox[2] / 2),
                int(text_offset[1]) - 150,
            ),
            lyrics_img,
        )
        text_images.append(image_copy)

    print(f"Before FFMPEG... {time.time() - start_time:.2f} seconds.")
    if not deactivate_animation:
        for idx in tqdm(range(frames), desc="Generating frames"):
            im = generate_frame(
                idx,
                rate,
                sr,
                stride,
                bars,
                envs,
                speed,
                smooth,
                tmp,
                fg_color,
                fg_color2,
                bg_color,
                size,
                text_image,
                len(frames),
                offset,
            )
            # https://jdhao.github.io/2019/07/20/pil_jpeg_image_quality/
            im.save(ffmpeg_pipe.stdin, "bmp")
    else:
        for idx in tqdm(range(len(text_images)), desc="Generating frames"):
            text_images[idx].convert("RGB").save(ffmpeg_pipe.stdin, "bmp")

    print(f"Generated frames in {time.time() - start_time:.2f} seconds.")
    return frames


def visualize(
    audio: Path,
    tmp: Path,
    out: Path,
    png_path: Path | None = None,
    is_square: bool | None = False,
    seek=None,
    duration=None,
    rate=2,
    bars=50,
    speed=4,
    audio_time=0.4,
    oversample=3,
    fg_color=(62, 71, 175),
    fg_color2=(62, 71, 175),
    bg_color=(1, 1, 1),
    size=(400, 600),
    stereo=False,
    text="",
    image_transparency=0.2,
    timestamps=[],
    tags=None,
    title=None,
    deactivate_animation=True,
    aligned_text=[],
):
    """
    Generate the visualisation for the `audio` file, using a `tmp` folder and saving the final
    video in `out`.
    `seek` and `durations` gives the extract location if any.
    `rate` is the framerate of the output video.

    `bars` is the number of bars in the animation.
    `speed` is the base speed of transition. Depending on volume, actual speed will vary
        between 0.5 and 2 times it.
    `time` amount of audio shown at once on a frame.
    `oversample` higher values will lead to more frequent changes.
    `fg_color` is the rgb color to use for the foreground.
    `fg_color2` is the rgb color to use for the second wav if stereo is set.
    `bg_color` is the rgb color to use for the background.
    `size` is the `(width, height)` in pixels to generate.
    `stereo` is whether to create 2 waves.
    """
    # https://hamelot.io/visualization/using-ffmpeg-to-convert-a-set-of-images-into-a-video/
    # https://stackoverflow.com/questions/55800185/my-ffmpeg-output-always-add-extra-30s-of-silence-at-the-end
    input_lst = (
        [
            "ffmpeg",
            "-y",
            "-loglevel",
            "panic",
            "-r",
            str(rate),
            "-f",
            "image2pipe",
            "-s",
            f"{size[0]}x{size[1]}",
            "-i",
            "-",
            "-i",
            audio,
            "-c:v",
            "copy",
            "-c:a",
            "aac",
            "-b:a",
            "192k",
            "-map",
            "0:v",
            "-map",
            "1:a",
            "-vcodec",
            "libx264",
            "-crf",
            "28",
            "-preset",
            "ultrafast",
            "-pix_fmt",
            "yuv420p",
            "-shortest",
            "-fflags",
            "+shortest",
            "-max_interleave_delta",
            "100M",
        ]
        + (["-t", f"{round(duration, 2)}"] if duration is not None else [])
        + [out.resolve()]
    )
    # print(input_lst)
    pipe = sp.Popen(input_lst, stdin=sp.PIPE)

    if timestamps:
        frame_offset = 0
        duration = 0
        elapsed_duration = 0
        for i in range(len(timestamps)):
            start_time = timestamps[i]["t"]
            lyric = timestamps[i]["l"]
            elapsed_duration += duration
            int(duration * rate)
            # frame_offset += frames
            if i < len(timestamps) - 1:
                next_start_time = timestamps[i + 1]["t"]
                duration = next_start_time - start_time
            else:
                duration = None
            frame_offset += visualize_section(
                audio,
                tmp,
                pipe,
                png_path,
                is_square,
                start_time,
                duration,
                rate,
                bars,
                speed,
                audio_time,
                oversample,
                fg_color,
                fg_color2,
                bg_color,
                size,
                stereo,
                lyric,
                image_transparency,
                frame_offset,
                tags,
                title,
                deactivate_animation=deactivate_animation,
                aligned_text=aligned_text,
            )
    else:
        visualize_section(
            audio,
            tmp,
            pipe,
            png_path,
            is_square,
            seek,
            duration,
            rate,
            bars,
            speed,
            audio_time,
            oversample,
            fg_color,
            fg_color2,
            bg_color,
            size,
            stereo,
            text,
            image_transparency,
            0,
            tags,
            title,
            deactivate_animation=deactivate_animation,
            aligned_text=aligned_text,
        )

    pipe.stdin.close()
    pipe.wait()


def render_demucs_audio(
    audio: Audio,
    text: str,
    out: Path | str | None = None,
    png_path: Path | None = None,
    is_square: bool | None = False,
    image_transparency: float = 0.6,
    timestamps=[],
    size=(512, 768),
    tags: str | None = None,
    title: str | None = None,
    stereo: bool = True,
    deactivate_animation: bool = True,
    aligned_text=[],
) -> Path:
    """Demucs an audio file and render a video with a waveform visualization."""
    start_time = time.time()
    print("Audio duration is ", audio.duration_s)
    with tempfile.TemporaryDirectory() as temp_dir:
        # TODO: this process is a bit unnecessary since we already have the audio
        # However, this requires a change in visualize API
        # this is a wav operation but the wav isn't uploaded to s3
        orig_wav_path = Path(temp_dir) / "orig.wav"
        print(f"Before to wav {time.time() - start_time:.2f} seconds.")
        audio.write_wav(str(orig_wav_path))
        print(f"Before final conversion of full video in {time.time() - start_time:.2f} seconds.")
        visualize(
            orig_wav_path,
            Path(temp_dir),
            Path(temp_dir) / "out.mp4",
            png_path=png_path,
            is_square=is_square,
            duration=round(audio.duration_s, 2),
            size=size,
            bg_color=(0, 3, 33),
            rate=5,
            stereo=stereo,
            text=text,
            # text="",
            tags=tags,
            title=title,
            image_transparency=image_transparency,
            timestamps=timestamps,
            deactivate_animation=deactivate_animation,
            aligned_text=aligned_text,
        )
        out_path = Path(temp_dir) / "out.mp4"
        print(f"Generated full video in {time.time() - start_time:.2f} seconds.")

        if out:
            shutil.copy(out_path, out)
            return Path(out)
        return out_path