# n_qubit_combos = [2,4,6,8,10]
# n_stratum_combos = [2,4,6,8,10]

# file_name = './{n_qubits}q_{n_stratums}s.json'
# # Total 25 data variations.
# # I want a 5 by 5 plot where the top row is 2 qubits [2,4,6,8,10] stratums (5 cols), second row is 4 qubits (5 cols) and so forth
# # In the first 5 by 5 plot make the x-axis simply the len of avg_loss array and call it epochs since each epcoh shows 1 loss value
# # In the second 5 by 5 plot make the x-axis the evals_range array and call it forward passes (on QML one epoch can require different number of forward passes as gradients are coputed differently)
# # The first line can be called DFS QNG and uses the dfs_qng data the second line is just called QNG
# # For both lines use the std_loss as the standard deviation around the training lines as they were ran 5 times in total

# # Below is the data strucutre of each file. Start with reading each file then creating the two 5 by 5 plots. 
# # Make lines thick and visible 
# # I dont want you to plot this I want to plot it on my own computer. No need for fancy styling minimal code is better

# {
#   "exact_gse": 0,
#   "config": {
#     "n_qubits": 2,
#     "n_stratums": 2,
#     "n_epochs": 30,
#     "step_size": 0.1,
#     "n_runs": 5
#   },
#   "dfs_qng": {
#     "avg_loss": [],
#     "std_loss": [],
#     "evals_range": []
#   },
#   "standard_qng": {
#     "avg_loss": [],
#     "std_loss": [],
#     "evals_range": []
#   }
# }

# NOTE For the egt data the data structure looks like what is attached below.
# I.e. it uses the field exact_qng for avg / std loss and it does not have the evals range array
# {
#   "exact_gse": ,
#   "config": {
#     "n_qubits": ,
#     "n_epochs": ,
#     "step_size": ,
#     "n_runs": 
#   },
#   "exact_qng": {
#     "avg_loss": [],
#     "std_loss": []
#   }
# }






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

n_qubits = [2, 4, 6, 8, 10]
n_stratums = [2, 4, 6, 8, 10]

data = {}
for nq in n_qubits:
    for ns in n_stratums:
        with open(f'./VQE_DATA/{nq}q_{ns}s.json') as f:
            data[(nq, ns)] = json.load(f)

egt_data = {}
for nq in n_qubits:
    for ns in n_stratums:
        with open(f'./EGT_DATA/{nq}q.json') as f:
            egt_data[(nq, ns)] = json.load(f)

# Load Adam data strictly for 2 qubits
adam_data = {}
for nq in n_qubits:
    for ns in n_stratums:
        if True:
            try:
                with open(f'./VQE_DATA/adam01_{nq}q_{ns}s.json') as f:
                    adam_data[(nq, ns)] = json.load(f)
            except FileNotFoundError:
                pass

