"""
SMT Magic State Distillation - Extended Verification
=====================================================

Three additional tests:
1. Circuit-level noise: errors during the Clifford gates, not just input
2. Output error rate verification: does the model reproduce 35p^3?
3. Round 5: ~810k qubits, ~1.6M variables - definitive scaling test

Requirements: z3-solver
"""

import time
import json
import itertools
from z3 import *


def create_error_vars(n_qubits, prefix="e"):
    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 reed_muller_stabilizers():
    x_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],
    ]
    z_stabs = list(x_stabs)
    return x_stabs, z_stabs


def build_syndrome_constraints(x_errors, z_errors, x_stabs, z_stabs):
    constraints = []
    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])
            constraints.append(syn == False)
    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])
            constraints.append(syn == False)
    return constraints


# ===========================================================
# TEST 1: Circuit-level noise
# ===========================================================

def test_circuit_level_noise():
    """
    Model errors occurring DURING the distillation circuit,
    not just on the input T-states.
    
    The 15-to-1 distillation circuit consists of ~30 CNOT gates
    applied to 15 qubits. We inject a symbolic error after each
    CNOT and check if Z3 still handles it.
    
    For the [[15,1,3]] code, the encoding circuit uses CNOTs
    to create the encoded state. We model errors after each gate
    as additional symbolic Pauli variables that compose with
    the qubit's existing error.
    """
    print("\n" + "=" * 70)
    print("TEST 1: Circuit-Level Noise")
    print("=" * 70)
    
    x_stabs, z_stabs = reed_muller_stabilizers()
    n_qubits = 15
    
    # Model: each qubit has input error + gate errors from each CNOT it participates in
    # The [[15,1,3]] encoding uses roughly 30 CNOTs
    # Each CNOT can introduce an error on control and target
    
    # Simplified CNOT schedule for [[15,1,3]] (representative, not exact)
    cnot_schedule = [
        (0, 1), (2, 3), (4, 5), (6, 7), (8, 9), (10, 11), (12, 13),
        (0, 2), (1, 3), (4, 6), (5, 7), (8, 10), (9, 11), (12, 14),
        (0, 4), (1, 5), (2, 6), (3, 7), (8, 12), (9, 13), (10, 14),
        (0, 8), (1, 9), (2, 10), (3, 11), (4, 12), (5, 13), (6, 14),
    ]
    
    n_gates = len(cnot_schedule)
    
    for n_blocks in [1, 16, 241]:
        start = time.time()
        s = Solver()
        s.set("timeout", 120000)
        
        all_error_vars = []
        
        for blk in range(n_blocks):
            # Input errors
            x_in, z_in = create_error_vars(n_qubits, f"b{blk}_in")
            
            # Gate errors: each CNOT introduces potential errors
            gate_x_errors = []
            gate_z_errors = []
            for g in range(n_gates):
                gx = [Bool(f"b{blk}_g{g}_X_{i}") for i in range(n_qubits)]
                gz = [Bool(f"b{blk}_g{g}_Z_{i}") for i in range(n_qubits)]
                gate_x_errors.append(gx)
                gate_z_errors.append(gz)
                
                # Most gate errors are False (no error on qubits not in this gate)
                ctrl, tgt = cnot_schedule[g]
                for q in range(n_qubits):
                    if q != ctrl and q != tgt:
                        s.add(gx[q] == False)
                        s.add(gz[q] == False)
            
            # Compose all errors: total error = input XOR gate1 XOR gate2 XOR ...
            total_x = list(x_in)
            total_z = list(z_in)
            for g in range(n_gates):
                for q in range(n_qubits):
                    total_x[q] = Xor(total_x[q], gate_x_errors[g][q])
                    total_z[q] = Xor(total_z[q], gate_z_errors[g][q])
            
            # Syndrome check on the composed error
            syn = build_syndrome_constraints(total_x, total_z, x_stabs, z_stabs)
            for c in syn:
                s.add(c)
            
            all_error_vars.extend(x_in + z_in)
            for g in range(n_gates):
                all_error_vars.extend(gate_x_errors[g] + gate_z_errors[g])
        
        # Nontrivial error
        s.add(Or(*all_error_vars))
        
        n_vars = len(all_error_vars)
        
        result = s.check()
        elapsed = time.time() - start
        
        print(f"\n  Blocks: {n_blocks}, Gates/block: {n_gates}")
        print(f"  Variables: {n_vars}")
        print(f"  Time: {elapsed:.4f}s")
        print(f"  Result: {result}")
        if n_vars > 0:
            print(f"  us/var: {elapsed * 1e6 / n_vars:.1f}")


