from qiskit import qpy
from qiskit.circuit import ParameterVector
from qiskit_machine_learning.neural_networks import SamplerQNN
from qiskit_machine_learning.algorithms.classifiers import NeuralNetworkClassifier
from qiskit.primitives import Sampler
import numpy as np
from huggingface_hub import hf_hub_download


def load_qnn_model(repo_id="squ11z1/Two-Moons"):
 
    # Download files
    circuit_path = hf_hub_download(repo_id=repo_id, filename="circuit.qpy")
    weights_path = hf_hub_download(repo_id=repo_id, filename="weights.npy")
    
    # Load circuit
    with open(circuit_path, 'rb') as f:
        circuit = qpy.load(f)[0]
    
    # Load weights
    weights = np.load(weights_path)
    
    return circuit, weights


def create_qnn_classifier(circuit, weights):

    # Get parameters from circuit
    input_params = [p for p in circuit.parameters if p.name.startswith('x')]
    weight_params = [p for p in circuit.parameters if p.name.startswith('w')]

    # Parity interpretation function
    def parity(x):
        # Convert to binary string and count ones
        return bin(x).count("1") % 2

    # Create sampler (using V2 to avoid deprecation)
    from qiskit.primitives import StatevectorSampler as Sampler

    sampler = Sampler()

    # Create QNN
    qnn = SamplerQNN(
        circuit=circuit,
        input_params=input_params,
        weight_params=weight_params,
        interpret=parity,
        output_shape=2,
        sampler=sampler
    )
    
    # Create classifier
    classifier = NeuralNetworkClassifier(
        neural_network=qnn,
        optimizer=None  # Pre-trained
    )
    
    # Set weights
    classifier._fit_result = type('obj', (object,), {'x': weights})
    
    return classifier


def predict(X, repo_id="squ11z1/Two-Moons"):
   
    # Load model
    circuit, weights = load_qnn_model(repo_id)
    
    # Create classifier
    classifier = create_qnn_classifier(circuit, weights)
    
    # Predict
    predictions = classifier.predict(X)
    
    return predictions


# Example usage
if __name__ == "__main__":
    print("="*70)
    print("Quantum Neural Network - Inference Example")
    print("="*70)
    
    # Example: Load and predict
    repo_id = "squ11z1/Two-Moons"
    
    print("\n1. Loading model from Hugging Face...")
    circuit, weights = load_qnn_model(repo_id)
    print(f"Circuit: {circuit.num_qubits} qubits, depth {circuit.depth()}")
    print(f"Weights: {weights}")
    
    print("\n2. Creating classifier...")
    classifier = create_qnn_classifier(circuit, weights)
    print(f" Classifier ready")
    
    print("\n3. Making predictions...")
    # Example data (Two Moons-like points)
    X_test = np.array([
        [0.5, 0.2],   # Class 1
        [-0.5, 0.5],  # Class 0
        [1.0, 0.0],   # Class 1
        [-1.0, 0.8]   # Class 0
    ])
    
    predictions = classifier.predict(X_test)
    
    print(f"\n   Input data:")
    for i, x in enumerate(X_test):
        print(f"     Sample {i+1}: {x}")
    
    print(f"\n   Predictions: {predictions}")
    print(f"     (0 = Negative class, 1 = Positive class)")
    
    print("\n" + "="*70)
    print("Inference complete!")
    print("="*70)