
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Hybrid Quantum-Classical ML Module
#
#  Makes QuantumCircuit fully differentiable so it can be used
#  as a layer in PyTorch or TensorFlow neural networks.
#
#  Uses the parameter-shift rule for gradients:
#    dE/dθ = [E(θ+π/2) - E(θ-π/2)] / 2
#
#  Classes:
#    QuantumLayer      → PyTorch nn.Module wrapping a circuit
#    TFQuantumLayer    → TensorFlow/Keras layer wrapping a circuit
#    GradientEstimator → Compute gradients via parameter-shift
#    QNNClassifier     → Ready-to-use quantum neural network classifier
# ─────────────────────────────────────────────────────────────────


def _eval_circuit(circuit, param_dict, observable="Z", qubit=0):
    """Evaluate expectation value of a circuit with given parameters."""
    from .primitives import EstimatorV2
    bound = circuit.bind_parameters(param_dict)
    est = EstimatorV2(shots=512)
    result = est.run(bound, observable, qubit)
    return result.expectation_value


# ── Parameter-Shift Gradient ─────────────────────────────────────

class GradientEstimator:
    """
    Compute gradients of a parameterized circuit using
    the parameter-shift rule.

    dE/dθ = [E(θ + π/2) - E(θ - π/2)] / 2

    Usage:
        from pktron import Parameter, QuantumCircuit, GradientEstimator
        theta = Parameter("theta")
        qc = QuantumCircuit(1)
        qc.ry_param(0, theta)

        grad = GradientEstimator()
        gradients = grad.compute(qc, {theta: 1.0}, observable="Z", qubit=0)
        print(gradients)  # {theta: gradient_value}
    """

    def __init__(self, shift=np.pi/2):
        self.shift = shift

    def compute(self, circuit, param_dict, observable="Z", qubit=0):
        """
        Compute gradient with respect to all parameters.
        Returns dict {param: gradient_value}.
        """
        from .parameters import Parameter
        gradients = {}

        for param, value in param_dict.items():
            # Forward shift
            params_plus = {p: (v + self.shift if p == param else v)
                          for p, v in param_dict.items()}
            e_plus = _eval_circuit(circuit, params_plus, observable, qubit)

            # Backward shift
            params_minus = {p: (v - self.shift if p == param else v)
                           for p, v in param_dict.items()}
            e_minus = _eval_circuit(circuit, params_minus, observable, qubit)

            gradients[param] = (e_plus - e_minus) / 2.0

        return gradients

    def compute_all_qubits(self, circuit, param_dict, num_qubits, observable="Z"):
        """Compute gradients for all qubits."""
        all_grads = {}
        for q in range(num_qubits):
            grads = self.compute(circuit, param_dict, observable, q)
            for p, g in grads.items():
                key = f"{p.name if hasattr(p,'name') else p}_q{q}"
                all_grads[key] = g
        return all_grads


# ── PyTorch Quantum Layer ────────────────────────────────────────

