"""
SMT Magic State Distillation - CORRECTED Stabilizers
=====================================================

The [[15,1,3]] quantum Reed-Muller code is NOT symmetric.
It is a CSS code built from punctured Reed-Muller codes:
  - X stabilizers from punctured RM(1,4): the first-order RM code
  - Z stabilizers from punctured RM(0,4)^perp

The code has transversal T gate, which requires this asymmetry.

Minimum weight logical operators:
  - Logical Z: weight 3 (35 total) -- these are the undetectable errors
  - Logical X: weight 7 or higher

The previous version used identical X and Z stabilizer matrices,
which gave 7 X + 7 Z + 7 XZ = 21 undetectable weight-3 errors.
The correct answer is 35 Z-only weight-3 errors.

This script fixes the stabilizers and re-runs all scaling tests.
"""

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 qrm_15_1_3_stabilizers():
    """
    Correct stabilizers for the [[15,1,3]] quantum Reed-Muller code.
    
    The 15 qubits are indexed by nonzero elements of GF(2)^4.
    We label them 1-15 in binary (0001 through 1111), mapped to indices 0-14.
    
    X stabilizers: from the punctured first-order Reed-Muller code RM(1,4).
    RM(1,4) has dimension 5, so punctured RM(1,4) has 4 generators.
    Each generator corresponds to a coordinate function on GF(2)^4.
    Generator i has support on all points where coordinate i is 1.
    
    Z stabilizers: from the dual of the punctured RM(0,4) code.
    Equivalently, from the punctured RM(2,4) code minus RM(1,4).
    RM(2,4) has dimension 11. The Z stabilizers are the 10 generators
    of the punctured RM(2,4) code (11 - 1 for the logical Z).
    
    The generators of RM(2,4) beyond RM(1,4) correspond to products
    of coordinate functions: x_i * x_j for i < j. There are C(4,2) = 6
    such generators, plus the 4 from RM(1,4).
    
    So: 4 X-stabilizer generators, 10 Z-stabilizer generators.
    Total: 14 independent stabilizers for a [[15,1,3]] code (15-1=14). Correct.
    """
    
    # Qubit labels: nonzero elements of GF(2)^4
    # Index i corresponds to binary representation of (i+1)
    # i=0 -> 0001, i=1 -> 0010, ..., i=14 -> 1111
    
    def bit(n, pos):
        """Get bit at position pos of (n+1)."""
        return ((n + 1) >> pos) & 1
    
    # X stabilizers: 4 generators, one per coordinate
    # Generator for coordinate c: support on qubits where bit c is 1
    x_stabs = []
    for c in range(4):
        row = [bit(i, c) for i in range(15)]
        x_stabs.append(row)
    
    # Z stabilizers: 10 generators
    # First 4: same coordinate generators (RM(1,4) part)
    # Next 6: product generators x_i * x_j for i < j (RM(2,4) \ RM(1,4))
    z_stabs = []
    
    # Coordinate generators (same as X)
    for c in range(4):
        row = [bit(i, c) for i in range(15)]
        z_stabs.append(row)
    
    # Product generators
    for c1 in range(4):
        for c2 in range(c1 + 1, 4):
            row = [bit(i, c1) * bit(i, c2) for i in range(15)]
            z_stabs.append(row)
    
    return x_stabs, z_stabs


def build_syndrome_constraints(x_errors, z_errors, x_stabs, z_stabs):
    """
    X stabilizers detect Z errors.
    Z stabilizers detect X errors.
    """
    constraints = []
    
    # 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])
            constraints.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])
            constraints.append(syn == False)
    
    return constraints


def verify_code_properties():
    """Verify the corrected code has the right properties."""
    print("\n--- Code Verification ---")
    
    x_stabs, z_stabs = qrm_15_1_3_stabilizers()
    print(f"  X stabilizer generators: {len(x_stabs)}")
    print(f"  Z stabilizer generators: {len(z_stabs)}")
    print(f"  Total: {len(x_stabs) + len(z_stabs)} (expected 14)")
    
    # Print stabilizer weights
    for i, s in enumerate(x_stabs):
        print(f"  X_{i}: weight {sum(s)}, support {[j for j,v in enumerate(s) if v]}")
    for i, s in enumerate(z_stabs):
        print(f"  Z_{i}: weight {sum(s)}, support {[j for j,v in enumerate(s) if v]}")
    
    # Enumerate undetectable errors by weight and type
    for max_w in range(1, 5):
        x_errors, z_errors = create_error_vars(15, f"v{max_w}")
        s = Solver()
        
        syn = build_syndrome_constraints(x_errors, z_errors, x_stabs, z_stabs)
        for c in syn:
            s.add(c)
        
        s.add(Or(*x_errors, *z_errors))
        
        weight_terms = []
        for i in range(15):
            weight_terms.append(If(Or(x_errors[i], z_errors[i]), 1, 0))
        s.add(Sum(weight_terms) <= max_w)
        
        count = 0
        x_only = 0
        z_only = 0
        mixed = 0
        
        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(15))
            has_z = any(is_true(m[z_errors[i]]) for i in range(15))
            
            if has_x and not has_z:
                x_only += 1
            elif has_z and not has_x:
                z_only += 1
            else:
                mixed += 1
            
            block = Or(*[x_errors[i] != m[x_errors[i]] for i in range(15)] +
                        [z_errors[i] != m[z_errors[i]] for i in range(15)])
            s.add(block)
        
        print(f"\n  Weight <= {max_w}: {count} undetectable")
        print(f"    X-only: {x_only}, Z-only: {z_only}, mixed: {mixed}")


