"""
SMT Magic State Distillation - Round 4 Extension
Tests 54,240 qubits / 108,480 SMT variables
"""

import time
import json
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_15_1_3_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)  # same pattern for CSS code
    return x_stabs, z_stabs


def build_syndrome_constraints(x_errors, z_errors, x_stabs, z_stabs):
    syndrome_zero = []
    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)
    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 test_round(round_num, timeout_sec=300):
    """Test a specific concatenation round."""
    
    x_stabs, z_stabs = reed_muller_15_1_3_stabilizers()
    n_qubits = 15
    
    # Calculate structure
    levels = round_num
    blocks_per_level = []
    total_blocks = 0
    for lev in range(levels):
        n = 15 ** (levels - 1 - lev)
        blocks_per_level.append(n)
        total_blocks += n
    
    total_qubits = total_blocks * n_qubits
    total_vars = 2 * total_qubits
    
    print(f"\n  Round {round_num}")
    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"    Timeout: {timeout_sec}s")
    print(f"    Building constraints...", end="", flush=True)
    
    start = time.time()
    
    s = Solver()
    s.set("timeout", timeout_sec * 1000)
    
    # Store all block error vars indexed by (level, block_index)
    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)
            
            # Connect to previous level (if not level 0)
            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 % 500 == 0:
                print(f" {block_count}...", end="", flush=True)
    
    build_time = time.time() - start
    print(f" done ({build_time:.1f}s)")
    
    # Nontrivial error somewhere
    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...", end="", flush=True)
    solve_start = time.time()
    result = s.check()
    solve_time = time.time() - solve_start
    total_time = time.time() - start
    
    print(f" done")
    print(f"    Build time: {build_time:.4f}s")
    print(f"    Solve time: {solve_time:.4f}s")
    print(f"    Total time: {total_time:.4f}s")
    print(f"    Result: {result}")
    
    if str(result) == 'unknown':
        print(f"    *** TIMEOUT at {timeout_sec}s ***")
    
    return {
        'round': round_num,
        '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 SCALING TEST")
    print("=" * 70)
    
    results = []
    
    for r in [1, 2, 3, 4]:
        timeout = 300 if r <= 3 else 600  # 10 min timeout for round 4
        res = test_round(r, timeout_sec=timeout)
        results.append(res)
    
    print(f"\n{'='*70}")
    print("SCALING SUMMARY")
    print(f"{'='*70}")
    print(f"  {'Round':>6s}  {'Blocks':>8s}  {'Qubits':>8s}  {'Vars':>10s}  {'Solve(s)':>10s}  {'Total(s)':>10s}  {'Ratio':>8s}")
    
    prev_total = None
    for r in results:
        ratio = ""
        if prev_total and prev_total > 0:
            ratio = f"{r['total_time']/prev_total:.1f}x"
        print(f"  {r['round']:>6d}  {r['blocks']:>8d}  {r['qubits']:>8d}  {r['variables']:>10d}  {r['solve_time']:>10.4f}  {r['total_time']:>10.4f}  {ratio:>8s}")
        prev_total = r['total_time']
    
    print(f"\n{'='*70}")
    if all(r['sat'] == 'sat' for r in results):
        print("All rounds solved. SMT verification scales across 4 concatenation levels.")
        print("The exponential blowup in entanglement tracking does NOT occur")
        print("when using Pauli frame error propagation through stabilizer constraints.")
    elif any(r['sat'] == 'unknown' for r in results):
        print("Timeout detected. The scaling hits a wall at this concatenation depth.")
    print(f"{'='*70}")
    
    with open('smt_distillation_round4.json', 'w') as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to smt_distillation_round4.json")


if __name__ == '__main__':
    main()