def make_plot(x_key, x_label):
    fig, axes = plt.subplots(5, 5, figsize=(18, 16))
    fig.subplots_adjust(hspace=0.3, wspace=0.35, top=0.93, left=0.07)

    # Column titles
    for j, ns in enumerate(n_stratums):
        axes[0][j].set_title(f'{ns} Stratums', fontsize=12, fontweight='bold')

    # Row labels
    for i, nq in enumerate(n_qubits):
        axes[i][0].annotate(f'{nq}Q', xy=(0, 0.5), xytext=(-45, 0),
            xycoords='axes fraction', textcoords='offset points',
            fontsize=12, fontweight='bold', ha='center', va='center', rotation=90)

    for i, nq in enumerate(n_qubits):
        for j, ns in enumerate(n_stratums):
            ax = axes[i][j]
            d = data[(nq, ns)]
            gse = d['exact_gse']

            # 1. Plot the standard and DFS QNG lines
            for key, label, color in [('dfs_qng', 'DFS QNG', 'tab:blue'), ('standard_qng', 'QNG', 'tab:orange')]:
                avg = np.array(d[key]['avg_loss'])
                std = np.array(d[key]['std_loss'])
                x = np.arange(len(avg)) if x_key == 'epochs' else np.array(d[key]['evals_range'])
                ax.plot(x, avg, label=label, color=color, linewidth=2.5)
                ax.fill_between(x, avg - std, avg + std, alpha=0.2, color=color)

            # 2. Add EGT QNG line to both plots
            egt_d = egt_data[(nq, ns)]['exact_qng']
            avg_egt = np.array(egt_d['avg_loss'])
            std_egt = np.array(egt_d['std_loss'])
            
            if x_key == 'epochs':
                x_egt = np.arange(len(avg_egt))
            else:
                # Calculate cumulative forward passes: cost per epoch * epoch number
                cost_per_epoch = 2 * ((2 ** nq) - 1)
                x_egt = np.arange(1, len(avg_egt) + 1) * cost_per_epoch

            ax.plot(x_egt, avg_egt, label='EGT QNG', color='tab:green', linewidth=2.5)
            ax.fill_between(x_egt, avg_egt - std_egt, avg_egt + std_egt, alpha=0.2, color='tab:green')

            # 3. Add Adam line ONLY if nq == 2 and data exists
            if (nq, ns) in adam_data:
                ad = adam_data[(nq, ns)].get('adam', adam_data[(nq, ns)])
                if 'avg_loss' in ad:
                    avg_adam = np.array(ad['avg_loss'])
                    std_adam = np.array(ad.get('std_loss', np.zeros_like(avg_adam)))
                    
                    if x_key == 'epochs':
                        # Truncate Adam data to 30 epochs here
                        avg_adam = avg_adam[:30]
                        std_adam = std_adam[:30]
                        x_adam = np.arange(len(avg_adam))
                    else:
                        x_adam = np.array(ad['evals_range'])

                    ax.plot(x_adam, avg_adam, label='Adam', color='tab:red', linewidth=2.5)
                    ax.fill_between(x_adam, avg_adam - std_adam, avg_adam + std_adam, alpha=0.2, color='tab:red')

            ax.axhline(y=gse, color='black', linestyle='--', linewidth=1.5, label='Exact GSE')
            ax.tick_params(labelsize=9)
            if i == 4:
                ax.set_xlabel(x_label, fontsize=10)
            if i == 0 and j == 0:
                ax.legend(fontsize=8)
    return fig

fig1 = make_plot('epochs', 'Epochs')
fig2 = make_plot('evals', 'Forward Passes')
plt.show()













############################## TABLE
import json
import numpy as np

n_qubits = [2, 4, 6, 8, 10]
n_stratums = [2, 4, 6, 8, 10]

data = {}
for nq in n_qubits:
    for ns in n_stratums:
        with open(f'./VQE_DATA/{nq}q_{ns}s.json') as f:
            data[(nq, ns)] = json.load(f)

egt_data = {}
for nq in n_qubits:
    for ns in n_stratums:
        with open(f'./EGT_DATA/{nq}q.json') as f:
            egt_data[(nq, ns)] = json.load(f)

# Load Adam data using the updated path structure
adam_data = {}
for nq in n_qubits:
    for ns in n_stratums:
        try:
            with open(f'./VQE_DATA/adam01_{nq}q_{ns}s.json') as f:
                adam_data[(nq, ns)] = json.load(f)
        except FileNotFoundError:
            pass

# Compute regret for each (qubit, strata, optimizer)
results = {}
budget = 30  # Fixed epoch budget

