import stim
import numpy as np

# Load circuit
circuit = stim.Circuit.from_file("datasets/sc_d3_r3_p0010_z.stim")

# Extract DEM - this is what BP needs
dem = circuit.detector_error_model(decompose_errors=True)
print(f"Detectors: {dem.num_detectors}")      # 24
print(f"Error mechanisms: {dem.num_errors}")  # 286
print(f"Observables: {dem.num_observables}")  # 1

def build_parity_check_matrix(dem):
    """Convert DEM to parity check matrix H and prior probabilities."""
    errors = []
    for inst in dem.flattened():
        if inst.type == 'error':
            prob = inst.args_copy()[0]
            dets = [t.val for t in inst.targets_copy() if t.is_relative_detector_id()]
            obs = [t.val for t in inst.targets_copy() if t.is_logical_observable_id()]
            errors.append({'prob': prob, 'detectors': dets, 'observables': obs})

    n_detectors = dem.num_detectors
    n_errors = len(errors)

    # Parity check matrix
    H = np.zeros((n_detectors, n_errors), dtype=np.uint8)
    # Prior error probabilities (for BP initialization)
    priors = np.zeros(n_errors)
    # Which errors flip the logical observable
    obs_flip = np.zeros(n_errors, dtype=np.uint8)

    for j, e in enumerate(errors):
        priors[j] = e['prob']
        for d in e['detectors']:
            H[d, j] = 1
        if e['observables']:
            obs_flip[j] = 1

    return H, priors, obs_flip

H, priors, obs_flip = build_parity_check_matrix(dem)
print(f"H shape: {H.shape}")  # (24, 286)

# Compile sampler
sampler = circuit.compile_detector_sampler()

# Sample detection events + observable flip
n_shots = 1000
samples = sampler.sample(n_shots, append_observables=True)

# Split into syndrome and observable
syndromes = samples[:, :-1]           # shape: (n_shots, n_detectors)
actual_obs_flips = samples[:, -1]     # shape: (n_shots,)

print(f"Syndrome shape: {syndromes.shape}")
print(f"Example syndrome: {syndromes[0]}")

def bp_decode(H, syndrome, priors, max_iter=50, damping=0.5):
    """
    Belief Propagation decoder (min-sum variant).

    Args:
        H: Parity check matrix (n_detectors, n_errors)
        syndrome: Detection events (n_detectors,)
        priors: Prior error probabilities (n_errors,)
        max_iter: Maximum BP iterations
        damping: Message damping factor

    Returns:
        estimated_errors: Most likely error pattern (n_errors,)
        soft_output: Log-likelihood ratios (n_errors,)
    """
    n_checks, n_vars = H.shape

    # Initialize LLRs from priors: LLR = log((1-p)/p)
    llr_prior = np.log((1 - priors) / priors)

    # Messages: check-to-variable and variable-to-check
    # ... BP message passing iterations ...

    # Hard decision
    estimated_errors = (soft_output < 0).astype(int)

    return estimated_errors, soft_output

# Decode each syndrome
for i in range(n_shots):
    syndrome = syndromes[i]
    estimated_errors, _ = bp_decode(H, syndrome, priors)

    # Predict observable flip
    predicted_obs_flip = np.dot(estimated_errors, obs_flip) % 2

    # Check if decoding succeeded
    success = (predicted_obs_flip == actual_obs_flips[i])

def evaluate_decoder(decoder_fn, circuit, n_shots=10000):
    """Evaluate decoder logical error rate."""
    dem = circuit.detector_error_model(decompose_errors=True)
    H, priors, obs_flip = build_parity_check_matrix(dem)

    sampler = circuit.compile_detector_sampler()
    samples = sampler.sample(n_shots, append_observables=True)
    syndromes = samples[:, :-1]
    actual_obs = samples[:, -1]

    errors = 0
    for i in range(n_shots):
        est_errors, _ = decoder_fn(H, syndromes[i], priors)
        pred_obs = np.dot(est_errors, obs_flip) % 2
        if pred_obs != actual_obs[i]:
            errors += 1

    return errors / n_shots

# logical_error_rate = evaluate_decoder(bp_decode, circuit)