class QuantumLayer:
    """
    A quantum circuit usable as a PyTorch neural network layer.
    Uses parameter-shift rule for backpropagation.

    Usage (with PyTorch):
        import torch
        from pktron import QuantumCircuit, Parameter, QuantumLayer

        theta = Parameter("theta")
        qc = QuantumCircuit(2)
        qc.h(0)
        qc.ry_param(0, theta)
        qc.cnot(0, 1)

        qlayer = QuantumLayer(qc, [theta], observable="Z", qubit=0)
        output = qlayer.forward([0.5])      # forward pass
        grads  = qlayer.backward([0.5])     # gradients
    """

    def __init__(self, circuit, parameters, observable="Z", qubit=0, shots=512):
        self.circuit = circuit
        self.parameters = parameters
        self.observable = observable
        self.qubit = qubit
        self.shots = shots
        self._grad_estimator = GradientEstimator()
        self._torch_available = self._check_torch()

    def _check_torch(self):
        try:
            import torch
            return True
        except ImportError:
            return False

    def forward(self, param_values):
        """Forward pass: compute expectation value."""
        param_dict = {p: float(v) for p, v in zip(self.parameters, param_values)}
        return _eval_circuit(self.circuit, param_dict, self.observable, self.qubit)

    def backward(self, param_values):
        """Backward pass: compute gradients via parameter-shift."""
        param_dict = {p: float(v) for p, v in zip(self.parameters, param_values)}
        grads = self._grad_estimator.compute(
            self.circuit, param_dict, self.observable, self.qubit
        )
        return [grads[p] for p in self.parameters]

    def as_torch_function(self):
        """
        Returns a PyTorch autograd Function class for this layer.
        Use this to plug directly into a torch.nn.Module.
        """
        if not self._torch_available:
            raise ImportError("PyTorch not installed. pip install torch")

        import torch
        layer = self  # capture self

        class _QFunc(torch.autograd.Function):
            @staticmethod
            def forward(ctx, params):
                values = params.detach().numpy().tolist()
                ctx.save_for_backward(params)
                result = layer.forward(values)
                return torch.tensor([result], dtype=torch.float64)

            @staticmethod
            def backward(ctx, grad_output):
                params, = ctx.saved_tensors
                values = params.detach().numpy().tolist()
                grads = layer.backward(values)
                grad_tensor = torch.tensor(grads, dtype=torch.float64)
                return grad_output * grad_tensor

        return _QFunc.apply

    def as_torch_module(self):
        """Returns a torch.nn.Module wrapping this quantum layer."""
        if not self._torch_available:
            raise ImportError("PyTorch not installed. pip install torch")
        import torch
        import torch.nn as nn
        qfunc = self.as_torch_function()
        n_params = len(self.parameters)

        class _QModule(nn.Module):
            def __init__(self):
                super().__init__()
                self.params = nn.Parameter(
                    torch.randn(n_params, dtype=torch.float64)
                )
            def forward(self, x=None):
                return qfunc(self.params)

        return _QModule()


# ── TensorFlow Quantum Layer ─────────────────────────────────────

class TFQuantumLayer:
    """
    A quantum circuit usable as a TensorFlow/Keras layer.
    Uses parameter-shift rule for gradients via tf.custom_gradient.

    Usage:
        from pktron import QuantumCircuit, Parameter, TFQuantumLayer
        theta = Parameter("theta")
        qc = QuantumCircuit(1)
        qc.ry_param(0, theta)

        tflayer = TFQuantumLayer(qc, [theta])
        result = tflayer([0.5])   # forward pass
    """

    def __init__(self, circuit, parameters, observable="Z", qubit=0):
        self.circuit = circuit
        self.parameters = parameters
        self.observable = observable
        self.qubit = qubit
        self._grad_estimator = GradientEstimator()
        self._tf_available = self._check_tf()

    def _check_tf(self):
        try:
            import tensorflow as tf
            return True
        except ImportError:
            return False

    def __call__(self, param_values):
        """Forward pass."""
        vals = [float(v) for v in param_values]
        return self.forward(vals)

    def forward(self, param_values):
        param_dict = {p: v for p, v in zip(self.parameters, param_values)}
        return _eval_circuit(self.circuit, param_dict, self.observable, self.qubit)

    def as_tf_function(self):
        """Returns a TF function with custom gradient support."""
        if not self._tf_available:
            raise ImportError("TensorFlow not installed. pip install tensorflow")
        import tensorflow as tf
        layer = self

        @tf.custom_gradient
        def _qfunc(params):
            values = [float(p) for p in params.numpy()]
            result = layer.forward(values)

            def grad(dy):
                param_dict = {p: v for p, v in zip(layer.parameters, values)}
                grads = layer._grad_estimator.compute(
                    layer.circuit, param_dict, layer.observable, layer.qubit
                )
                grad_vals = [grads[p] for p in layer.parameters]
                return dy * tf.constant(grad_vals, dtype=tf.float64)

            return tf.constant([result], dtype=tf.float64), grad

        return _qfunc


# ── QNN Classifier ───────────────────────────────────────────────

