# noise_levels = [0.001, 0.01, 0.1]

# for noise_level in noise_levels:
#     path = f'./Data/noise_qng({noise_level}).json'
#     Read the JSON file. The file when read has the format below
#     So we ran QNG training experiments at 3 different noise noise_levels
#     at each noise level i ran qng for 1, 5 and 10 layers. 
#     the layers_i object is an array with 5 training curves. 
#     These 5 training curves we should take the mean for and also add the STD
#     I want a 3x3 plot where first row is noise level 0.001, second row is noise leve 0.01 and so forth
#     column 1 is for 1 layer, column 2 is for 2 layers and so forth. Please plot this for me
#     below is the strucutre the data will have when reading the paths.

# data = {
#     "layers_1": [
#         [],
#         [],
#         [],
#         [],
#         []
#     ],
#     "layers_5": [
#         [],
#         [],
#         [],
#         [],
#         []
#     ],
#     "layers_10": [
#         [],
#         [],
#         [],
#         [],
#         []
#     ]
# }

import json
import numpy as np
import matplotlib.pyplot as plt
import matplotlib

matplotlib.rcParams.update({
    "font.size": 12,
    "axes.labelsize": 14,
    "axes.titlesize": 15,
    "xtick.labelsize": 12,
    "ytick.labelsize": 12,
    "legend.fontsize": 12,
    "figure.dpi": 150,
})

noise_levels = [0.001, 0.01, 0.1]
layer_counts = [1, 5, 10]

fig, axes = plt.subplots(3, 3, figsize=(10, 8), sharey="row", layout="constrained")

for i, noise_level in enumerate(noise_levels):
    path = f"./Data/noise_qng({noise_level}).json"
    path_2 = f"./Data/regular_qng({noise_level}).json"

    with open(path, "r") as f:
        data = json.load(f)
    with open(path_2, "r") as f:
        data_2 = json.load(f)

    for j, n_layers in enumerate(layer_counts):
        ax = axes[i][j]
        key = f"layers_{n_layers}"

        # Noise QNG
        curves = np.array(data[key])
        mean = curves.mean(axis=0)
        std = curves.std(axis=0)
        steps = np.arange(len(mean))
        ax.plot(steps, mean, color="#1f77b4", linewidth=2.5, linestyle="--", label="Noise QNG")
        ax.fill_between(steps, mean - std, mean + std, alpha=0.15, color="#1f77b4")

        # Regular QNG
        curves_2 = np.array(data_2[key])
        mean_2 = curves_2.mean(axis=0)
        std_2 = curves_2.std(axis=0)
        steps_2 = np.arange(len(mean_2))
        ax.plot(steps_2, mean_2, color="#d62728", linewidth=2.5, label="Regular QNG")
        ax.fill_between(steps_2, mean_2 - std_2, mean_2 + std_2, alpha=0.15, color="#d62728")

        ax.grid(True, alpha=0.5, linewidth=0.8)

        # Column titles (top row only)
        if i == 0:
            ax.set_title(f"{n_layers} Layer{'s' if n_layers > 1 else ''}")

        # Row labels (first column only)
        if j == 0:
            ax.set_ylabel(f"$q = {noise_level}$")

        # Legend on top-left subplot, upper right
        if i == 0 and j == 0:
            ax.legend(frameon=True, fancybox=False, edgecolor="black",
                      framealpha=1.0, loc="upper right")

fig.supxlabel("Epochs", fontsize=14, x=0.57)
fig.supylabel("Energy", fontsize=14, y=0.53)

plt.show()