quantum_classifier.py

"""
Quantum Classifier using PennyLane and Variational Quantum Circuits (VQC).

This implements a binary classifier using:
- AngleEmbedding for data encoding
- Variational layers with trainable parameters
- Hybrid quantum-classical training loop
"""

import numpy as np
import pennylane as qml
from pennylane import numpy as pnp
from pennylane.optimize import AdamOptimizer
import time
from sklearn.metrics import accuracy_score
from utils.data_generator import generate_concentric_circles
from utils.visualization import plot_decision_boundary, plot_training_curves


class QuantumClassifier:
    """
    Variational Quantum Circuit Classifier.

    Args:
        n_qubits: Number of qubits (should match number of features)
        n_layers: Number of variational layers
        device_name: PennyLane device ('default.qubit', 'lightning.qubit')
    """

    def __init__(self, n_qubits=4, n_layers=3, device_name='default.qubit'):
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        self.device = qml.device(device_name, wires=n_qubits)

        # Initialize parameters randomly
        self.params = pnp.random.uniform(0, np.pi, (n_layers, n_qubits, 3), requires_grad=True)

        # Create the quantum circuit
        self.qnode = qml.QNode(self._circuit, self.device, interface='autograd')

        # History tracking
        self.history = {'loss': [], 'accuracy': []}

    def _circuit(self, params, x):
        """
        Quantum circuit with encoding and variational layers.

        Args:
            params: Trainable parameters for variational layers
            x: Input features to encode

        Returns:
            Expectation value of PauliZ on first qubit
        """
        # Data Encoding Layer: AngleEmbedding
        # We need to pad or repeat features to match n_qubits
        if len(x) < self.n_qubits:
            # Pad with zeros or repeat features
            x_padded = np.pad(x, (0, self.n_qubits - len(x)), mode='constant')
        elif len(x) > self.n_qubits:
            # Take first n_qubits features
            x_padded = x[:self.n_qubits]
        else:
            x_padded = x

        # Encode data as rotation angles
        for i in range(self.n_qubits):
            qml.RY(x_padded[i], wires=i)

        # Variational Layers
        for layer in range(self.n_layers):
            # Parameterized rotation gates
            for i in range(self.n_qubits):
                qml.Rot(params[layer, i, 0], params[layer, i, 1], params[layer, i, 2], wires=i)

            # Entangling layer (CNOT gates in a ring)
            for i in range(self.n_qubits):
                qml.CNOT(wires=[i, (i + 1) % self.n_qubits])

        # Measurement: Expectation value of PauliZ on first qubit
        return qml.expval(qml.PauliZ(0))

    def _cost_function(self, params, X, y):
        """
        Cost function for training (cross-entropy loss).

        Args:
            params: Circuit parameters
            X: Input features
            y: True labels (0 or 1)

        Returns:
            Average loss over all samples
        """
        predictions = pnp.array([self.qnode(params, x) for x in X])

        # Convert predictions from [-1, 1] to [0, 1]
        predictions = (predictions + 1) / 2

        # Binary cross-entropy loss
        # Add small epsilon to avoid log(0)
        epsilon = 1e-10
        predictions = pnp.clip(predictions, epsilon, 1 - epsilon)

        loss = -pnp.mean(
            y * pnp.log(predictions) + (1 - y) * pnp.log(1 - predictions)
        )

        return loss

    def predict_proba(self, X):
        """
        Predict class probabilities.

        Args:
            X: Input features

        Returns:
            Probability of class 1
        """
        predictions = np.array([self.qnode(self.params, x) for x in X])
        # Convert from [-1, 1] to [0, 1]
        proba = (predictions + 1) / 2
        return proba

    def predict(self, X):
        """
        Predict class labels.

        Args:
            X: Input features

        Returns:
            Predicted labels (0 or 1)
        """
        proba = self.predict_proba(X)
        return (proba >= 0.5).astype(int)

    def fit(self, X_train, y_train, n_epochs=100, learning_rate=0.01, batch_size=None, verbose=True):
        """
        Train the quantum classifier.

        Args:
            X_train: Training features
            y_train: Training labels
            n_epochs: Number of training epochs
            learning_rate: Learning rate for optimizer
            batch_size: Batch size (None for full batch)
            verbose: Print training progress

        Returns:
            Training history
        """
        optimizer = AdamOptimizer(stepsize=learning_rate)

        start_time = time.time()

        for epoch in range(n_epochs):
            # Full batch gradient descent or mini-batch
            if batch_size is None:
                # Full batch
                self.params, loss = optimizer.step_and_cost(
                    lambda p: self._cost_function(p, X_train, y_train),
                    self.params
                )
            else:
                # Mini-batch (simple implementation)
                indices = np.random.choice(len(X_train), batch_size, replace=False)
                X_batch = X_train[indices]
                y_batch = y_train[indices]

                self.params, loss = optimizer.step_and_cost(
                    lambda p: self._cost_function(p, X_batch, y_batch),
                    self.params
                )

                # Calculate full training loss for history
                loss = self._cost_function(self.params, X_train, y_train)

            # Calculate accuracy
            y_pred = self.predict(X_train)
            accuracy = accuracy_score(y_train, y_pred)

            # Store history
            self.history['loss'].append(float(loss))
            self.history['accuracy'].append(float(accuracy))

            if verbose and (epoch + 1) % 10 == 0:
                print(f"Epoch {epoch + 1}/{n_epochs} - Loss: {loss:.4f} - Accuracy: {accuracy:.4f}")

        training_time = time.time() - start_time

        if verbose:
            print(f"\nTraining completed in {training_time:.2f} seconds")

        return self.history

    def evaluate(self, X_test, y_test):
        """
        Evaluate the model on test data.

        Args:
            X_test: Test features
            y_test: Test labels

        Returns:
            Dictionary with evaluation metrics
        """
        y_pred = self.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)

        return {
            'accuracy': accuracy,
            'predictions': y_pred
        }