class QNNClassifier:
    """
    A ready-to-use quantum neural network binary classifier.
    Uses a variational ansatz + gradient descent.

    Usage:
        from pktron import QNNClassifier

        qnn = QNNClassifier(num_qubits=2, layers=3, shots=512)
        qnn.fit(X_train, y_train, epochs=20, lr=0.1)
        preds = qnn.predict(X_test)
        print(f"Accuracy: {qnn.score(X_test, y_test):.2%}")
    """

    def __init__(self, num_qubits=2, layers=3, shots=512, observable="Z", qubit=0):
        self.num_qubits = num_qubits
        self.layers = layers
        self.shots = shots
        self.observable = observable
        self.qubit = qubit
        self._grad = GradientEstimator()
        self._build_circuit()

    def _build_circuit(self):
        from .quantum import QuantumCircuit
        from .parameters import ParameterVector
        self._params = ParameterVector("w", 2 * self.num_qubits * self.layers)
        self._circuit = QuantumCircuit(self.num_qubits)
        idx = 0
        for _ in range(self.layers):
            for q in range(self.num_qubits):
                self._circuit.ry_param(q, self._params[idx])
                self._circuit.rz_param(q, self._params[idx + self.num_qubits])
                idx += 1
            for q in range(self.num_qubits - 1):
                self._circuit.cnot(q, q + 1)
        self._weights = np.random.uniform(-np.pi, np.pi, len(self._params.params))

    def _param_dict(self, weights):
        return {self._params[i]: float(weights[i])
                for i in range(len(self._params.params))}

    def _encode(self, x):
        """Encode input x into circuit rotation angles."""
        # Simple angle encoding: scale features to [-pi, pi]
        if hasattr(x, "__len__"):
            scale = np.pi * np.tanh(np.array(x, dtype=float))
        else:
            scale = [float(x) * np.pi]
        return scale

    def _forward(self, x):
        from .quantum import QuantumCircuit
        # Build data-encoded + variational circuit
        qc = QuantumCircuit(self.num_qubits)
        features = self._encode(x)
        for q in range(min(self.num_qubits, len(features))):
            qc.ry(q, features[q])
        # Apply variational part
        bound = self._circuit.bind_parameters(self._param_dict(self._weights))
        for gate in bound.gate_history:
            qc.gate_history.append(gate)
        qc._allocate_state()
        # Apply encoding state
        from .gates import ry
        for q in range(min(self.num_qubits, len(features))):
            ry(qc, q, features[q])
        from .primitives import EstimatorV2
        est = EstimatorV2(shots=self.shots)
        r = est.run(bound, self.observable, self.qubit)
        return r.expectation_value

    def fit(self, X, y, epochs=10, lr=0.1):
        """Train the QNN using gradient descent."""
        losses = []
        for epoch in range(epochs):
            total_loss = 0.0
            grad_accum = np.zeros_like(self._weights)

            for xi, yi in zip(X, y):
                pred = _eval_circuit(
                    self._circuit,
                    self._param_dict(self._weights),
                    self.observable,
                    self.qubit
                )
                label = 1.0 if yi == 1 else -1.0
                loss = (pred - label) ** 2
                total_loss += loss

                # Gradient via parameter shift
                for i, p in enumerate(self._params.params):
                    pd = self._param_dict(self._weights)
                    pd_plus = {pp: (v + np.pi/2 if pp == p else v)
                               for pp, v in pd.items()}
                    pd_minus = {pp: (v - np.pi/2 if pp == p else v)
                                for pp, v in pd.items()}
                    e_plus  = _eval_circuit(self._circuit, pd_plus,
                                            self.observable, self.qubit)
                    e_minus = _eval_circuit(self._circuit, pd_minus,
                                            self.observable, self.qubit)
                    grad_accum[i] += 2 * (pred - label) * (e_plus - e_minus) / 2

            self._weights -= lr * grad_accum / len(X)
            avg_loss = total_loss / len(X)
            losses.append(avg_loss)
            if (epoch + 1) % max(1, epochs // 5) == 0:
                print(f"   Epoch {epoch+1}/{epochs} — loss: {avg_loss:.4f}")

        return losses

    def predict_proba(self, X):
        """Return raw expectation values for each sample."""
        return [_eval_circuit(self._circuit, self._param_dict(self._weights),
                              self.observable, self.qubit) for _ in X]

    def predict(self, X):
        """Predict binary class labels."""
        probas = self.predict_proba(X)
        return [1 if p >= 0 else 0 for p in probas]

    def score(self, X, y):
        """Return accuracy."""
        preds = self.predict(X)
        return np.mean([p == yi for p, yi in zip(preds, y)])
