"""Independent implementation of arXiv:2503.16678, Fig.2b and Tables 2/3.

Inputs are physical (t,x,y), outputs (u,v,p). No reference fields enter forward.
The three groups of n circuit angles are independent, as specified by 3*n*L.
"""
import pennylane as qml
import torch
from torch import nn


class QCPINN(nn.Module):
    def __init__(self, input_dim=3, output_dim=3, hidden_dim=50,
                 n_qubits=5, quantum_layers=1, backend="default.qubit", shots=None):
        super().__init__()
        if shots is not None:
            raise ValueError("The paper requires analytic simulation: shots=None")
        self.n_qubits = n_qubits
        self.quantum_layers = quantum_layers
        self.preprocessor = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.Tanh(),
                                         nn.Linear(hidden_dim, n_qubits))
        self.quantum_angles = nn.Parameter(torch.empty(quantum_layers, 3*n_qubits))
        self.postprocessor = nn.Sequential(nn.Linear(n_qubits, hidden_dim), nn.Tanh(),
                                          nn.Linear(hidden_dim, output_dim))
        for block in (self.preprocessor, self.postprocessor):
            for module in block:
                if isinstance(module, nn.Linear):
                    nn.init.xavier_normal_(module.weight)
                    nn.init.zeros_(module.bias)
        nn.init.xavier_normal_(self.quantum_angles)
        device = qml.device(backend, wires=n_qubits, shots=None)

        @qml.qnode(device, interface="torch", diff_method="backprop")
        def circuit(features, angles):
            qml.AngleEmbedding(features, wires=range(n_qubits), rotation="X")
            for layer in range(quantum_layers):
                for wire in range(n_qubits):
                    qml.RX(angles[layer, wire], wires=wire)
                for wire in range(n_qubits):
                    qml.RZ(angles[layer, n_qubits+wire], wires=wire)
                # Ring order matches Fig.2b, with distinct angles for every CRX.
                edges = [(n_qubits-1, 0)] + [(i-1, i) for i in range(n_qubits-1, 0, -1)]
                for index, edge in enumerate(edges):
                    qml.CRX(angles[layer, 2*n_qubits+index], wires=edge)
            return tuple(qml.expval(qml.PauliZ(w)) for w in range(n_qubits))

        self.circuit = circuit

    def forward(self, coordinates):
        if coordinates.ndim != 2 or coordinates.shape[1] != 3:
            raise ValueError(f"Expected (batch,3) (t,x,y); got {coordinates.shape}")
        encoded = self.preprocessor(coordinates)
        features = torch.stack(self.circuit(encoded, self.quantum_angles), dim=-1)
        return self.postprocessor(features.to(dtype=coordinates.dtype))