# ===========================================================
# TEST 2: Output error rate verification
# ===========================================================

def test_output_error_rate():
    """
    Verify the distillation guarantee: input error p -> output ~35 p^3.
    
    Count the number of weight-w undetectable error patterns.
    The output error probability is:
      p_out = sum over w>=3 of (n_undetectable(w) * p^w * (1-p)^(15-w))
    
    For small p: p_out ~ n_undetectable(3) * p^3
    The known result is p_out ~ 35 p^3.
    So we should find n_undetectable(3) = 35 ... but we found 21 earlier.
    
    The discrepancy: our earlier count was for the stabilizer code alone.
    The distillation protocol also has X-type and Z-type logical errors
    that must be counted separately. Let's do a careful enumeration.
    """
    print("\n" + "=" * 70)
    print("TEST 2: Output Error Rate Verification")
    print("=" * 70)
    
    x_stabs, z_stabs = reed_muller_stabilizers()
    n_qubits = 15
    
    # Enumerate ALL undetectable errors up to weight 4
    for max_w in range(1, 5):
        x_errors, z_errors = create_error_vars(n_qubits, f"w{max_w}")
        
        s = Solver()
        
        syn = build_syndrome_constraints(x_errors, z_errors, x_stabs, z_stabs)
        for c in syn:
            s.add(c)
        
        # Nontrivial
        s.add(Or(*x_errors, *z_errors))
        
        # Weight constraint
        weight_terms = []
        for i in range(n_qubits):
            has_err = Or(x_errors[i], z_errors[i])
            weight_terms.append(If(has_err, 1, 0))
        s.add(Sum(weight_terms) <= max_w)
        
        # Count by enumerating
        count = 0
        x_only = 0
        z_only = 0
        y_count = 0
        mixed = 0
        start = time.time()
        
        while s.check() == sat and count < 50000:
            m = s.model()
            count += 1
            
            has_x = any(is_true(m[x_errors[i]]) for i in range(n_qubits))
            has_z = any(is_true(m[z_errors[i]]) for i in range(n_qubits))
            
            if has_x and not has_z:
                x_only += 1
            elif has_z and not has_x:
                z_only += 1
            else:
                mixed += 1
            
            # Block this solution
            block = Or(*[x_errors[i] != m[x_errors[i]] for i in range(n_qubits)] +
                        [z_errors[i] != m[z_errors[i]] for i in range(n_qubits)])
            s.add(block)
        
        elapsed = time.time() - start
        print(f"\n  Weight <= {max_w}: {count} undetectable errors ({elapsed:.3f}s)")
        print(f"    X-only: {x_only}, Z-only: {z_only}, mixed(XZ): {mixed}")
        
        if max_w == 3:
            # The distillation output error coefficient
            # Each undetectable error of weight w contributes to output error
            # For the 15-to-1 protocol, the coefficient is C(15,w) * f(w)
            # where f(w) accounts for which errors are logical
            print(f"\n    The distillation protocol output error:")
            print(f"    p_out ~ {count} * p^3 for weight-3 errors")
            print(f"    (Known result: ~35 p^3)")
            if count > 0:
                print(f"    Our coefficient: {count}")
                print(f"    Note: exact coefficient depends on which errors")
                print(f"    are logical vs trivial (in the stabilizer group)")


