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

# ==========================================
# Hardware & ma-QAOA Setup
# ==========================================
n_qubits = 2
p_layers = 2  # Number of alternating QAOA layers

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

# Define MaxCut Ring Graph
edges = [(i, (i + 1) % n_qubits) for i in range(n_qubits)]
n_edges = len(edges)
n_nodes = n_qubits

# In ma-QAOA: Each layer has `n_edges` parameters for Cost, and `n_nodes` for Mixer
total_params = p_layers * (n_edges + n_nodes)

# Minimizing \sum Z_i Z_j maximizes the cut
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]

# ==========================================
# Define the Strata (Blocks of mutually commuting generators)
# ==========================================
stratum_slices = []
current_idx = 0
for _ in range(p_layers):
    # Cost Stratum: All ZZ gates commute
    stratum_slices.append(slice(current_idx, current_idx + n_edges))
    current_idx += n_edges
    
    # Mixer Stratum: All RX gates commute
    stratum_slices.append(slice(current_idx, current_idx + n_nodes))
    current_idx += n_nodes

n_stratums = len(stratum_slices)

# ==========================================
# Multi-Angle QAOA Circuit
# ==========================================
@qml.qnode(dev, interface="autograd")
def ma_qaoa_circuit(flat_params):
    # 1. Initial superposition state
    for i in range(n_qubits):
        qml.Hadamard(wires=i)
        
    idx = 0
    # 2. Alternating ma-QAOA layers
    for layer in range(p_layers):
        
        # Cost Layer: Each edge gets its own parameter
        for e_idx, (i, j) in enumerate(edges):
            qml.IsingZZ(2 * flat_params[idx], wires=[i, j])
            idx += 1
            
        # Mixer Layer: Each node gets its own parameter
        for n_idx in range(n_nodes):
            qml.RX(2 * flat_params[idx], wires=n_idx)
            idx += 1
            
    return qml.expval(H_cost)

# ==========================================
# DFS-GQNG Step (Exact QNG)
# ==========================================
def dfs_gqng_step_maqaoa(qnode, flat_params, stepsize=0.01, epsilon=1e-8):
    new_params = np.copy(flat_params)
    
    # Iterate through our pre-defined commuting strata
    for idx_slice in stratum_slices:
        grad = qml.grad(qnode)(new_params)
        g = qml.metric_tensor(qnode)(new_params)
        
        # Isolate the block for the active stratum
        grad_sub = grad[idx_slice]
        g_sub = g[idx_slice, idx_slice]
        
        # Regularize and invert the block
        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)
        
        # Update active parameters while keeping others frozen
        new_params[idx_slice] = new_params[idx_slice] - stepsize * np.dot(g_sub_inv, grad_sub)
        
    return new_params

# ==========================================
# Experimental Setup
# ==========================================
epochs = 30
step_size = 0.1
n_runs = 5

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

print(f"Exact GSE calculated as: {exact_gse:.4f}")
print(f"Total Parameters: {total_params}")
print(f"Total Strata: {n_stratums} (Sizes: {[s.stop - s.start for s in stratum_slices]})")
print(f"Starting ma-QAOA experiment over {n_runs} runs...\n")

for run in range(n_runs):
    print(f"--- Run {run + 1}/{n_runs} ---")
    
    # Initialize random parameters (Fair starting point for both)
    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)

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

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

# ==========================================
# Statistical Aggregation
# ==========================================
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)

# ==========================================
# Cost Calculations for Plot 2
# ==========================================
# Standard QNG Cost: 2P (grad) + 2P^2 (metric tensor)
evals_per_epoch_qng = (2 * total_params) + (2 * (total_params ** 2))

# DFS-GQNG Cost: Sum over all strata: [2S (grad) + 2S^2 (metric tensor)]
evals_per_epoch_dfs = 0
for s in stratum_slices:
    size = s.stop - s.start
    evals_per_epoch_dfs += (2 * size) + (2 * (size ** 2))

evals_range_qng = [epoch * evals_per_epoch_qng for epoch in range(epochs)]
evals_range_dfs = [epoch * evals_per_epoch_dfs for epoch in range(epochs)]

print(f"\nStandard QNG Circuit Evals per Epoch: {evals_per_epoch_qng}")
print(f"DFS-GQNG Circuit Evals per Epoch: {evals_per_epoch_dfs}")

# ==========================================
# Plotting 1: Convergence vs Epochs
# ==========================================
plt.figure(figsize=(9, 6))
epochs_range = range(epochs)

plt.plot(epochs_range, avg_loss_dfs, marker='o', color='purple', label='DFS-GQNG (Exact QNG)')
plt.fill_between(epochs_range, avg_loss_dfs - std_loss_dfs, avg_loss_dfs + std_loss_dfs, color='purple', alpha=0.2)

plt.plot(epochs_range, avg_loss_qng, marker='^', color='orange', label='Standard QNG')
plt.fill_between(epochs_range, avg_loss_qng - std_loss_qng, avg_loss_qng + std_loss_qng, color='orange', alpha=0.2)

plt.axhline(y=exact_gse, color='black', linestyle=':', label=f'Exact GSE ({exact_gse:.4f})')
plt.title(f"Optimizer Comparison: DFS-GQNG vs Standard QNG (ma-QAOA)\n(Average over {n_runs} runs)")
plt.xlabel("Epoch")
plt.ylabel("Expected Energy")
plt.legend()
plt.grid(True)
plt.show()

# ==========================================
# Plotting 2: Convergence vs Circuit Evaluations
# ==========================================
plt.figure(figsize=(9, 6))

plt.plot(evals_range_dfs, avg_loss_dfs, marker='o', color='purple', label='DFS-GQNG (Exact QNG)')
plt.fill_between(evals_range_dfs, avg_loss_dfs - std_loss_dfs, avg_loss_dfs + std_loss_dfs, color='purple', alpha=0.2)

plt.plot(evals_range_qng, avg_loss_qng, marker='^', color='orange', label='Standard QNG')
plt.fill_between(evals_range_qng, avg_loss_qng - std_loss_qng, avg_loss_qng + std_loss_qng, color='orange', alpha=0.2)

plt.axhline(y=exact_gse, color='black', linestyle=':', label=f'Exact GSE ({exact_gse:.4f})')
plt.title(f"Computational Cost Comparison: DFS-GQNG vs Standard QNG (ma-QAOA)\n(Average over {n_runs} runs)")
plt.xlabel("Circuit Evaluations (Hardware Cost)")
plt.ylabel("Expected Energy")
plt.legend()
plt.grid(True)
plt.show()

# ==========================================
# Save Results
# ==========================================
folder = "./dfsvqe-maQAOA"
os.makedirs(folder, exist_ok=True)

filename = f"{n_qubits}q_{p_layers}p_maQAOA.json"
filepath = os.path.join(folder, filename)

results = {
    "exact_gse": float(exact_gse),
    "config": {
        "n_qubits": n_qubits,
        "p_layers": p_layers,
        "n_stratums": n_stratums,
        "n_epochs": epochs,
        "step_size": step_size,
        "n_runs": n_runs,
    },
    "dfs_qng": {
        "avg_loss": avg_loss_dfs.tolist(),
        "std_loss": std_loss_dfs.tolist(),
        "evals_range": evals_range_dfs,
    },
    "standard_qng": {
        "avg_loss": avg_loss_qng.tolist(),
        "std_loss": std_loss_qng.tolist(),
        "evals_range": evals_range_qng,
    },
}

with open(filepath, "w") as f:
    json.dump(results, f, indent=2)

print(f"Results saved to {filepath}")