import json
from functools import cache

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import hsv_to_rgb

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)
plt.margins(x=0.025)


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, linewidth=1.0)
    plt.fill_between(
        X,
        Y - Ya,
        Y,
        alpha=0.5,
        color="green",
        linewidth=0.0,
    )
    plt.fill_between(
        X,
        Y,
        Y + Yd,
        alpha=0.5,
        color="red",
        linewidth=0.0,
    )
    plt.ylim(bottom=0)
    plt.savefig("/code/metrics/metrics.png")


@cache
def get_m():
    return len(cors)


@cache
def get_n():
    return len(cors[0])


@cache
def get_y():
    return np.array(cors)


@cache
def _get_x(px: bool):
    if px:
        return abs(np.diff(get_y(), axis=1, prepend=0)).sum(axis=0).cumsum()
    else:
        return np.arange(get_n())


def get_x(*, px: bool):
    return _get_x(px)


@cache
def get_w():
    return get_y().max(axis=1)


@cache
def _get_z(pc: bool):
    M = get_m()
    if pc:
        W = get_w()
        I = W.cumsum()
        assert I.shape == (M,)
        return (I - W / 2) / I.max() * M
    else:
        return np.arange(M)


def get_z(*, pc: bool):
    return _get_z(pc)


@cache
def _get_c(pc: bool):
    M = get_m()
    Z = get_z(pc=pc)
    W = get_w()
    return hsv_to_rgb(np.array([3.0 * Z / M % 1, 0.4 + Z * 0, 0.6 + 0.15 * (Z / M)]).transpose())


def get_c(*, pc: bool):
    return _get_c(pc)


def render_cors(*, px, pc):
    M = get_m()
    N = get_n()
    Y = get_y()
    X = get_x(px=px)
    assert Y.shape == (M, N)
    plt.clf()
    C = get_c(pc=pc)
    plt.stackplot(
        X,
        Y,
        colors=C,
        aa=False,
        linewidth=0.0,
    )
    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)