# ===========================================================
# TEST 3: Round 5 scaling
# ===========================================================

def test_round_5():
    """
    Round 5: 15^4 + 15^3 + 15^2 + 15 + 1 = 54,241 blocks
    Total qubits: 813,615
    SMT variables: 1,627,230
    """
    print("\n" + "=" * 70)
    print("TEST 3: Round 5 Scaling")
    print("=" * 70)
    
    x_stabs, z_stabs = reed_muller_stabilizers()
    n_qubits = 15
    
    levels = 5
    blocks_per_level = [15**(levels-1-lev) for lev in range(levels)]
    total_blocks = sum(blocks_per_level)
    total_qubits = total_blocks * n_qubits
    total_vars = 2 * total_qubits
    
    print(f"  Levels: {levels}")
    print(f"  Blocks per level: {blocks_per_level}")
    print(f"  Total blocks: {total_blocks}")
    print(f"  Total qubits: {total_qubits}")
    print(f"  SMT variables: {total_vars}")
    print(f"  Building constraints...", flush=True)
    
    start = time.time()
    
    s = Solver()
    s.set("timeout", 600000)  # 10 min
    
    all_x = {}
    all_z = {}
    block_count = 0
    
    for lev in range(levels):
        n_blocks = blocks_per_level[lev]
        for b in range(n_blocks):
            x_e, z_e = create_error_vars(n_qubits, f"L{lev}_{b}")
            all_x[(lev, b)] = x_e
            all_z[(lev, b)] = z_e
            
            syn = build_syndrome_constraints(x_e, z_e, x_stabs, z_stabs)
            for c in syn:
                s.add(c)
            
            if lev > 0:
                for q in range(15):
                    prev_block = b * 15 + q
                    if (lev - 1, prev_block) in all_x:
                        prev_has_error = Or(*all_x[(lev-1, prev_block)], *all_z[(lev-1, prev_block)])
                        s.add(x_e[q] == prev_has_error)
            
            block_count += 1
            if block_count % 5000 == 0:
                elapsed = time.time() - start
                print(f"    {block_count}/{total_blocks} blocks ({elapsed:.1f}s)", flush=True)
    
    build_time = time.time() - start
    print(f"  Build complete: {build_time:.1f}s")
    
    # Nontrivial error
    any_parts = []
    for key in all_x:
        any_parts.append(Or(*all_x[key], *all_z[key]))
    s.add(Or(*any_parts))
    
    print(f"  Solving...", flush=True)
    solve_start = time.time()
    result = s.check()
    solve_time = time.time() - solve_start
    total_time = time.time() - start
    
    print(f"  Solve time: {solve_time:.4f}s")
    print(f"  Total time: {total_time:.1f}s")
    print(f"  Result: {result}")
    
    if str(result) == 'unknown':
        print(f"  *** TIMEOUT ***")
    else:
        print(f"  us/var (solve only): {solve_time * 1e6 / total_vars:.1f}")
    
    return {
        'round': 5,
        'blocks': total_blocks,
        'qubits': total_qubits,
        'variables': total_vars,
        'build_time': build_time,
        'solve_time': solve_time,
        'total_time': total_time,
        'sat': str(result),
    }


def main():
    print("=" * 70)
    print("SMT MAGIC STATE DISTILLATION - EXTENDED VERIFICATION")
    print("=" * 70)
    
    # Test 1: Circuit-level noise
    test_circuit_level_noise()
    
    # Test 2: Output error rate
    test_output_error_rate()
    
    # Test 3: Round 5
    r5 = test_round_5()
    
    print(f"\n{'='*70}")
    print("ALL TESTS COMPLETE")
    print(f"{'='*70}")
    
    results = {'round5': r5}
    with open('smt_extended_results.json', 'w') as f:
        json.dump(results, f, indent=2)


if __name__ == '__main__':
    main()