import pennylane as qml
from pennylane import numpy as np
import matplotlib.pyplot as plt
import networkx as nx
import os
import json


n_qubits = 10   
p_layers = [1,2,3,4,5] 

dev = qml.device("lightning.gpu", wires=n_qubits+1)

for p in p_layers:
    # --- 3-REGULAR GRAPH GENERATION ---
    graph = nx.random_regular_graph(3, n_qubits, seed=42)
    edges = list(graph.edges)

    n_edges = len(edges)
    n_nodes = n_qubits

    total_params = p * (n_edges + n_nodes)

    # Hamiltonian Construction
    obs_cost = [qml.PauliZ(i) @ qml.PauliZ(j) for i, j in edges]
    coeffs_cost = [1.0] * n_edges
    H_cost = qml.Hamiltonian(coeffs_cost, obs_cost)

    # Compute EXACT Ground State Energy (GSE)
    H_matrix = qml.matrix(H_cost)
    exact_gse = np.linalg.eigvalsh(H_matrix)[0]


    stratum_slices = []
    current_idx = 0
    for _ in range(p):
        stratum_slices.append(slice(current_idx, current_idx + n_edges))
        current_idx += n_edges
        stratum_slices.append(slice(current_idx, current_idx + n_nodes))
        current_idx += n_nodes

    n_stratums = len(stratum_slices)


    @qml.qnode(dev, interface="autograd")
    def ma_qaoa_circuit(flat_params):
        for i in range(n_qubits):
            qml.Hadamard(wires=i)
        idx = 0
        for layer in range(p):
            for i, j in edges:
                qml.IsingZZ(2 * flat_params[idx], wires=[i, j])
                idx += 1
            for n_idx in range(n_nodes):
                qml.RX(2 * flat_params[idx], wires=n_idx)
                idx += 1
        return qml.expval(H_cost)


    def dfs_gqng_step_maqaoa(qnode, flat_params, stepsize=0.01, epsilon=1e-8):
        new_params = np.copy(flat_params)
        # Optimization: Calculate metric tensor once per step
        g_full = qml.metric_tensor(qnode)(new_params)
        for idx_slice in stratum_slices:
            grad = qml.grad(qnode)(new_params)
            grad_sub = grad[idx_slice]
            g_sub = g_full[idx_slice, idx_slice]

            sub_size = g_sub.shape[0]
            g_sub_reg = g_sub + epsilon * np.eye(sub_size)
            g_sub_inv = np.linalg.inv(g_sub_reg)

            new_params[idx_slice] = new_params[idx_slice] - stepsize * np.dot(g_sub_inv, grad_sub)
        return new_params

    # Experiment Setup
    epochs = 30
    a_epochs = 100
    step_size = 0.1 
    n_runs = 5

    all_loss_dfs = np.zeros((n_runs, epochs))
    all_loss_qng = np.zeros((n_runs, epochs))
    all_loss_adam = np.zeros((n_runs, a_epochs)) 

    evals_per_epoch_adam = 2 * total_params

    print(f"--- Configuration: {n_qubits} Qubits, p={p} ---")

    for run in range(n_runs):
        print(f"Run {run + 1}/{n_runs}...")
        initial_array = np.random.uniform(0, 2 * np.pi, total_params)
        params_dfs = np.array(initial_array, requires_grad=True)
        params_qng = np.array(initial_array, requires_grad=True)
        params_adam = np.array(initial_array, requires_grad=True)

        # DFS-GQNG Loop
        for epoch in range(epochs):
            all_loss_dfs[run, epoch] = float(ma_qaoa_circuit(params_dfs))
            params_dfs = dfs_gqng_step_maqaoa(ma_qaoa_circuit, params_dfs, stepsize=step_size)

        # Standard QNG Loop
        opt = qml.QNGOptimizer(stepsize=step_size)
        for epoch in range(epochs):
            all_loss_qng[run, epoch] = float(ma_qaoa_circuit(params_qng))
            params_qng, _ = opt.step_and_cost(ma_qaoa_circuit, params_qng)

        # Adam Loop
        opt_adam = qml.AdamOptimizer(stepsize=step_size)
        for epoch in range(a_epochs):
            all_loss_adam[run, epoch] = float(ma_qaoa_circuit(params_adam))
            params_adam, _ = opt_adam.step_and_cost(ma_qaoa_circuit, params_adam)

    # Statistics
    avg_loss_dfs = np.mean(all_loss_dfs, axis=0)
    std_loss_dfs = np.std(all_loss_dfs, axis=0)
    avg_loss_qng = np.mean(all_loss_qng, axis=0)
    std_loss_qng = np.std(all_loss_qng, axis=0)
    avg_loss_adam = np.mean(all_loss_adam, axis=0)
    std_loss_adam = np.std(all_loss_adam, axis=0)

    # Cost Calculation
    evals_per_epoch_qng = (2 * total_params) + (2 * (total_params ** 2))
    evals_per_epoch_dfs = sum([(2 * (s.stop-s.start)) + (2 * (s.stop-s.start)**2) for s in stratum_slices])

    # Plotting
    plt.figure(figsize=(10, 6))
    plt.plot(range(epochs), avg_loss_dfs, label=f'DFS-GQNG ({epochs} eps)', color='purple')
    plt.plot(range(epochs), avg_loss_qng, label=f'Standard QNG ({epochs} eps)', color='orange')
    plt.plot(range(a_epochs), avg_loss_adam, label=f'Adam ({a_epochs} eps)', color='teal', linestyle='--')
    plt.axhline(y=exact_gse, color='black', linestyle=':', label='Exact GSE')
    plt.title(f"MaxCut 3-Regular: {n_qubits}q, p={p}")
    plt.xlabel("Epochs")
    plt.ylabel("Energy")
    plt.legend()
    plt.show()

    # Save Results 
    folder = "/content/drive/MyDrive/3MAX"
    os.makedirs(folder, exist_ok=True)
    filepath = os.path.join(folder, f"3maxreg_{n_qubits}q_{p}p_maQAOA.json")

    results = {
        "exact_gse": float(exact_gse),
        "config": {
            "n_qubits": n_qubits,
            "p_layers": p,
            "n_stratums": n_stratums,
            "n_runs": n_runs,
            "epochs_qng": epochs,
            "epochs_adam": a_epochs,
            "step_size": step_size
        },
        "dfs_qng": {
            "avg_loss": avg_loss_dfs.tolist(),
            "std_loss": std_loss_dfs.tolist(),
            "evals_per_epoch": evals_per_epoch_dfs
        },
        "standard_qng": {
            "avg_loss": avg_loss_qng.tolist(),
            "std_loss": std_loss_qng.tolist(),
            "evals_per_epoch": evals_per_epoch_qng
        },
        "adam": {
            "avg_loss": avg_loss_adam.tolist(),
            "std_loss": std_loss_adam.tolist(),
            "evals_per_epoch": evals_per_epoch_adam
        }
    }

    with open(filepath, "w") as f:
        json.dump(results, f, indent=2)
        print(f"Results for p={p} saved to {filepath}\n")