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

n_qubits = [4, 6, 8, 10]
n_stratums = [1,2,3,4,5]

data = {}
for nq in n_qubits:
    for ns in n_stratums:
        with open(f'./3MAX-DATA/3maxreg_{nq}q_{ns}p_maQAOA.json') as f:
            data[(nq, ns)] = json.load(f)

adamdata = {}
for nq in n_qubits:
    for ns in n_stratums:
        with open(f'./3MAX-DATA/3maxreg_{nq}q_{ns}p_maQAOA.json') as f:
            adamdata[(nq, ns)] = json.load(f)

def make_plot(x_key, x_label):
    fig, axes = plt.subplots(4, 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.arange(len(avg)) * d[key]["evals_per_epoch"]
                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. Plot the Adam line
            ad = adamdata[(nq, ns)]
            adam_key = 'adam'  # Adjust this key if it's named differently in your JSON
            avg_adam = np.array(ad[adam_key]['avg_loss'])[0:30]
            std_adam = np.array(ad[adam_key]['std_loss'])[0:30]
            x_adam = np.arange(len(avg_adam))[0:30] if x_key == 'epochs' else np.arange(len(avg)) * d[key]["evals_per_epoch"] #np.array(ad[adam_key]['evals_range'])
            ax.plot(x_adam, avg_adam, label='Adam', color='red', linewidth=2.5)
            ax.fill_between(x_adam, avg_adam - std_adam, avg_adam + std_adam, alpha=0.2, color='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 = [4, 6, 8, 10]
n_stratums = [1, 2, 3, 4, 5]

data = {}
for nq in n_qubits:
    for ns in n_stratums:
        with open(f'./3MAX-DATA/3maxreg_{nq}q_{ns}p_maQAOA.json') as f:
            data[(nq, ns)] = json.load(f)

adamdata = {}
for nq in n_qubits:
    for ns in n_stratums:
        with open(f'./3MAX-DATA/3maxreg_{nq}q_{ns}p_maQAOA.json') as f:
            adamdata[(nq, ns)] = json.load(f)

# Compute regret for each (qubit, strata, optimizer)
results = {}
for nq in n_qubits:
    for ns in n_stratums:
        d = data[(nq, ns)]
        ad = adamdata[(nq, ns)]
        gse = d['exact_gse']

        dfs_avg = np.array(d['dfs_qng']['avg_loss'])
        qng_avg = np.array(d['standard_qng']['avg_loss'])
        adam_avg = np.array(ad['adam']['avg_loss'])[0:30]

        dfs_std_arr = np.array(d['dfs_qng']['std_loss'])
        qng_std_arr = np.array(d['standard_qng']['std_loss'])
        adam_std_arr = np.array(ad['adam']['std_loss'])[0:30]

        # Since all run for 30 epochs, they have the same array lengths.
        # We integrate over the epochs directly.
        epochs = np.arange(len(dfs_avg))
        budget = epochs[-1]

        dfs_regret = np.trapezoid(dfs_avg - gse, x=epochs) / budget
        qng_regret = np.trapezoid(qng_avg - gse, x=epochs) / budget
        adam_regret = np.trapezoid(adam_avg - gse, x=epochs) / budget
        
        dfs_regret_std = np.trapezoid(dfs_std_arr, x=epochs) / budget
        qng_regret_std = np.trapezoid(qng_std_arr, x=epochs) / budget
        adam_regret_std = np.trapezoid(adam_std_arr, x=epochs) / budget

        results[(nq, ns)] = {
            'DFS-QNG': (dfs_regret, dfs_regret_std),
            'QNG':     (qng_regret, qng_regret_std),
            'Adam':    (adam_regret, adam_regret_std),
        }

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

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

def bold_winner(nq, ns):
    """Return the optimizer name with the lowest mean for this cell."""
    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 < best_val:
                best_val = val
                best_opt = opt
    return best_opt

def fmt_cell_bold(nq, ns, opt):
    if opt not in results.get((nq, ns), {}):
        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 QAOA MaxCut3. Budget equals total epochs. Best per cell in \textbf{bold}.}
\label{tab:qaoa_regret}
\resizebox{\textwidth}{!}{%
\begin{tabular}{cl|ccccc}
\toprule
$n$ & Optimizer & 1 Stratum & 2 Strata & 3 Strata & 4 Strata & 5 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_bold(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}""")