"""
SMT-Based Verification of Magic State Distillation Circuits
============================================================

Tests whether Z3 can formally verify the error properties of
the 15-to-1 magic state distillation protocol, and whether
the approach scales to concatenated rounds.

The 15-to-1 protocol:
  - 15 noisy T-states in, 1 cleaner T-state out
  - Based on the [[15,1,3]] Reed-Muller code
  - Clifford circuit + T-state injection
  - Input error p -> output error ~35 p^3

The SMT approach:
  - Each qubit's error is a symbolic boolean triple (X, Z, Y=XZ)
  - Clifford gates permute Pauli errors deterministically
  - Stabilizer measurements detect certain error patterns
  - Z3 enumerates undetectable error patterns
  - If this scales polynomially in circuit size, the gap is closed
  - If it blows up exponentially, the gap is confirmed

Requirements: z3-solver (pip install z3-solver)
"""

import time
import json
from z3 import *


def create_error_vars(n_qubits, prefix="e"):
    """Create symbolic error variables for n qubits.
    
    Each qubit has two boolean error variables:
      X error (bit flip) and Z error (phase flip)
      Y error = X AND Z (both)
    """
    x_errors = [Bool(f"{prefix}_X_{i}") for i in range(n_qubits)]
    z_errors = [Bool(f"{prefix}_Z_{i}") for i in range(n_qubits)]
    return x_errors, z_errors


def apply_cnot(x_errors, z_errors, control, target):
    """CNOT propagates X errors forward and Z errors backward.
    
    X_target' = X_target XOR X_control
    Z_control' = Z_control XOR Z_target
    """
    new_x = list(x_errors)
    new_z = list(z_errors)
    new_x[target] = Xor(x_errors[target], x_errors[control])
    new_z[control] = Xor(z_errors[control], z_errors[target])
    return new_x, new_z


def apply_hadamard(x_errors, z_errors, qubit):
    """Hadamard swaps X and Z errors."""
    new_x = list(x_errors)
    new_z = list(z_errors)
    new_x[qubit] = z_errors[qubit]
    new_z[qubit] = x_errors[qubit]
    return new_x, new_z


def apply_phase(x_errors, z_errors, qubit):
    """S gate: X -> Y = XZ, Z -> Z.
    In error tracking: Z' = Z XOR X (since SXS^dag = Y = iXZ)
    """
    new_x = list(x_errors)
    new_z = list(z_errors)
    new_z[qubit] = Xor(z_errors[qubit], x_errors[qubit])
    return new_x, new_z


def reed_muller_15_1_3_stabilizers():
    """
    The [[15,1,3]] quantum Reed-Muller code stabilizer generators.
    
    This code is the basis of the 15-to-1 magic state distillation.
    It has 14 stabilizer generators (X-type and Z-type).
    
    The X stabilizers (generators of the punctured Reed-Muller code):
    Qubits indexed 1-15, using the standard construction from
    the classical [15,11,3] Hamming code and its even subcode.
    """
    # X-type stabilizers for [[15,1,3]]
    # These are the rows of the X stabilizer matrix
    x_stabs = [
        [0,0,0,0,0,0,0,1,1,1,1,1,1,1,1],  # X on qubits 8-15
        [0,0,0,1,1,1,1,0,0,0,0,1,1,1,1],  # X on qubits 4-7,12-15
        [0,1,1,0,0,1,1,0,0,1,1,0,0,1,1],  # X on qubits 2,3,6,7,10,11,14,15
        [1,0,1,0,1,0,1,0,1,0,1,0,1,0,1],  # X on qubits 1,3,5,7,9,11,13,15
        [0,0,0,0,0,0,0,0,0,0,0,1,1,1,1],  # 
        [0,0,0,0,0,0,0,0,0,1,1,0,0,1,1],  #
        [0,0,0,0,0,0,0,0,1,0,1,0,1,0,1],  #
    ]
    
    # Z-type stabilizers (same pattern)
    z_stabs = [
        [0,0,0,0,0,0,0,1,1,1,1,1,1,1,1],
        [0,0,0,1,1,1,1,0,0,0,0,1,1,1,1],
        [0,1,1,0,0,1,1,0,0,1,1,0,0,1,1],
        [1,0,1,0,1,0,1,0,1,0,1,0,1,0,1],
        [0,0,0,0,0,0,0,0,0,0,0,1,1,1,1],
        [0,0,0,0,0,0,0,0,0,1,1,0,0,1,1],
        [0,0,0,0,0,0,0,0,1,0,1,0,1,0,1],
    ]
    
    return x_stabs, z_stabs


