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

# Define hardware and strata
n_qubits = 2
n_stratums = 2

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


# Define the Hamiltonian: Z_i Z_{i+1} + X_i
coeffs = [-1.0] * n_qubits + [-1.0] * n_qubits
obs = [qml.PauliZ(i) @ qml.PauliZ((i + 1) % n_qubits) for i in range(n_qubits)] + \
      [qml.PauliX(i) for i in range(n_qubits)]
H = qml.Hamiltonian(coeffs, obs)

# Compute the EXACT Ground State Energy (GSE) for the current Hamiltonian
H_matrix = qml.matrix(H)
eigenvalues = np.linalg.eigvalsh(H_matrix)
exact_gse = eigenvalues[0]

@qml.qnode(dev, interface="autograd")
def qnode(flat_params):
    # flat_params is a 1D array of size (n_stratums * n_qubits)
    for s in range(n_stratums):
        for q in range(n_qubits):
            # Calculate the exact flat index for this stratum and qubit
            idx = s * n_qubits + q

            # Even stratum = RZ, Odd stratum = RX
            if s % 2 == 0:
                qml.RZ(flat_params[idx], wires=q)
            else:
                qml.RX(flat_params[idx], wires=q)

        # Apply linear CNOTs
        for q in range(n_qubits - 1):
            qml.CNOT(wires=[q, q + 1])

    return qml.expval(H)

def dfs_gqng_step(qnode, flat_params, stepsize=0.01, epsilon=1e-8):
    new_params = np.copy(flat_params)
    for s in range(n_stratums):
        start_idx = s * n_qubits
        end_idx = start_idx + n_qubits
        idx = slice(start_idx, end_idx)

        grad = qml.grad(qnode)(new_params)
        g = qml.metric_tensor(qnode)(new_params)

        grad_sub = grad[idx]
        g_sub = g[idx, idx]

        g_sub_reg = g_sub + epsilon * np.eye(n_qubits)
        g_sub_inv = np.linalg.inv(g_sub_reg)

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

# ==========================================
# Experimental Setup
# ==========================================
epochs = 30
step_size = 0.1
n_runs = 5
total_params = n_stratums * n_qubits

# Matrices to store loss histories for all runs
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"Starting rigorous experiment over {n_runs} runs...\n")

for run in range(n_runs):
    print(f"--- Run {run + 1}/{n_runs} ---")
    # Initialize random parameters for this specific run (No fixed seed!)
    # Both optimizers start from the EXACT same random location for a fair comparison
    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)

    # --- Exact QNG (DFS-GQNG) Optimization Loop ---
    for epoch in range(epochs):
        loss = qnode(params_dfs)
        all_loss_dfs[run, epoch] = loss
        params_dfs = dfs_gqng_step(qnode, params_dfs, stepsize=step_size, epsilon=1e-8)

    # --- Standard QNG Optimization Loop ---
    opt = qml.QNGOptimizer(stepsize=step_size)
    for epoch in range(epochs):
        loss = qnode(params_qng)
        all_loss_qng[run, epoch] = loss
        params_qng, _ = opt.step_and_cost(qnode, 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: M * [2S (grad) + 2S^2 (metric tensor)]
# Here, Stratum size S = n_qubits
evals_per_epoch_dfs = n_stratums * ((2 * n_qubits) + (2 * (n_qubits ** 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"Standard 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\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\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 ########################
import json
import os

folder = "./dfsvqe"
os.makedirs(folder, exist_ok=True)

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

results = {
    "exact_gse": float(exact_gse),
    "config": {
        "n_qubits": n_qubits,
        "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}")