simulation.py

# """
# simulation.py — Monte Carlo threshold simulation for D(S3).

# Supports two noise models:
#   Code capacity:     T=1 round, p_meas=0 (perfect syndromes).
#   Phenomenological:  T=L rounds, p_meas=p (noisy syndromes).
#                      Uses JIT decoder with commit/defer rule.

# Trial structure:
#   1. Prepare D(S3) ground state.
#   2. Record initial logical observables.
#   3. For T rounds: apply errors, take syndrome measurement, run JIT step.
#   4. Take final NOISELESS syndrome round, run JIT step.
#   5. Finalize: force-commit remaining clusters, decode eta globally.
#   6. Logical error = any change in homological invariants.

# L must not be divisible by 3 (Z3 logical wraps trivially).
# """

# import numpy as np
# import matplotlib
# matplotlib.use('Agg')
# import matplotlib.pyplot as plt
# import time
# from datetime import datetime
# import os

# from lattice import DS3Lattice
# from errors import DepolarizingError
# from decoder import JITDecoder


# def run_trial(L, T, error_model, rng):
#     """
#     Single Monte Carlo trial with JIT decoding.

#     At each timestep: apply errors → measure (with noise) → JIT step.
#     The JIT receives p_meas so D(Z3) measurements within ungauged
#     patches also experience measurement noise.

#     Returns True if logical error occurred.
#     """
#     lattice = DS3Lattice(L)
#     init_obs = lattice.get_logical_observables()
#     decoder = JITDecoder(lattice)
#     p_meas = error_model.p_meas

#     # T noisy rounds with JIT decoding at each step
#     for step in range(T):
#         error_model.apply(lattice, rng)
#         det = lattice.step(p_meas=p_meas, rng=rng)
#         decoder.process_timestep(det, det['t'], p_meas=p_meas, rng=rng)

#     # Final noiseless round
#     det = lattice.step(p_meas=0.0, rng=rng)
#     decoder.process_timestep(det, det['t'], p_meas=0.0, rng=rng)

#     # Finalize: complete patches, force-commit, decode eta
#     decoder.finalize()

#     return lattice.has_logical_error(init_obs)


# def run_sweep(L_values, p_values, n_trials=500, T=1, p_meas_factor=0.0,
#               seed=42, verbose=True):
#     """
#     Sweep over (L, p).

#     Parameters
#     ----------
#     T : int
#         Number of noisy rounds per trial.
#         T=1 for code capacity, T=L for phenomenological.
#     p_meas_factor : float
#         Measurement noise as multiple of p.
#         0.0 = perfect syndromes (code capacity).
#         1.0 = p_meas = p (phenomenological).

#     Returns results[L][p] = (p_L, std_err).
#     """
#     rng = np.random.default_rng(seed)
#     results = {L: {} for L in L_values}
#     total = len(L_values) * len(p_values)
#     done = 0
#     t0 = time.time()
#     for L in L_values:
#         T_actual = T if T > 1 else 1  # allow T=L convention
#         for p in p_values:
#             p_meas = p * p_meas_factor
#             em = DepolarizingError(p, p_meas)
#             n_err = sum(run_trial(L, T_actual, em, rng) for _ in range(n_trials))
#             pL = n_err / n_trials
#             se = np.sqrt(pL * (1 - pL) / n_trials)
#             results[L][p] = (pL, se)
#             done += 1
#             if verbose:
#                 elapsed = time.time() - t0
#                 print(f'  L={L:2d} p={p:.3f} pL={pL:.4f}±{se:.4f} '
#                       f'[{done}/{total}] [{elapsed:.0f}s]', flush=True)
#     return results


# def run_sweep_phenomenological(L_values, p_values, n_trials=500,
#                                 seed=42, verbose=True):
#     """
#     Sweep with phenomenological noise: T=L rounds, p_meas=p.
#     """
#     rng = np.random.default_rng(seed)
#     results = {L: {} for L in L_values}
#     total = len(L_values) * len(p_values)
#     done = 0
#     t0 = time.time()
#     for L in L_values:
#         for p in p_values:
#             em = DepolarizingError(p, p)  # p_meas = p
#             n_err = sum(run_trial(L, L, em, rng) for _ in range(n_trials))
#             pL = n_err / n_trials
#             se = np.sqrt(pL * (1 - pL) / n_trials)
#             results[L][p] = (pL, se)
#             done += 1
#             if verbose:
#                 elapsed = time.time() - t0
#                 print(f'  L={L:2d} T={L} p={p:.3f} pL={pL:.4f}±{se:.4f} '
#                       f'[{done}/{total}] [{elapsed:.0f}s]', flush=True)
#     return results