def build_syndrome_constraints(x_errors, z_errors, x_stabs, z_stabs):
    """
    Build Z3 constraints for syndrome measurements.
    
    An X stabilizer detects Z errors: syndrome bit = XOR of Z errors
    on the stabilizer's support.
    
    A Z stabilizer detects X errors: syndrome bit = XOR of X errors
    on the stabilizer's support.
    
    For distillation to succeed, ALL syndromes must be 0 (no detection).
    We want to find error patterns that are UNDETECTABLE.
    """
    syndrome_zero = []
    
    # X stabilizers detect Z errors
    for stab in x_stabs:
        support = [i for i, s in enumerate(stab) if s == 1]
        if support:
            syn = z_errors[support[0]]
            for j in support[1:]:
                syn = Xor(syn, z_errors[j])
            syndrome_zero.append(syn == False)
    
    # Z stabilizers detect X errors
    for stab in z_stabs:
        support = [i for i, s in enumerate(stab) if s == 1]
        if support:
            syn = x_errors[support[0]]
            for j in support[1:]:
                syn = Xor(syn, x_errors[j])
            syndrome_zero.append(syn == False)
    
    return syndrome_zero


def count_undetectable_errors(n_qubits, x_stabs, z_stabs, max_weight=None):
    """
    Use Z3 to enumerate all undetectable error patterns.
    
    An undetectable error is one where all syndromes are 0
    but the error is nontrivial (not the identity).
    
    For the [[15,1,3]] code with distance 3, the minimum weight
    undetectable error has weight 3. We expect:
    - Weight 1: 0 undetectable (all detected)
    - Weight 2: 0 undetectable (all detected)  
    - Weight 3: some undetectable (logical operators)
    
    Returns count of undetectable errors at each weight.
    """
    x_errors, z_errors = create_error_vars(n_qubits)
    
    solver = Solver()
    
    # Syndrome = 0 constraints
    syn_constraints = build_syndrome_constraints(x_errors, z_errors, x_stabs, z_stabs)
    for c in syn_constraints:
        solver.add(c)
    
    # At least one error (not the identity)
    any_error = Or(*x_errors, *z_errors)
    solver.add(any_error)
    
    # Optional weight constraint
    if max_weight is not None:
        weight_vars = []
        for i in range(n_qubits):
            has_error = Or(x_errors[i], z_errors[i])
            weight_vars.append(If(has_error, 1, 0))
        total_weight = Sum(weight_vars)
        solver.add(total_weight <= max_weight)
    
    # Enumerate solutions
    solutions = []
    count = 0
    start = time.time()
    
    while solver.check() == sat and count < 10000:  # cap at 10k
        model = solver.model()
        count += 1
        
        # Extract the error pattern
        pattern = {}
        weight = 0
        for i in range(n_qubits):
            xe = is_true(model[x_errors[i]])
            ze = is_true(model[z_errors[i]])
            if xe or ze:
                ptype = 'Y' if (xe and ze) else ('X' if xe else 'Z')
                pattern[i] = ptype
                weight += 1
        
        solutions.append({'weight': weight, 'pattern': pattern})
        
        # Block this solution
        block = Or(*[x_errors[i] != model[x_errors[i]] for i in range(n_qubits)] +
                    [z_errors[i] != model[z_errors[i]] for i in range(n_qubits)])
        solver.add(block)
    
    elapsed = time.time() - start
    
    # Count by weight
    weight_counts = {}
    for s in solutions:
        w = s['weight']
        weight_counts[w] = weight_counts.get(w, 0) + 1
    
    return {
        'total': count,
        'by_weight': weight_counts,
        'elapsed': elapsed,
        'capped': count >= 10000,
    }