def main():
    """Main function to demonstrate quantum classifier."""
    print("=" * 60)
    print("Quantum Classifier with PennyLane")
    print("=" * 60)

    # Generate dataset
    print("\n1. Generating concentric circles dataset...")
    X_train, X_test, y_train, y_test = generate_concentric_circles(n_samples=200, noise=0.1)
    print(f"   Training samples: {len(X_train)}, Test samples: {len(X_test)}")

    # Create quantum classifier
    print("\n2. Creating Quantum Classifier...")
    print(f"   - Qubits: 4")
    print(f"   - Variational layers: 3")
    print(f"   - Total parameters: {3 * 4 * 3}")

    qc = QuantumClassifier(n_qubits=4, n_layers=3, device_name='default.qubit')

    # Train the model
    print("\n3. Training Quantum Classifier...")
    start_time = time.time()
    history = qc.fit(X_train, y_train, n_epochs=100, learning_rate=0.01, verbose=True)
    training_time = time.time() - start_time

    # Evaluate on test set
    print("\n4. Evaluating on test set...")
    results = qc.evaluate(X_test, y_test)
    print(f"   Test Accuracy: {results['accuracy']:.4f}")

    # Plot training curves
    print("\n5. Generating visualizations...")
    plot_training_curves(history, title="Quantum Classifier Training",
                        save_path="../results/plots/quantum_training_curves.png")

    # Plot decision boundary
    plot_decision_boundary(
        np.vstack([X_train, X_test]),
        np.hstack([y_train, y_test]),
        qc.predict,
        title="Quantum Classifier Decision Boundary",
        save_path="../results/plots/quantum_decision_boundary.png"
    )

    print("\n6. Results Summary:")
    print(f"   - Training time: {training_time:.2f}s")
    print(f"   - Test accuracy: {results['accuracy']:.4f}")
    print(f"   - Final training loss: {history['loss'][-1]:.4f}")

    print("\n" + "=" * 60)
    print("Quantum classifier training completed!")
    print("=" * 60)

    return {
        'model': qc,
        'history': history,
        'results': results,
        'training_time': training_time
    }


if __name__ == "__main__":
    main()