# def plot_results(results, L_values, p_values, save_path, title=None):
#     fig, ax = plt.subplots(figsize=(10, 7))
#     cmap = plt.get_cmap('tab10')
#     for i, L in enumerate(sorted(L_values)):
#         ps = sorted(results[L].keys())
#         pLs = [results[L][p][0] for p in ps]
#         errs = [results[L][p][1] for p in ps]
#         ax.errorbar(ps, pLs, yerr=errs, marker='o', capsize=3,
#                     linewidth=2, label=f'L = {L}', color=cmap(i))
#     ax.axhline(0.5, color='grey', ls='--', alpha=0.5)
#     ax.set_xlabel('Physical error rate $p$', fontsize=14)
#     ax.set_ylabel('Logical error rate $p_L$', fontsize=14)
#     ax.set_title(title or r'D(S$_3$) Threshold', fontsize=15)
#     ax.legend(fontsize=12)
#     ax.grid(True, alpha=0.3)
#     ax.set_ylim(-0.02, 1.02)
#     plt.tight_layout()
#     plt.savefig(save_path, dpi=150, bbox_inches='tight')
#     print(f'Saved: {save_path}')


# def main():
#     print('=' * 60)
#     print('D(S3) Threshold Simulation — JIT Decoder')
#     print('=' * 60)

#     # Code-capacity mode
#     L_values = [4, 5, 7, 8]
#     p_values = [0.02, 0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18]
#     n_trials = 300
#     # T = 1
#     # p_meas_factor = 0.0

#     # Phenomenological mode
#     T = L  # convention: T=L rounds for phenomenological  
#     p_meas_factor = 0.02  # p_meas = p for phenomenological

#     print(f'Mode: code capacity (T={T}, p_meas={p_meas_factor}*p)')
#     print(f'L: {L_values}')
#     print(f'p: {[f"{p:.2f}" for p in p_values]}')
#     print(f'Trials: {n_trials}')
#     print()

#     t0 = time.time()
#     results = run_sweep(L_values, p_values, n_trials=n_trials,
#                         T=T, p_meas_factor=p_meas_factor,
#                         seed=2024, verbose=True)
#     print(f'\nTotal: {time.time()-t0:.0f}s')

#     script_dir = os.path.dirname(os.path.abspath(__file__))
#     plot_dir = os.path.join(script_dir, 'threshold_plots')
#     os.makedirs(plot_dir, exist_ok=True)
#     # plot_results(results, L_values, p_values,
#     #              os.path.join(plot_dir, 'threshold_plot.png'),
#     #              title=r'D(S$_3$) Code-Capacity Threshold (JIT Decoder)')

#     now = datetime.now()
#     formatted_time = now.strftime("%Y_%m_%d_%H_%M_%S")

#     if T == 1 and p_meas_factor == 0.0:
#         plot_title = r'D(S$_3$) Code-Capacity Threshold (JIT Decoder)'
#         plot_name = f'code_capacity_threshold_{formatted_time}.png'

#     else: 
#         plot_title = r'D(S$_3$) Phenomenological Threshold (JIT Decoder)'
#         plot_name = f'phenomenological_threshold_{formatted_time}.png'


#     plot_results(results, L_values, p_values,
#                  os.path.join(plot_dir, plot_name),
#                  title=plot_title)



#     print('\nSummary:')
#     hdr = f'{"p":>8}' + ''.join(f'  L={L:2d}   ' for L in sorted(L_values))
#     print(hdr)
#     for p in sorted(p_values):
#         row = f'{p:8.3f}'
#         for L in sorted(L_values):
#             row += f'  {results[L][p][0]:.4f}  '
#         print(row)


# if __name__ == '__main__':
#     main()

"""
simulation.py — Monte Carlo threshold simulation for D(S3).

Supports two noise models:
  Code capacity:     T=1 round, p_meas=0 (perfect syndromes).
  Phenomenological:  T=L rounds, p_meas=p (noisy syndromes).
                     Uses JIT decoder with commit/defer rule.

Trial structure:
  1. Prepare D(S3) ground state.
  2. Record initial logical observables.
  3. For T rounds: apply errors, take syndrome measurement, run JIT step.
  4. Take final NOISELESS syndrome round, run JIT step.
  5. Finalize: force-commit remaining clusters, decode eta globally.
  6. Logical error = any change in homological invariants.

L must not be divisible by 3 (Z3 logical wraps trivially).
"""

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import time
from datetime import datetime
import os

