#!/usr/bin/env python3
"""Working Guppy Grover implementation"""
from guppylang import guppy
from guppylang.std.builtins import result
from guppylang.std.quantum import qubit, h, x, cz, cx, measure
import math
import json

@guppy
def grover_n2_simple() -> None:
    """Grover for n=2, marking |11⟩ with single iteration"""
    # Initialize
    q0 = qubit()
    q1 = qubit()
    
    # Hadamard to create superposition
    h(q0)
    h(q1)
    
    # Oracle: mark |11⟩ with CZ
    cz(q0, q1)
    
    # Diffusion operator
    h(q0)
    h(q1)
    x(q0)
    x(q1)
    
    # CZ for diffusion
    h(q1)
    cx(q0, q1)
    h(q1)
    
    x(q0)
    x(q1)
    h(q0)
    h(q1)
    
    # Measure
    m0 = measure(q0)
    m1 = measure(q1)
    
    # Convert bool to int for result
    r0 = 1 if m0 else 0
    r1 = 1 if m1 else 0
    
    result("q0", r0)
    result("q1", r1)

def run_experiment(shots=1000):
    """Run Grover experiment and analyze results"""
    print(f"\nRunning Guppy/Selene Grover (n=2, marking |11⟩)...")
    
    # Create and run emulator
    sim = grover_n2_simple.emulator(n_qubits=2)
    sim = sim.with_shots(shots)
    results = sim.run()
    
    # Count outcomes
    outcomes = {"00": 0, "01": 0, "10": 0, "11": 0}
    
    for shot in results.results:
        q0_val = None
        q1_val = None
        
        for entry in shot.entries:
            if entry[0] == 'q0':
                q0_val = entry[1]
            elif entry[0] == 'q1':
                q1_val = entry[1]
        
        if q0_val is not None and q1_val is not None:
            outcome = f"{q0_val}{q1_val}"
            outcomes[outcome] = outcomes.get(outcome, 0) + 1
    
    # Calculate probabilities
    probs = {k: v/shots for k, v in outcomes.items()}
    
    return probs, outcomes

def main():
    print("="*60)
    print("GUPPY/SELENE QUANTUM EMULATOR - GROVER'S ALGORITHM")
    print("="*60)
    
    # Run with different shot counts to show consistency
    for shots in [100, 1000, 4000]:
        probs, counts = run_experiment(shots)
        
        print(f"\nShots: {shots}")
        print("-" * 30)
        for state, prob in sorted(probs.items()):
            bar = "█" * int(prob * 20)
            marked = "⭐" if state == "11" else ""
            print(f"|{state}⟩: {prob:.3f} [{bar:<20}] {marked}")
        
        success_rate = probs.get("11", 0)
        print(f"\nSuccess rate: {success_rate:.1%}")
    
    # Final high-precision run
    print("\n" + "="*60)
    print("FINAL HIGH-PRECISION RUN (10000 shots)")
    print("="*60)
    
    probs, counts = run_experiment(10000)
    
    results = {
        "backend": "guppy/selene",
        "n_qubits": 2,
        "marked_state": "11",
        "k_iterations": 1,
        "shots": 10000,
        "probabilities": probs,
        "success_rate": probs.get("11", 0),
        "theoretical_success": 1.0  # For n=2, k=1
    }
    
    print(f"\nFinal Results:")
    print(json.dumps(results, indent=2))
    
    # Save results
    with open("quantum/guppy/results/guppy_grover_final.json", "w") as f:
        json.dump(results, f, indent=2)
    
    print(f"\n✅ Results saved to quantum/guppy/results/guppy_grover_final.json")
    
    # Compare with theory
    print(f"\n📊 Performance Analysis:")
    print(f"  Theoretical success rate: 100.0%")
    print(f"  Achieved success rate:    {results['success_rate']:.1%}")
    print(f"  Difference:               {abs(1.0 - results['success_rate']):.1%}")

if __name__ == "__main__":
    main()