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

# Suppress Pennylane's DepolarizingChannel warnings
warnings.filterwarnings("ignore", message="Operator DepolarizingChannel does not define a decomposition")

# --- 1. Circuit Definitions (1D Params for Correct Stratum Slicing) ---
def build_qnodes(n_qubits, n_layers, noise_level=0.1):
    dev_clean = qml.device("default.qubit", wires=n_qubits+1)
    dev_noisy = qml.device("default.mixed", wires=n_qubits+1)

    coeffs = []
    obs = []
    for i in range(n_qubits):
        coeffs.extend([-1.0, -1.0])
        obs.append(qml.PauliZ(i) @ qml.PauliZ((i + 1) % n_qubits))
        obs.append(qml.PauliX(i))
    H = qml.Hamiltonian(coeffs, obs)

    def circuit_ansatz(flat_params, apply_noise=False):
        # We flatten the parameters to ensure slicing cleanly isolates 
        # mutually commuting gates (Stratum 1: RX layer 1, Stratum 2: RZ layer 1...)
        idx = 0
        for l in range(n_layers):
            # RX Stratum
            for i in range(n_qubits):
                qml.RX(flat_params[idx], wires=i)
                if apply_noise and noise_level > 0:
                    qml.DepolarizingChannel(noise_level, wires=i)
                idx += 1
            # RZ Stratum
            for i in range(n_qubits):
                qml.RZ(flat_params[idx], wires=i)
                if apply_noise and noise_level > 0:
                    qml.DepolarizingChannel(noise_level, wires=i)
                idx += 1
            if n_qubits > 1:
                for i in range(n_qubits):
                    qml.CNOT(wires=[i, (i + 1) % n_qubits])
        return qml.expval(H)

    # We need a clean QNode for the Metric Tensor to avoid Pennylane failures
    @qml.qnode(dev_clean, interface="autograd")
    def qnode_clean(flat_params):
        return circuit_ansatz(flat_params, apply_noise=False)

    # We use a noisy QNode for evaluating the physical loss and gradients
    @qml.qnode(dev_noisy, interface="autograd")
    def qnode_noisy(flat_params):
        return circuit_ansatz(flat_params, apply_noise=True)

    return qnode_clean, qnode_noisy

# --- 2. DFS-GQNG Implementation ---
def dfs_gqng_step(qnode_clean, qnode_noisy, flat_params, n_qubits, n_layers, stepsize=0.01, epsilon=1e-8):
    new_params = np.copy(flat_params)
    
    # Evaluate noisy gradients (Hardware Simulation)
    grad = qml.grad(qnode_noisy)(new_params)
    
    # Evaluate clean metric tensor (Geometric Structure invariance)
    g = qml.metric_tensor(qnode_clean)(new_params)

    # 2 strata per layer (1 RX stratum and 1 RZ stratum)
    n_stratums = 2 * n_layers
    
    for s in range(n_stratums):
        # Each stratum has exactly `n_qubits` commuting parameters
        start_idx = s * n_qubits
        end_idx = start_idx + n_qubits
        idx = slice(start_idx, end_idx)

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

        # Regularize and invert block
        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

# --- 3. Training & Experiment Logic ---
def train_noisy_models(n_qubits, n_layers, n_epochs=50, noise_level=0.1):
    qnode_clean, qnode_noisy = build_qnodes(n_qubits, n_layers, noise_level)

    # Total parameters are 2 (RX, RZ) * n_layers * n_qubits
    total_params = 2 * n_layers * n_qubits
    init_params = np.random.uniform(0, 2 * np.pi, total_params)
    
    # Initialize BOTH optimizers at the exact same location
    params_std = np.array(init_params, requires_grad=True)
    params_dfs = np.array(init_params, requires_grad=True)

    std_qng = qml.QNGOptimizer(stepsize=0.01)

    std_losses = []
    dfs_losses = []

    for epoch in range(n_epochs):
        # DFS-GQNG Step
        loss_dfs = qnode_noisy(params_dfs)
        dfs_losses.append(loss_dfs)
        params_dfs = dfs_gqng_step(qnode_clean, qnode_noisy, params_dfs, n_qubits, n_layers, stepsize=0.01)

        # Standard QNG Step (Pennylane fallback handles clean metric tensor automatically)
        loss_std = qnode_noisy(params_std)
        std_losses.append(loss_std)
        params_std, _ = std_qng.step_and_cost(qnode_noisy, params_std)

    return std_losses, dfs_losses

def run_experiment_repeated(n_qubits, n_layers, n_epochs=50, n_runs=5, noise_level=0.1):
    all_std_runs = []
    all_dfs_runs = []

    for run in range(n_runs):
        print(f"  Starting Run {run+1}/{n_runs} for {n_layers} Layers...")
        std_loss, dfs_loss = train_noisy_models(n_qubits, n_layers, n_epochs, noise_level)
        all_std_runs.append(std_loss)
        all_dfs_runs.append(dfs_loss)

    return np.array(all_std_runs), np.array(all_dfs_runs)

# --- 4. Execution & Plotting ---
n_qubits = 4
n_epochs = 50
n_runs = 5
noise_level = 0.1
layer_configs = [1]

std_qng_logs = {}
dfs_qng_logs = {}

fig, axes = plt.subplots(1, len(layer_configs), figsize=(15, 5))

for idx, n_layers in enumerate(layer_configs):
    print(f"\nEvaluating {n_layers} Layers Configuration...")
    std_data, dfs_data = run_experiment_repeated(n_qubits, n_layers, n_epochs, n_runs, noise_level)

    config_label = f"layers_{n_layers}"
    std_qng_logs[config_label] = std_data.tolist()
    dfs_qng_logs[config_label] = dfs_data.tolist()

    # Calculate Mean and Standard Deviation for variance bands
    mean_std = np.mean(std_data, axis=0)
    std_std = np.std(std_data, axis=0)
    
    mean_dfs = np.mean(dfs_data, axis=0)
    std_dfs = np.std(dfs_data, axis=0)

    epochs = np.arange(n_epochs)

    # Plotting results for current configuration
    axes[idx].plot(epochs, mean_std, label=f'Standard QNG', color='orange')
    axes[idx].plot(epochs, mean_dfs, label=f'DFS-GQNG (Exact QNG)', color='purple', linestyle='--')

    # Shaded regions represent the variance
    axes[idx].fill_between(epochs, mean_std - std_std, mean_std + std_std, color='orange', alpha=0.2)
    axes[idx].fill_between(epochs, mean_dfs - std_dfs, mean_dfs + std_dfs, color='purple', alpha=0.2)

    axes[idx].set_title(f"{n_layers} Layers ({n_qubits*n_layers*2} Params)\nNoise p={noise_level}")
    axes[idx].set_xlabel("Epoch")
    axes[idx].set_ylabel("Expected Energy")
    axes[idx].legend()
    axes[idx].grid(True, alpha=0.3)

with open('std_qng_noisy.json', 'w') as f:
    json.dump(std_qng_logs, f, indent=4)

with open('dfs_qng_noisy.json', 'w') as f:
    json.dump(dfs_qng_logs, f, indent=4)

print("Finished! Logs saved to std_qng_noisy.json and dfs_qng_noisy.json")

plt.tight_layout()
plt.show()