def test_distillation_round(round_num, n_input_states):
    """
    Test SMT verification of one round of 15-to-1 distillation.
    
    Round 1: 15 input T-states, 1 output
    Round 2: 15 * 15 = 225 input T-states, 15 intermediate, 1 output
    
    We model each round's stabilizer check as an independent
    SMT problem and measure solve time.
    """
    print(f"\n  Round {round_num}: {n_input_states} input T-states")
    
    x_stabs, z_stabs = reed_muller_15_1_3_stabilizers()
    n_qubits = 15  # per distillation block
    
    if round_num == 1:
        # Single block: 15 qubits, 14 stabilizers
        start = time.time()
        x_errors, z_errors = create_error_vars(n_qubits, "r1")
        
        solver = Solver()
        syn = build_syndrome_constraints(x_errors, z_errors, x_stabs, z_stabs)
        for c in syn:
            solver.add(c)
        solver.add(Or(*x_errors, *z_errors))  # nontrivial error
        
        # Check satisfiability
        result = solver.check()
        elapsed = time.time() - start
        
        print(f"    Qubits: {n_qubits}")
        print(f"    Stabilizers: {len(x_stabs) + len(z_stabs)}")
        print(f"    SMT variables: {2 * n_qubits}")
        print(f"    SAT result: {result}")
        print(f"    Time: {elapsed:.4f}s")
        
        return {'round': 1, 'qubits': n_qubits, 'time': elapsed, 'sat': str(result)}
    
    elif round_num == 2:
        # 15 parallel blocks feeding into one final block
        # Total qubits: 15 * 15 = 225 (but only 15 per block)
        # We model this as 15 independent SMT checks + 1 final check
        # The question: does the COMPOSITION blow up?
        
        # First: 15 independent blocks
        block_times = []
        for b in range(15):
            start = time.time()
            x_e, z_e = create_error_vars(n_qubits, f"r2_b{b}")
            s = Solver()
            syn = build_syndrome_constraints(x_e, z_e, x_stabs, z_stabs)
            for c in syn:
                s.add(c)
            s.add(Or(*x_e, *z_e))
            s.check()
            block_times.append(time.time() - start)
        
        # Then: the final block with correlated inputs
        # The output errors of the 15 blocks become the input errors
        # of the final block. This is where composition matters.
        start = time.time()
        
        # Create error vars for all 15 blocks + final block
        all_x = []
        all_z = []
        s_final = Solver()
        
        for b in range(15):
            x_e, z_e = create_error_vars(n_qubits, f"r2c_b{b}")
            all_x.append(x_e)
            all_z.append(z_e)
            
            # Each block's syndrome must be 0 (passed distillation)
            syn = build_syndrome_constraints(x_e, z_e, x_stabs, z_stabs)
            for c in syn:
                s_final.add(c)
        
        # The output error of each block is the logical error
        # (undetectable by the block's stabilizers).
        # This becomes the input error for the final block.
        # The logical X operator of [[15,1,3]] is X on all 15 qubits.
        # The logical Z operator is Z on all 15 qubits.
        
        # Final block error vars
        final_x, final_z = create_error_vars(n_qubits, "final")
        
        # Connect: final block's input error on qubit b =
        # logical error of block b
        # Logical X = X on all qubits -> detected unless weight >= 3
        # For simplicity: the output error of block b is a single
        # boolean (did a logical error occur in block b?)
        logical_errors = []
        for b in range(15):
            # A logical error in block b means: the block had an
            # undetectable error of weight >= 3 that acts as a
            # logical operator. We model this as: at least one
            # of the block's qubits has an error AND syndrome = 0
            block_has_error = Or(*all_x[b], *all_z[b])
            logical_errors.append(block_has_error)
            
            # Connect to final block input
            s_final.add(final_x[b] == logical_errors[b])
        
        # Final block syndrome must also be 0
        syn_final = build_syndrome_constraints(final_x, final_z, x_stabs, z_stabs)
        for c in syn_final:
            s_final.add(c)
        
        # At least one error somewhere
        any_error = Or(*[Or(*all_x[b], *all_z[b]) for b in range(15)],
                       *final_x, *final_z)
        s_final.add(any_error)
        
        result = s_final.check()
        compose_time = time.time() - start
        
        total_vars = 2 * n_qubits * 16  # 15 blocks + 1 final
        total_constraints = (len(x_stabs) + len(z_stabs)) * 16 + 15  # syndromes + connections
        
        print(f"    Total qubits modeled: {n_qubits * 16}")
        print(f"    SMT variables: {total_vars}")
        print(f"    SMT constraints: ~{total_constraints}")
        print(f"    Block times: min={min(block_times):.4f}s, max={max(block_times):.4f}s, avg={sum(block_times)/len(block_times):.4f}s")
        print(f"    Composition time: {compose_time:.4f}s")
        print(f"    Total time: {sum(block_times) + compose_time:.4f}s")
        print(f"    Composition SAT: {result}")
        
        return {
            'round': 2,
            'qubits': n_qubits * 16,
            'block_time_avg': sum(block_times) / len(block_times),
            'compose_time': compose_time,
            'total_time': sum(block_times) + compose_time,
            'sat': str(result),
            'variables': total_vars,
        }
    
    elif round_num == 3:
        # 15^2 = 225 blocks at level 1, 15 at level 2, 1 at level 3
        # Total: 225 + 15 + 1 = 241 blocks, each 15 qubits
        # Total qubits: 3615
        # This is where we test if it blows up
        
        n_blocks = 225 + 15 + 1
        total_qubits = n_blocks * n_qubits
        total_vars = 2 * total_qubits
        
        start = time.time()
        
        s = Solver()
        s.set("timeout", 60000)  # 60 second timeout
        
        all_blocks_x = []
        all_blocks_z = []
        
        # Level 1: 225 independent blocks
        for b in range(225):
            x_e, z_e = create_error_vars(n_qubits, f"L1_{b}")
            all_blocks_x.append(x_e)
            all_blocks_z.append(z_e)
            syn = build_syndrome_constraints(x_e, z_e, x_stabs, z_stabs)
            for c in syn:
                s.add(c)
        
        # Level 2: 15 blocks, each fed by 15 level-1 blocks
        for b in range(15):
            x_e, z_e = create_error_vars(n_qubits, f"L2_{b}")
            all_blocks_x.append(x_e)
            all_blocks_z.append(z_e)
            syn = build_syndrome_constraints(x_e, z_e, x_stabs, z_stabs)
            for c in syn:
                s.add(c)
            
            # Connect level 1 outputs to level 2 inputs
            for q in range(15):
                l1_block = b * 15 + q
                l1_has_error = Or(*all_blocks_x[l1_block], *all_blocks_z[l1_block])
                s.add(x_e[q] == l1_has_error)
        
        # Level 3: 1 final block
        final_x, final_z = create_error_vars(n_qubits, "L3")
        syn = build_syndrome_constraints(final_x, final_z, x_stabs, z_stabs)
        for c in syn:
            s.add(c)
        
        for q in range(15):
            l2_block = 225 + q
            l2_has_error = Or(*all_blocks_x[l2_block], *all_blocks_z[l2_block])
            s.add(final_x[q] == l2_has_error)
        
        # Nontrivial error
        any_error_parts = [Or(*all_blocks_x[b], *all_blocks_z[b]) for b in range(len(all_blocks_x))]
        any_error_parts.extend([Or(*final_x, *final_z)])
        s.add(Or(*any_error_parts))
        
        result = s.check()
        elapsed = time.time() - start
        
        print(f"    Total blocks: {n_blocks}")
        print(f"    Total qubits modeled: {total_qubits}")
        print(f"    SMT variables: {total_vars}")
        print(f"    Time: {elapsed:.4f}s")
        print(f"    SAT: {result}")
        
        if str(result) == 'unknown':
            print(f"    TIMEOUT - solver could not decide in 60s")
        
        return {
            'round': 3,
            'blocks': n_blocks,
            'qubits': total_qubits,
            'variables': total_vars,
            'time': elapsed,
            'sat': str(result),
        }


