import threading
from pathlib import Path
import time

import numpy as np
import matplotlib.figure as mfigure

target = Path("/tmp/p_test")
target.mkdir(parents=True, exist_ok=True)

ω = 2
N = 1024 * 10

N_threads = 70


def gen_frames(steps):
    fig = mfigure.Figure(layout="constrained")
    ax = fig.subplots()
    ax.set_xlim(0, 2 * np.pi)
    ax.set_ylim(-1.1, 1.1)
    # ax.axis("off")

    x = np.linspace(0, 2 * np.pi, 1024 * 10)
    (ln,) = ax.plot(x, np.ones_like(x))

    fr_ann = ax.annotate(
        "step:        ",
        (1, 1),
        xytext=(-1, -5),
        xycoords="axes fraction",
        textcoords="offset points",
        ha="right",
        va="top",
        font="hack",
        visible=True,
    )
    φ_ann = ax.annotate(
        "φ:        ",
        (1, 0),
        xytext=(0, -5),
        xycoords=fr_ann,
        textcoords="offset points",
        ha="right",
        va="top",
        font="hack",
        visible=True,
    )

    for step in steps:
        φ = 2 * np.pi * step / 1024
        y = np.sin(x * ω + φ)
        ln.set_data(x, y)
        fr_ann.set_text(f"n: {step: 5d}")
        φ_ann.set_text(f"φ: {φ:.4f}")
        fig.savefig(target / f"sin_{step:05d}.png")


threads = []
for j in range(N_threads):
    threads.append(threading.Thread(target=gen_frames, args=(range(j, N, N_threads),)))

start = time.monotonic()
[t.start() for t in threads]
[t.join() for t in threads]
stop = time.monotonic()

print(f"Took {stop - start:.2f} for {N} frames with {N_threads} threads")