for nq in n_qubits:
    for ns in n_stratums:
        d = data[(nq, ns)]
        gse = d['exact_gse']

        # DFS-QNG (truncated to budget)
        dfs_avg = np.array(d['dfs_qng']['avg_loss'])[:budget]
        dfs_std = np.array(d['dfs_qng']['std_loss'])[:budget]
        x_dfs = np.arange(len(dfs_avg))
        dfs_regret = np.trapezoid(dfs_avg - gse, x=x_dfs) / budget
        dfs_regret_std = np.trapezoid(dfs_std, x=x_dfs) / budget

        # Standard QNG (truncated to budget)
        qng_avg = np.array(d['standard_qng']['avg_loss'])[:budget]
        qng_std = np.array(d['standard_qng']['std_loss'])[:budget]
        x_qng = np.arange(len(qng_avg))
        qng_regret = np.trapezoid(qng_avg - gse, x=x_qng) / budget
        qng_regret_std = np.trapezoid(qng_std, x=x_qng) / budget

        # EGT (truncated to budget)
        egt_d = egt_data[(nq, ns)]['exact_qng']
        egt_avg = np.array(egt_d['avg_loss'])[:budget]
        egt_std = np.array(egt_d['std_loss'])[:budget]
        x_egt = np.arange(len(egt_avg))
        
        if len(egt_avg) > 1:
            egt_regret = np.trapezoid(egt_avg - gse, x=x_egt) / budget
            egt_regret_std = np.trapezoid(egt_std, x=x_egt) / budget
        else:
            egt_regret = None
            egt_regret_std = None

        # Store baselines
        results[(nq, ns)] = {
            'DFS-QNG': (dfs_regret, dfs_regret_std),
            'QNG':     (qng_regret, qng_regret_std),
            'EGT':     (egt_regret, egt_regret_std),
        }

        # Adam (truncated to budget)
        if (nq, ns) in adam_data:
            ad = adam_data[(nq, ns)].get('adam', adam_data[(nq, ns)])
            if 'avg_loss' in ad:
                adam_avg = np.array(ad['avg_loss'])[:budget]
                adam_std = np.array(ad.get('std_loss', np.zeros_like(adam_avg)))[:budget]
                x_adam = np.arange(len(adam_avg))
                
                if len(adam_avg) > 1:
                    adam_regret = np.trapezoid(adam_avg - gse, x=x_adam) / budget
                    adam_regret_std = np.trapezoid(adam_std, x=x_adam) / budget
                    results[(nq, ns)]['Adam'] = (adam_regret, adam_regret_std)

optimizers = ['DFS-QNG', 'QNG', 'Adam', 'EGT']
n_opt = len(optimizers)

def bold_winner(nq, ns):
    best_opt = None
    best_val = float('inf')
    for opt in optimizers:
        if opt in results.get((nq, ns), {}):
            val = results[(nq, ns)][opt][0]
            if val is not None and val < best_val:
                best_val = val
                best_opt = opt
    return best_opt

def fmt_cell(nq, ns, opt):
    if opt not in results.get((nq, ns), {}) or results[(nq, ns)].get(opt, (None,))[0] is None:
        return '---'
    mean, std = results[(nq, ns)][opt]
    winner = bold_winner(nq, ns)
    if opt == winner:
        return f'\\textbf{{{mean:.2f}}}{{\\tiny$\\pm${std:.2f}}}'
    return f'{mean:.2f}{{\\tiny$\\pm${std:.2f}}}'

# Print LaTeX
print(r"""\begin{table}[t]
\centering
\caption{Normalized Area Above Optimal ($\downarrow$) on VQE Ising Model. Budget equals 30 epochs. Best per cell in \textbf{bold}.}
\label{tab:vqe_regret}
\footnotesize
\begin{tabular}{cl|ccccc}
\toprule
$n$ & Optimizer & 2 Strata & 4 Strata & 6 Strata & 8 Strata & 10 Strata \\
\midrule""")

for i, nq in enumerate(n_qubits):
    for j, opt in enumerate(optimizers):
        if j == 0:
            qubit_label = f'\\multirow{{{n_opt}}}{{*}}{{{nq}}}'
        else:
            qubit_label = ''
        cells = [fmt_cell(nq, ns, opt) for ns in n_stratums]
        row = f'{qubit_label} & {opt} & {" & ".join(cells)} \\\\'
        print(row)
    if i < len(n_qubits) - 1:
        print(r'\midrule')

print(r"""\bottomrule
\end{tabular}
\end{table}""")