import numpy as np
import pytest
import stim

import simmer


def test_decode_using_pymatching():
    circuit = stim.Circuit.generated('repetition_code:memory',
                                     rounds=3,
                                     distance=3,
                                     after_clifford_depolarization=0.05)
    result = simmer.sample_decode(
        circuit=circuit,
        num_shots=1000,
        decoder='pymatching',
    )
    assert result.num_discards == 0
    assert 1 <= result.num_errors <= 100
    assert result.num_shots == 1000


def test_pymatching_works_on_surface_code():
    stats = simmer.sample_decode(
        num_shots=1000,
        circuit=stim.Circuit.generated(
            "surface_code:rotated_memory_x",
            distance=3,
            rounds=15,
            after_clifford_depolarization=0.001,
        ),
        decoder="pymatching",
    )
    assert 0 <= stats.num_errors <= 50


def test_decode_using_internal_decoder():
    pytest.importorskip('gqec')

    circuit = stim.Circuit.generated('repetition_code:memory',
                                     rounds=3,
                                     distance=3,
                                     after_clifford_depolarization=0.05)
    result = simmer.sample_decode(
        circuit=circuit,
        num_shots=1000,
        decoder='internal',
    )
    assert result.num_discards == 0
    assert 1 <= result.num_errors <= 100
    assert result.num_shots == 1000


def test_decode_using_internal_decoder_correlated():
    pytest.importorskip('gqec')

    circuit = stim.Circuit.generated('repetition_code:memory',
                                     rounds=3,
                                     distance=3,
                                     after_clifford_depolarization=0.05)
    result = simmer.sample_decode(
        circuit=circuit,
        num_shots=1000,
        decoder='internal_correlated',
    )
    assert result.num_discards == 0
    assert 1 <= result.num_errors <= 100
    assert result.num_shots == 1000


def test_empty():
    result = simmer.sample_decode(
        circuit=stim.Circuit(),
        num_shots=1000,
        decoder='pymatching',
    )
    assert result.num_discards == 0
    assert result.num_shots == 1000
    assert result.num_errors == 0


def test_no_observables():
    c = stim.Circuit("""
        X_ERROR(0.1) 0
        M 0
        DETECTOR rec[-1]
    """)
    result = simmer.sample_decode(
        circuit=c,
        num_shots=1000,
        decoder='pymatching',
    )
    assert result.num_discards == 0
    assert result.num_shots == 1000
    assert result.num_errors == 0


@pytest.mark.parametrize("offset", range(8))
def test_observable_offsets_mod8(offset: int):
    c = stim.Circuit("""
        X_ERROR(0.1) 0
        MR 0
        DETECTOR rec[-1]
    """) * (8 + offset) + stim.Circuit("""
        X_ERROR(0.1) 0
        MR 0
        OBSERVABLE_INCLUDE(0) rec[-1]
    """)
    result = simmer.sample_decode(
        circuit=c,
        num_shots=1000,
        decoder='pymatching',
    )
    assert result.num_discards == 0
    assert result.num_shots == 1000
    assert 50 <= result.num_errors <= 150


def test_no_detectors():
    c = stim.Circuit("""
        X_ERROR(0.1) 0
        M 0
        OBSERVABLE_INCLUDE(0) rec[-1]
    """)
    result = simmer.sample_decode(
        circuit=c,
        num_shots=1000,
        decoder='pymatching',
    )
    assert result.num_discards == 0
    assert 50 <= result.num_errors <= 150


def test_no_detectors_with_post_mask():
    c = stim.Circuit("""
        X_ERROR(0.1) 0
        M 0
        OBSERVABLE_INCLUDE(0) rec[-1]
    """)
    result = simmer.sample_decode(
        circuit=c,
        post_mask=np.array([0], dtype=np.uint8),
        num_shots=1000,
        decoder='pymatching',
    )
    assert result.num_discards == 0
    assert 50 <= result.num_errors <= 150


def test_post_selection():
    c = stim.Circuit("""
        X_ERROR(0.6) 0
        M 0
        DETECTOR(2) rec[-1]
        OBSERVABLE_INCLUDE(0) rec[-1]

        X_ERROR(0.5) 1
        M 1
        DETECTOR(1) rec[-1]
        OBSERVABLE_INCLUDE(0) rec[-1]
        
        X_ERROR(0.1) 2
        M 2
        OBSERVABLE_INCLUDE(0) rec[-1]
    """)
    result = simmer.sample_decode(
        circuit=c,
        post_mask=simmer.post_selection_mask_from_last_detector_coords(circuit=c, last_coord_minimum=2),
        num_shots=2000,
        decoder='pymatching',
    )
    assert 1050 <= result.num_discards <= 1350
    assert 40 <= result.num_errors <= 160