from lattice import DS3Lattice
from errors import DepolarizingError
from decoder import JITDecoder


def run_trial(L, T, error_model, rng):
    """
    Single Monte Carlo trial with JIT decoding.

    At each timestep: apply errors → measure (with noise) → JIT step.
    The JIT receives p_meas so D(Z3) measurements within ungauged
    patches also experience measurement noise.

    Returns True if logical error occurred.
    """
    lattice = DS3Lattice(L)
    init_obs = lattice.get_logical_observables()
    decoder = JITDecoder(lattice)
    p_meas = error_model.p_meas

    # T noisy rounds with JIT decoding at each step
    for step in range(T):
        error_model.apply(lattice, rng)
        det = lattice.step(p_meas=p_meas, rng=rng)
        decoder.process_timestep(det, det['t'], p_meas=p_meas, rng=rng)

    # Final noiseless round
    det = lattice.step(p_meas=0.0, rng=rng)
    decoder.process_timestep(det, det['t'], p_meas=0.0, rng=rng)

    # Finalize: complete patches, force-commit, decode eta
    decoder.finalize()

    return lattice.has_logical_error(init_obs)


def run_sweep(L_values, p_values, n_trials=500, T=1, p_meas_factor=0.0,
              seed=42, verbose=True):
    """
    Sweep over (L, p).

    Parameters
    ----------
    T : int
        Number of noisy rounds per trial.
        T=1 for code capacity, T=L for phenomenological.
    p_meas_factor : float
        Measurement noise as multiple of p.
        0.0 = perfect syndromes (code capacity).
        1.0 = p_meas = p (phenomenological).

    Returns results[L][p] = (p_L, std_err).
    """
    rng = np.random.default_rng(seed)
    results = {L: {} for L in L_values}
    total = len(L_values) * len(p_values)
    done = 0
    t0 = time.time()
    for L in L_values:
        T_actual = T if T > 1 else 1  # allow T='L' convention
        for p in p_values:
            p_meas = p * p_meas_factor
            em = DepolarizingError(p, p_meas)
            n_err = sum(run_trial(L, T_actual, em, rng) for _ in range(n_trials))
            pL = n_err / n_trials
            se = np.sqrt(pL * (1 - pL) / n_trials)
            results[L][p] = (pL, se)
            done += 1
            if verbose:
                elapsed = time.time() - t0
                print(f'  L={L:2d} p={p:.3f} pL={pL:.4f}±{se:.4f} '
                      f'[{done}/{total}] [{elapsed:.0f}s]', flush=True)
    return results


def run_sweep_phenomenological(L_values, p_values, n_trials=500,
                                seed=42, verbose=True):
    """
    Sweep with phenomenological noise: T=L rounds, p_meas=p.
    """
    rng = np.random.default_rng(seed)
    results = {L: {} for L in L_values}
    total = len(L_values) * len(p_values)
    done = 0
    t0 = time.time()
    for L in L_values:
        for p in p_values:
            em = DepolarizingError(p, p)  # p_meas = p
            n_err = sum(run_trial(L, L, em, rng) for _ in range(n_trials))
            pL = n_err / n_trials
            se = np.sqrt(pL * (1 - pL) / n_trials)
            results[L][p] = (pL, se)
            done += 1
            if verbose:
                elapsed = time.time() - t0
                print(f'  L={L:2d} T={L} p={p:.3f} pL={pL:.4f}±{se:.4f} '
                      f'[{done}/{total}] [{elapsed:.0f}s]', flush=True)
    return results


def plot_results(results, L_values, p_values, save_path, title=None):
    fig, ax = plt.subplots(figsize=(10, 7))
    cmap = plt.get_cmap('tab10')
    for i, L in enumerate(sorted(L_values)):
        ps = sorted(results[L].keys())
        pLs = [results[L][p][0] for p in ps]
        errs = [results[L][p][1] for p in ps]
        ax.errorbar(ps, pLs, yerr=errs, marker='o', capsize=3,
                    linewidth=2, label=f'L = {L}', color=cmap(i))
    ax.axhline(0.5, color='grey', ls='--', alpha=0.5)
    ax.set_xlabel('Physical error rate $p$', fontsize=14)
    ax.set_ylabel('Logical error rate $p_L$', fontsize=14)
    ax.set_title(title or r'D(S$_3$) Threshold', fontsize=15)
    ax.legend(fontsize=12)
    ax.grid(True, alpha=0.3)
    ax.set_ylim(-0.02, 1.02)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f'Saved: {save_path}')


