
import numpy as np

try:
    import torch
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False

class DifferentiableCircuit:
    """Quantum circuit layer with real PyTorch gradient support via parameter shift rule."""

    def __init__(self, num_qubits, num_params):
        self.num_qubits = num_qubits
        self.num_params = num_params

    def expectation(self, params_np):
        """Compute <Z_0> expectation value using statevector simulation."""
        n = self.num_qubits
        sv = np.zeros(2**n, dtype=complex); sv[0] = 1.0
        for i, theta in enumerate(params_np):
            q = i % n
            c, s = np.cos(theta/2), np.sin(theta/2)
            RY = np.array([[c,-s],[s,c]])
            sv = sv.reshape([2]*n)
            sv = np.tensordot(RY, sv, axes=[[1],[q]])
            order = list(range(1,q+1))+[0]+list(range(q+1,n))
            sv = np.transpose(sv, order).reshape(2**n)
        Z0 = np.array([1 if not (i >> (n-1)) & 1 else -1 for i in range(2**n)], dtype=float)
        return float(np.real(sv.conj() @ (Z0 * sv)))

    def parameter_shift_grad(self, params_np, shift=np.pi/2):
        """Compute gradient using the parameter shift rule."""
        grads = np.zeros(len(params_np))
        for i in range(len(params_np)):
            p_plus = params_np.copy(); p_plus[i] += shift
            p_minus = params_np.copy(); p_minus[i] -= shift
            grads[i] = (self.expectation(p_plus) - self.expectation(p_minus)) / 2
        return grads

if HAS_TORCH:
    class TorchQuantumFunction(torch.autograd.Function):
        @staticmethod
        def forward(ctx, params, circuit):
            params_np = params.detach().numpy()
            ctx.save_for_backward(params)
            ctx.circuit = circuit
            val = circuit.expectation(params_np)
            return torch.tensor(val, dtype=torch.float64)

        @staticmethod
        def backward(ctx, grad_output):
            params, = ctx.saved_tensors
            grads = ctx.circuit.parameter_shift_grad(params.detach().numpy())
            return grad_output * torch.tensor(grads, dtype=torch.float64), None

    class TorchQuantumLayer(torch.nn.Module):
        def __init__(self, num_qubits, num_params):
            super().__init__()
            self.circuit = DifferentiableCircuit(num_qubits, num_params)
            self.params = torch.nn.Parameter(torch.rand(num_params, dtype=torch.float64))

        def forward(self):
            return TorchQuantumFunction.apply(self.params, self.circuit)