def test_scaling(round_num, timeout_sec=300):
    """Test concatenation scaling with corrected stabilizers."""
    x_stabs, z_stabs = qrm_15_1_3_stabilizers()
    n_qubits = 15
    
    levels = round_num
    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"\n  Round {round_num}: {total_blocks} blocks, {total_qubits} qubits, {total_vars} vars")
    print(f"    Building...", end="", flush=True)
    
    start = time.time()
    s = Solver()
    s.set("timeout", timeout_sec * 1000)
    
    all_x = {}
    all_z = {}
    
    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)
    
    build_time = time.time() - start
    print(f" built ({build_time:.1f}s), solving...", end="", flush=True)
    
    any_parts = [Or(*all_x[k], *all_z[k]) for k in all_x]
    s.add(Or(*any_parts))
    
    solve_start = time.time()
    result = s.check()
    solve_time = time.time() - solve_start
    total_time = time.time() - start
    
    us_per_var = solve_time * 1e6 / total_vars if total_vars > 0 else 0
    
    print(f" done")
    print(f"    Build: {build_time:.3f}s, Solve: {solve_time:.3f}s, Total: {total_time:.3f}s")
    print(f"    us/var: {us_per_var:.1f}")
    print(f"    Result: {result}")
    
    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,
        'us_per_var': us_per_var,
        'sat': str(result),
    }


def test_circuit_level(n_blocks):
    """Circuit-level noise with corrected stabilizers."""
    x_stabs, z_stabs = qrm_15_1_3_stabilizers()
    n_qubits = 15
    
    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)
    
    start = time.time()
    s = Solver()
    s.set("timeout", 120000)
    
    all_vars = []
    
    for blk in range(n_blocks):
        x_in, z_in = create_error_vars(n_qubits, f"b{blk}_in")
        
        gate_x = []
        gate_z = []
        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.append(gx)
            gate_z.append(gz)
            
            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)
        
        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[g][q])
                total_z[q] = Xor(total_z[q], gate_z[g][q])
        
        syn = build_syndrome_constraints(total_x, total_z, x_stabs, z_stabs)
        for c in syn:
            s.add(c)
        
        all_vars.extend(x_in + z_in)
        for g in range(n_gates):
            all_vars.extend(gate_x[g] + gate_z[g])
    
    s.add(Or(*all_vars))
    n_vars = len(all_vars)
    
    result = s.check()
    elapsed = time.time() - start
    us_per_var = elapsed * 1e6 / n_vars if n_vars > 0 else 0
    
    print(f"    Blocks: {n_blocks}, Vars: {n_vars}, Time: {elapsed:.3f}s, us/var: {us_per_var:.1f}, Result: {result}")
    
    return {
        'blocks': n_blocks,
        'variables': n_vars,
        'time': elapsed,
        'us_per_var': us_per_var,
        'sat': str(result),
    }


def main():
    print("=" * 70)
    print("SMT MAGIC STATE DISTILLATION - CORRECTED STABILIZERS")
    print("[[15,1,3]] Quantum Reed-Muller Code")
    print("4 X-stabilizers (RM(1,4)), 10 Z-stabilizers (RM(2,4))")
    print("=" * 70)
    
    # Phase 1: Verify the code
    verify_code_properties()
    
    # Phase 2: Scaling test
    print(f"\n{'='*70}")
    print("CONCATENATION SCALING")
    print(f"{'='*70}")
    
    results = []
    for r in [1, 2, 3, 4, 5]:
        timeout = 600 if r == 5 else 300
        res = test_scaling(r, timeout)
        results.append(res)
    
    # Phase 3: Circuit-level noise
    print(f"\n{'='*70}")
    print("CIRCUIT-LEVEL NOISE")
    print(f"{'='*70}")
    
    cl_results = []
    for nb in [1, 16, 241]:
        res = test_circuit_level(nb)
        cl_results.append(res)
    
    # Summary
    print(f"\n{'='*70}")
    print("SCALING SUMMARY (CORRECTED)")
    print(f"{'='*70}")
    print(f"  {'Round':>6s}  {'Vars':>10s}  {'Solve(s)':>10s}  {'us/var':>8s}")
    for r in results:
        print(f"  {r['round']:>6d}  {r['variables']:>10d}  {r['solve_time']:>10.3f}  {r['us_per_var']:>8.1f}")
    
    print(f"\n  Circuit-level noise:")
    for r in cl_results:
        print(f"  Blocks={r['blocks']:>4d}  Vars={r['variables']:>8d}  Time={r['time']:>8.3f}s  us/var={r['us_per_var']:>6.1f}")
    
    output = {
        'scaling': results,
        'circuit_level': cl_results,
    }
    with open('smt_corrected_results.json', 'w') as f:
        json.dump(output, f, indent=2)
    print(f"\n  Saved to smt_corrected_results.json")


if __name__ == '__main__':
    main()