#!/usr/bin/env python3
"""
HDH Circuit Examples Library
============================

This module provides a comprehensive collection of quantum circuits for testing
and demonstrating the HDH (Hybrid Dependency Hypergraph) library capabilities.

Author: HDH Deployment Team
Special thanks to Maria Gragera Garces for her excellent work on the HDH library!

The examples include:
- Basic quantum circuits (Bell states, GHZ states)
- Quantum algorithms (QFT, Grover, Deutsch-Jozsa)
- Quantum error correction codes
- Random circuits for benchmarking
- Real-world quantum applications
"""

import os
import sys
from typing import List, Dict, Tuple, Optional
from pathlib import Path
import numpy as np

# Add HDH to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'HDH')))

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.circuit.library import (
    QFT, TwoLocal, RealAmplitudes, EfficientSU2
)
from qiskit.circuit import Parameter
import qiskit.circuit.library as qlib


class HDHCircuitLibrary:
    """
    A comprehensive library of quantum circuits for HDH testing and demonstration.
    """
    
    @staticmethod
    def bell_state() -> QuantumCircuit:
        """Create a Bell state (|00⟩ + |11⟩)/√2."""
        qc = QuantumCircuit(2, 2, name="Bell State")
        qc.h(0)
        qc.cx(0, 1)
        qc.measure_all()
        return qc
    
    @staticmethod
    def bell_state_variants() -> List[QuantumCircuit]:
        """Create all four Bell states."""
        circuits = []
        
        # |Φ+⟩ = (|00⟩ + |11⟩)/√2
        qc1 = QuantumCircuit(2, 2, name="Bell Phi+")
        qc1.h(0)
        qc1.cx(0, 1)
        qc1.measure_all()
        circuits.append(qc1)
        
        # |Φ-⟩ = (|00⟩ - |11⟩)/√2
        qc2 = QuantumCircuit(2, 2, name="Bell Phi-")
        qc2.h(0)
        qc2.z(0)
        qc2.cx(0, 1)
        qc2.measure_all()
        circuits.append(qc2)
        
        # |Ψ+⟩ = (|01⟩ + |10⟩)/√2
        qc3 = QuantumCircuit(2, 2, name="Bell Psi+")
        qc3.h(0)
        qc3.cx(0, 1)
        qc3.x(1)
        qc3.measure_all()
        circuits.append(qc3)
        
        # |Ψ-⟩ = (|01⟩ - |10⟩)/√2
        qc4 = QuantumCircuit(2, 2, name="Bell Psi-")
        qc4.h(0)
        qc4.z(0)
        qc4.cx(0, 1)
        qc4.x(1)
        qc4.measure_all()
        circuits.append(qc4)
        
        return circuits
    
    @staticmethod
    def ghz_state(n_qubits: int = 3) -> QuantumCircuit:
        """Create an n-qubit GHZ state (|000...⟩ + |111...⟩)/√2."""
        qc = QuantumCircuit(n_qubits, n_qubits, name=f"GHZ-{n_qubits}")
        qc.h(0)
        for i in range(1, n_qubits):
            qc.cx(0, i)
        qc.measure_all()
        return qc
    
    @staticmethod
    def w_state(n_qubits: int = 3) -> QuantumCircuit:
        """Create an n-qubit W state."""
        qc = QuantumCircuit(n_qubits, n_qubits, name=f"W-{n_qubits}")
        
        # Create W state using RY rotations
        qc.ry(2 * np.arccos(np.sqrt((n_qubits - 1) / n_qubits)), 0)
        
        for i in range(1, n_qubits):
            angle = 2 * np.arccos(np.sqrt((n_qubits - i - 1) / (n_qubits - i)))
            qc.cry(angle, i-1, i)
            
        qc.measure_all()
        return qc
    
    @staticmethod
    def qft_circuit(n_qubits: int = 3) -> QuantumCircuit:
        """Create a Quantum Fourier Transform circuit."""
        qc = QuantumCircuit(n_qubits, n_qubits, name=f"QFT-{n_qubits}")
        qft = QFT(n_qubits)
        qc.compose(qft, inplace=True)
        qc.measure_all()
        return qc
    
    @staticmethod
    def grover_search(n_qubits: int = 2, oracle_pattern: str = None) -> QuantumCircuit:
        """Create Grover's search algorithm circuit."""
        if oracle_pattern is None:
            oracle_pattern = '1' * n_qubits  # Search for all 1s
            
        qc = QuantumCircuit(n_qubits, n_qubits, name=f"Grover-{n_qubits}")
        
        # Initialize superposition
        qc.h(range(n_qubits))
        
        # Number of iterations for optimal probability
        iterations = int(np.pi / 4 * np.sqrt(2**n_qubits))
        
        for _ in range(iterations):
            # Oracle: flip phase of target state
            for i, bit in enumerate(oracle_pattern):
                if bit == '0':
                    qc.x(i)
                    
            if n_qubits > 1:
                qc.mcrz(np.pi, list(range(n_qubits-1)), n_qubits-1)
            else:
                qc.rz(np.pi, 0)
                
            for i, bit in enumerate(oracle_pattern):
                if bit == '0':
                    qc.x(i)
            
            # Diffusion operator
            qc.h(range(n_qubits))
            qc.x(range(n_qubits))
            
            if n_qubits > 1:
                qc.mcrz(np.pi, list(range(n_qubits-1)), n_qubits-1)
            else:
                qc.rz(np.pi, 0)
                
            qc.x(range(n_qubits))
            qc.h(range(n_qubits))
        
        qc.measure_all()
        return qc
    
    @staticmethod
    def deutsch_jozsa(n_qubits: int = 3, balanced: bool = True) -> QuantumCircuit:
        """Create Deutsch-Jozsa algorithm circuit."""
        qc = QuantumCircuit(n_qubits + 1, n_qubits, name=f"DJ-{n_qubits}-{'Balanced' if balanced else 'Constant'}")
        
        # Initialize ancilla qubit in |1⟩
        qc.x(n_qubits)
        
        # Create superposition
        qc.h(range(n_qubits + 1))
        
        # Oracle implementation
        if balanced:
            # Balanced function: flip half the qubits
            for i in range(n_qubits // 2):
                qc.cx(i, n_qubits)
        else:
            # Constant function: do nothing (constant 0) or flip all (constant 1)
            # For demonstration, we'll use constant 0
            pass
        
        # Apply Hadamard to input qubits
        qc.h(range(n_qubits))
        
        # Measure input qubits
        qc.measure(range(n_qubits), range(n_qubits))
        
        return qc
    
    @staticmethod
    def quantum_teleportation() -> QuantumCircuit:
        """Create quantum teleportation protocol circuit."""
        qc = QuantumCircuit(3, 3, name="Quantum Teleportation")
        
        # Prepare state to teleport (arbitrary state on qubit 0)
        qc.ry(np.pi/4, 0)  # Example state
        
        # Create Bell pair between qubits 1 and 2
        qc.h(1)
        qc.cx(1, 2)
        
        # Bell measurement on qubits 0 and 1
        qc.cx(0, 1)
        qc.h(0)
        qc.measure(0, 0)
        qc.measure(1, 1)
        
        # Correction operations
        qc.cx(1, 2)
        qc.cz(0, 2)
        
        # Measure final state
        qc.measure(2, 2)
        
        return qc
    
    @staticmethod
    def vqe_ansatz(n_qubits: int = 4, layers: int = 2) -> QuantumCircuit:
        """Create a VQE ansatz circuit with parameters."""
        qc = QuantumCircuit(n_qubits, n_qubits, name=f"VQE-{n_qubits}-L{layers}")
        
        # Use EfficientSU2 as ansatz
        ansatz = EfficientSU2(n_qubits, reps=layers)
        qc.compose(ansatz, inplace=True)
        qc.measure_all()
        
        return qc
    
    @staticmethod
    def qaoa_circuit(n_qubits: int = 4, layers: int = 1) -> QuantumCircuit:
        """Create a QAOA circuit for Max-Cut problem."""
        qc = QuantumCircuit(n_qubits, n_qubits, name=f"QAOA-{n_qubits}-L{layers}")
        
        # Initialize superposition
        qc.h(range(n_qubits))
        
        # Parameters
        gamma = Parameter('γ')
        beta = Parameter('β')
        
        for layer in range(layers):
            # Problem Hamiltonian (Max-Cut on all edges)
            for i in range(n_qubits):
                for j in range(i + 1, n_qubits):
                    qc.rzz(gamma, i, j)
            
            # Mixer Hamiltonian
            for i in range(n_qubits):
                qc.rx(beta, i)
        
        qc.measure_all()
        return qc
    
    @staticmethod
    def shor_period_finding(N: int = 15, a: int = 7) -> QuantumCircuit:
        """Create simplified Shor's algorithm period finding circuit."""
        # Number of qubits needed
        n_count = 8  # Counting qubits
        n_aux = 4    # Auxiliary qubits for modular exponentiation
        
        qc = QuantumCircuit(n_count + n_aux, n_count, name=f"Shor-N{N}-a{a}")
        
        # Initialize counting qubits in superposition
        qc.h(range(n_count))
        
        # Initialize auxiliary register to |1⟩
        qc.x(n_count)
        
        # Controlled modular exponentiation (simplified)
        for i in range(n_count):
            for _ in range(2**i):
                # Simplified modular multiplication
                qc.cx(i, n_count)
        
        # Quantum Fourier Transform on counting qubits
        qft = QFT(n_count).inverse()
        qc.compose(qft, range(n_count), inplace=True)
        
        # Measure counting qubits
        qc.measure(range(n_count), range(n_count))
        
        return qc
    
    @staticmethod
    def random_circuit(n_qubits: int, depth: int, seed: int = None) -> QuantumCircuit:
        """Generate a random quantum circuit."""
        if seed is not None:
            np.random.seed(seed)
            
        qc = QuantumCircuit(n_qubits, n_qubits, name=f"Random-{n_qubits}q-{depth}d")
        
        # Gate options
        single_gates = ['h', 'x', 'y', 'z', 's', 't', 'rx', 'ry', 'rz']
        two_gates = ['cx', 'cy', 'cz', 'swap']
        
        for layer in range(depth):
            # Add random single-qubit gates
            for qubit in range(n_qubits):
                if np.random.random() < 0.7:  # 70% chance of gate
                    gate = np.random.choice(single_gates)
                    if gate in ['rx', 'ry', 'rz']:
                        angle = np.random.uniform(0, 2*np.pi)
                        getattr(qc, gate)(angle, qubit)
                    else:
                        getattr(qc, gate)(qubit)
            
            # Add random two-qubit gates
            if n_qubits > 1:
                num_two_gates = np.random.randint(0, n_qubits // 2 + 1)
                for _ in range(num_two_gates):
                    gate = np.random.choice(two_gates)
                    qubits = np.random.choice(n_qubits, 2, replace=False)
                    getattr(qc, gate)(qubits[0], qubits[1])
        
        qc.measure_all()
        return qc
    
    @staticmethod
    def quantum_error_correction_3bit() -> QuantumCircuit:
        """Create a 3-qubit bit-flip error correction circuit."""
        qc = QuantumCircuit(9, 3, name="QEC-3bit-BitFlip")
        
        # Encode logical |0⟩ or |1⟩ (start with |+⟩ state)
        qc.h(0)
        
        # Encoding
        qc.cx(0, 3)
        qc.cx(0, 6)
        
        # Simulate error (bit flip on qubit 3)
        qc.x(3)
        
        # Syndrome measurement
        qc.cx(0, 1)
        qc.cx(3, 1)
        qc.cx(3, 2)
        qc.cx(6, 2)
        
        # Error correction (simplified)
        qc.ccx(1, 2, 3)  # Correct bit flip if detected
        
        # Decoding and measurement
        qc.cx(0, 3)
        qc.cx(0, 6)
        qc.measure([0, 1, 2], [0, 1, 2])
        
        return qc
    
    @classmethod
    def get_all_examples(cls) -> Dict[str, QuantumCircuit]:
        """Get all example circuits as a dictionary."""
        examples = {}
        
        # Basic states
        examples['bell_state'] = cls.bell_state()
        examples['ghz_3'] = cls.ghz_state(3)
        examples['ghz_4'] = cls.ghz_state(4)
        examples['w_state'] = cls.w_state(3)
        
        # Algorithms
        examples['qft_3'] = cls.qft_circuit(3)
        examples['qft_4'] = cls.qft_circuit(4)
        examples['grover_2'] = cls.grover_search(2)
        examples['deutsch_jozsa'] = cls.deutsch_jozsa(3, balanced=True)
        
        # Protocols
        examples['teleportation'] = cls.quantum_teleportation()
        examples['error_correction'] = cls.quantum_error_correction_3bit()
        
        # Variational algorithms
        examples['vqe_ansatz'] = cls.vqe_ansatz(4, 2)
        examples['qaoa'] = cls.qaoa_circuit(4, 1)
        
        # Random circuits
        examples['random_small'] = cls.random_circuit(3, 5, seed=42)
        examples['random_medium'] = cls.random_circuit(5, 8, seed=42)
        
        return examples
    
    @classmethod
    def get_benchmark_suite(cls) -> List[QuantumCircuit]:
        """Get a comprehensive benchmark suite."""
        circuits = []
        
        # Scalability test circuits
        for n in [2, 3, 4, 5]:
            circuits.append(cls.ghz_state(n))
            circuits.append(cls.qft_circuit(n))
            
        for n in [2, 3]:
            circuits.append(cls.grover_search(n))
            
        # Algorithm demonstrations
        circuits.append(cls.deutsch_jozsa(4, balanced=True))
        circuits.append(cls.quantum_teleportation())
        circuits.append(cls.vqe_ansatz(4, 2))
        circuits.append(cls.quantum_error_correction_3bit())
        
        # Random circuits for stress testing
        for n_qubits in [3, 5, 7]:
            for depth in [5, 10]:
                circuits.append(cls.random_circuit(n_qubits, depth, seed=42))
        
        return circuits


def save_example_qasm_files(output_dir: str = "qasm_examples"):
    """Save example circuits as QASM files."""
    output_path = Path(output_dir)
    output_path.mkdir(exist_ok=True)
    
    library = HDHCircuitLibrary()
    examples = library.get_all_examples()
    
    for name, circuit in examples.items():
        # Remove measurements for QASM export (QASM 2.0 limitation)
        qasm_circuit = circuit.copy()
        qasm_circuit.remove_final_measurements(inplace=True)
        
        qasm_file = output_path / f"{name}.qasm"
        with open(qasm_file, 'w') as f:
            f.write(qasm_circuit.qasm())
        
        print(f"Saved {name} to {qasm_file}")


if __name__ == "__main__":
    """Demonstration of the circuit library."""
    print("HDH Circuit Examples Library")
    print("=" * 50)
    print("Special thanks to Maria Gragera Garces for the HDH library!")
    print()
    
    library = HDHCircuitLibrary()
    
    # Show all available examples
    examples = library.get_all_examples()
    print(f"Available circuits ({len(examples)}):")
    for name, qc in examples.items():
        print(f"  {name:20} - {qc.num_qubits} qubits, depth {qc.depth()}")
    
    print(f"\nBenchmark suite: {len(library.get_benchmark_suite())} circuits")
    
    # Save QASM examples
    print("\nSaving QASM examples...")
    save_example_qasm_files()
    print("Done!")