import json

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import hsv_to_rgb
from scipy.io import wavfile

plt.rcParams["figure.figsize"] = [18, 9]
plt.style.use("dark_background")
plt.subplots_adjust(left=0.05, right=0.99, top=0.95, bottom=0.05)


with open("metrics.json", "r") as file:
    metrics = json.load(file)
    entries = metrics["entries"]
    cors = metrics["cors"]


def render_ploc():
    X, Y, Ya, Yd = np.array(entries).transpose()
    plt.clf()
    plt.plot(X, Y)
    plt.fill_between(X, Y, Y - Ya, alpha=0.5, color="green")
    plt.fill_between(X, Y, Y + Yd, alpha=0.5, color="red")
    plt.savefig("/code/metrics/metrics.png")


def render_cors(*, px, pc):
    M = len(cors)
    N = len(cors[0])
    Y = np.array(cors)
    if px:
        X = abs(np.diff(Y, axis=1, prepend=0)).sum(axis=0).cumsum()
    else:
        X = np.arange(N)
    assert Y.shape == (M, N)
    W = Y.max(axis=1)
    I = W.cumsum()
    assert I.shape == (M,)
    plt.clf()
    if pc:
        Z = (I - W / 2) / I.max() * M
    else:
        Z = np.arange(M)
    colours = hsv_to_rgb(np.array([3.0 * Z / M % 1, 0.4 + Z * 0, 0.6 + 0.15 * (Z / M)]).transpose())
    plt.stackplot(
        X,
        Y,
        colors=colours,
        aa=False,
    )
    plt.savefig(f"/code/metrics/metrics-px{int(px)}-pc{int(pc)}.png")


render_ploc()
render_cors(px=False, pc=False)
render_cors(px=False, pc=True)
render_cors(px=True, pc=False)
render_cors(px=True, pc=True)