def main():
    print("=" * 70)
    print("SMT VERIFICATION OF MAGIC STATE DISTILLATION")
    print("Testing Z3 scalability across concatenation rounds")
    print("=" * 70)
    
    # Phase 1: Verify the [[15,1,3]] code properties
    print("\n--- Phase 1: Code Properties ---")
    x_stabs, z_stabs = reed_muller_15_1_3_stabilizers()
    print(f"  [[15,1,3]] code: 15 qubits, {len(x_stabs)} X-stabilizers, {len(z_stabs)} Z-stabilizers")
    
    # Count undetectable errors by weight
    print("\n  Enumerating undetectable errors...")
    for max_w in [1, 2, 3]:
        result = count_undetectable_errors(15, x_stabs, z_stabs, max_weight=max_w)
        print(f"    Weight <= {max_w}: {result['total']} undetectable, {result['elapsed']:.4f}s")
        if result['by_weight']:
            for w, c in sorted(result['by_weight'].items()):
                print(f"      weight {w}: {c}")
    
    # Phase 2: Scalability test across concatenation rounds
    print(f"\n{'='*70}")
    print("--- Phase 2: Concatenation Scaling ---")
    print(f"{'='*70}")
    
    results = []
    
    r1 = test_distillation_round(1, 15)
    results.append(r1)
    
    r2 = test_distillation_round(2, 225)
    results.append(r2)
    
    r3 = test_distillation_round(3, 3375)
    results.append(r3)
    
    # Phase 3: Summary
    print(f"\n{'='*70}")
    print("SCALING SUMMARY")
    print(f"{'='*70}")
    print(f"  {'Round':>6s}  {'Qubits':>8s}  {'Variables':>10s}  {'Time (s)':>10s}  {'Scaling':>10s}")
    
    prev_time = None
    for r in results:
        t = r.get('total_time', r.get('time', 0))
        v = r.get('variables', 2 * r.get('qubits', 0))
        scaling = ""
        if prev_time and prev_time > 0 and t > 0:
            ratio = t / prev_time
            scaling = f"{ratio:.1f}x"
        print(f"  {r['round']:>6d}  {r.get('qubits', 0):>8d}  {v:>10d}  {t:>10.4f}  {scaling:>10s}")
        prev_time = t
    
    print(f"\n{'='*70}")
    print("INTERPRETATION")
    print(f"{'='*70}")
    print("If time scales polynomially (ratio roughly constant):")
    print("  SMT verification of magic state distillation is tractable.")
    print("  Hugh's gap (exponential blowup in entanglement tracking)")
    print("  does NOT occur with this approach.")
    print()
    print("If time scales exponentially (ratio grows with round):")
    print("  The exponential blowup IS real.")
    print("  SMT solvers cannot handle the entanglement structure")
    print("  of concatenated magic state circuits.")
    print(f"{'='*70}")
    
    # Save results
    with open('smt_distillation_results.json', 'w') as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to smt_distillation_results.json")


if __name__ == '__main__':
    main()