def main():
    import argparse
    parser = argparse.ArgumentParser(
        description='D(S3) Threshold Simulation — JIT Decoder')
    parser.add_argument('--mode', choices=['code_capacity', 'phenomenological', 'both'],
                        default='both', help='Which threshold to compute')
    parser.add_argument('--trials', type=int, default=500,
                        help='Number of trials per (L, p) point')
    parser.add_argument('--seed', type=int, default=2024)
    args = parser.parse_args()

    script_dir = os.path.dirname(os.path.abspath(__file__))
    plot_dir = os.path.join(script_dir, 'threshold_plots')
    os.makedirs(plot_dir, exist_ok=True)

    now = datetime.now()
    formatted_time = now.strftime("%Y_%m_%d_%H_%M_%S")

    # ── Code capacity ──
    if args.mode in ('code_capacity', 'both'):
        print('=' * 60)
        print('Code-Capacity Threshold (T=1, p_meas=0)')
        print('=' * 60)

        # L_values = [4, 7, 10, 13]
        L_values = [4, 10, 17]
        # p_values = [0.08, 0.10, 0.12, 0.14, 0.16, 0.18, 0.20, 0.22, 0.225, 0.23, 0.235, 0.24]
        p_values = [0.20, 0.21, 0.22, 0.225, 0.23, 0.235, 0.24]
        n_trials = args.trials

        print(f'L: {L_values}')
        print(f'p: {[f"{p:.2f}" for p in p_values]}')
        print(f'Trials: {n_trials}')
        print()

        t0 = time.time()
        results = run_sweep(L_values, p_values, n_trials=n_trials,
                            T=1, p_meas_factor=0.0,
                            seed=args.seed, verbose=True)
        print(f'\nTotal: {time.time()-t0:.0f}s')

        # plot_results(results, L_values, p_values,
        #              os.path.join(plot_dir, 'code_capacity.png'),
        #              title=r'D(S$_3$) Code-Capacity Threshold (HDRG)')
        
        script_dir = os.path.dirname(os.path.abspath(__file__))
        plot_dir = os.path.join(script_dir, 'threshold_plots')
        os.makedirs(plot_dir, exist_ok=True)

        plot_title = r'D(S$_3$) Code-Capacity Threshold (JIT Decoder)'
        plot_name = f'code_capacity_threshold_{formatted_time}.png'
        plot_results(results, L_values, p_values,
                     os.path.join(plot_dir, plot_name),
                     title=plot_title)

        _print_table(results, L_values, p_values)

    # ── Phenomenological ──
    if args.mode in ('phenomenological', 'both'):
        print('\n' + '=' * 60)
        print('Phenomenological Threshold (T=L, p_meas=p)')
        print('=' * 60)

        L_values = [4, 7, 10, 13]
        p_values = [0.001, 0.002, 0.003, 0.004, 0.005, 0.006,
                    0.007, 0.008, 0.010, 0.012, 0.015]
        n_trials = args.trials

        print(f'L: {L_values}')
        print(f'p: {[f"{p:.3f}" for p in p_values]}')
        print(f'Trials: {n_trials}')
        print()

        t0 = time.time()
        results = run_sweep_phenomenological(L_values, p_values,
                                             n_trials=n_trials,
                                             seed=args.seed, verbose=True)
        print(f'\nTotal: {time.time()-t0:.0f}s')

        plot_title = r'D(S$_3$) Phenomenological Threshold (JIT Decoder)'
        plot_name = f'phenomenological_threshold_{formatted_time}.png'
        plot_results(results, L_values, p_values,
                     os.path.join(plot_dir, plot_name),
                     title=plot_title)

        _print_table(results, L_values, p_values)


def _print_table(results, L_values, p_values):
    """Print a summary table of results."""
    print('\nSummary:')
    hdr = f'{"p":>8}' + ''.join(f'  L={L:<3d}    ' for L in sorted(L_values))
    print(hdr)
    print('-' * len(hdr))
    for p in sorted(p_values):
        row = f'{p:8.4f}'
        for L in sorted(L_values):
            pL, se = results[L][p]
            row += f'  {pL:.4f}±{se:.4f}'
        print(row)


if __name__ == '__main__':
    main()