import json
import math

import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter, LogLocator
losses = []
with open("models/losses.json", "r") as f:
    losses += json.load(f)

print(losses[8])

avglosses = []
iters = []

chunksize = 1
iter = 100

while iter + chunksize < len(losses):
    avg = 0
    for i in range(round(chunksize)):
        avg += losses[iter + i]
    avg = avg / round(chunksize)

    avglosses.append(avg)      # riktig loss, inte log(avg)
    iters.append(iter)         # riktigt iterationsnummer

    iter += round(chunksize)
    chunksize *= 1.01
    print(iter, avg)




fig, ax = plt.subplots()
ax.plot(iters, avglosses)
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlabel("Iteration")
ax.set_ylabel("Träningsloss")

import numpy as np

# Sätt egna y-ticks baserat på datans faktiska min/max
ymin, ymax = min(avglosses), max(avglosses)
yticks = np.arange(math.floor(ymin*10)/10, math.ceil(ymax*10)/10 + 0.1, 0.1)
ax.set_yticks(yticks)
ax.yaxis.set_major_formatter(ScalarFormatter())
ax.yaxis.set_minor_formatter(plt.NullFormatter())  # slipp rörig minor-text
# Grid för både major och minor
ax.grid(True, which='major', linestyle='-', linewidth=0.7, alpha=0.7)
ax.grid(True, which='minor', linestyle=':', linewidth=0.5, alpha=0.4)



plt.tight_layout()
plt.title("Träningsloss för microbatch")
plt.show()
plt.close()