import numpy as np
import os
import time

# --- M1 OPTIMIZATIONS ---
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["QISKIT_IN_PARALLEL"] = "TRUE"

from sklearn.svm import SVC
from sklearn.metrics import roc_auc_score
from qiskit.circuit import ParameterVector
from qiskit.circuit.library import ZZFeatureMap, RealAmplitudes
from qiskit_machine_learning.kernels import TrainableFidelityQuantumKernel
from qiskit_machine_learning.utils.loss_functions import SVCLoss
from qiskit_algorithms.optimizers import COBYLA
from qiskit_aer import AerSimulator

# --- CONFIGURATION ---
N_QUBITS = 8         
TRAIN_SIZE = 150     # Smaller subset for alignment (speed) now 150 for even faster
TEST_SIZE = 300
MAX_ITERS = 50       # Optimization steps

OUTPUT_DIR = "optimized_kernel_results"
if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)

def main():
    print(f"🚀 Starting Quantum Kernel Optimization (QKA)...")
    print("   Goal: Calculate optimal kernel geometry instead of guessing.")

    # 1. LOAD DATA
    possible_paths = ['vG.0.1/qgan_data_optimized.npz', 'qgan_data_optimized.npz']
    data_path = next((p for p in possible_paths if os.path.exists(p)), None)
    if not data_path: 
        print("❌ Error: qgan_data_optimized.npz not found.")
        return

    data = np.load(data_path)
    X_train_full = data['X_train']
    y_train_full = data['y_train']
    X_test_full = data['X_test']
    y_test_full = data['y_test']

    # 2. SUBSAMPLE (Balanced)
    pos_idx = np.where(y_train_full == 1)[0]
    neg_idx = np.where(y_train_full == 0)[0]
    train_idx = np.concatenate([
        np.random.choice(pos_idx, TRAIN_SIZE // 2, replace=False),
        np.random.choice(neg_idx, TRAIN_SIZE // 2, replace=False)
    ])
    np.random.shuffle(train_idx)
    
    X_train = X_train_full[train_idx]
    y_train = y_train_full[train_idx]
    X_test = X_test_full[:TEST_SIZE]
    y_test = y_test_full[:TEST_SIZE]

    print(f"   📊 Training Alignment on {len(X_train)} samples...")

    # 3. DEFINE TRAINABLE FEATURE MAP
    # Circuit = Data Encoding (ZZ) + Trainable Rotation (RealAmplitudes)
    # This acts as a "Quantum Lens" that the optimizer focuses.
    print("   ⚛️  Initializing Trainable Kernel...")
    
    fm = ZZFeatureMap(N_QUBITS, reps=2, entanglement='linear')
    ansatz = RealAmplitudes(N_QUBITS, reps=1)
    
    combined_circuit = fm.compose(ansatz)
    
    # FIX: Only pass the training_parameters. 
    # The class automatically infers that the other params (fm.parameters) are data inputs.
    quant_kernel = TrainableFidelityQuantumKernel(
        feature_map=combined_circuit,
        training_parameters=ansatz.parameters
    )

    # 4. OPTIMIZE KERNEL (ALIGNMENT)
    print("   🧠 Optimizing Kernel Geometry (COBYLA)...")
    
    # SVCLoss minimizes the classification error of the kernel
    loss_func = SVCLoss(C=1.0, gamma='auto')
    optimizer = COBYLA(maxiter=MAX_ITERS)
    
    from qiskit_machine_learning.kernels.algorithms import QuantumKernelTrainer
    
    trainer = QuantumKernelTrainer(
        quantum_kernel=quant_kernel, 
        loss=loss_func, 
        optimizer=optimizer,
        initial_point=[0.1] * len(ansatz.parameters)
    )
    
    start_time = time.time()
    
    # Train the Kernel parameters
    results = trainer.fit(X_train, y_train)
    
    print(f"   ✅ Optimization Complete in {time.time() - start_time:.1f}s")
    print(f"   🎯 Final Loss: {results.optimal_value:.4f}")
    
    # 5. RUN FINAL SVC WITH OPTIMIZED KERNEL
    print("   🏆 Training Final SVM with Optimized Kernel...")
    
    optimized_kernel = results.quantum_kernel
    
    # We use the optimized kernel to fit the full SVM
    qsvc = SVC(kernel=optimized_kernel.evaluate, probability=True)
    qsvc.fit(X_train, y_train)
    
    # Evaluate
    test_probs = qsvc.predict_proba(X_test)[:, 1]
    test_auc = roc_auc_score(y_test, test_probs)
    
    print("\n" + "="*40)
    print("🚀 OPTIMIZED QUANTUM KERNEL RESULTS")
    print("="*40)
    print(f"✅ Test AUC:      {test_auc:.4f}")
    print(f"📉 Linear Base:   0.7500")
    print("="*40)
    
    # Check if we beat the baseline
    if test_auc > 0.75:
        print("🎉 SUCCESS: Optimization found a better topology than brute force.")
        np.save(f"{OUTPUT_DIR}/optimized_weights.npy", results.optimal_point)
    else:
        print("🔸 Result: Optimization matched or slightly underperformed baseline.")

if __name__ == "__main__":
    main()