import numpy as np
from ldpc import bp_decoder
from ldpc import bposd_decoder
from bposd.css import css_code
from decoders import first_min_bp_decoder, bp_uf_decoder, uf_decoder
from scipy.sparse import csr_array,csr_matrix

def syndrome_history_map(H, width):
    rep_code = np.zeros((width, width), dtype=np.uint8)
    for i in range(width):
        rep_code[i][i] = 1
        if i+1 < width:
            rep_code[i+1][i] = 1

    return csr_matrix(np.concatenate((np.kron(np.eye(width), H), np.kron(rep_code, np.eye(H.shape[0]))), axis=1))

def code_capacity(
    code,
    phys_err_rate,
    target_num_logical_errors
):
    p = phys_err_rate
    hz = code.hz
    n = hz.shape[1]

    #decoder = bp_uf_decoder(hz, p)
    #decoder = uf_decoder(hz, p)
    decoder = bposd_decoder(
        hz,
        error_rate = p,
        bp_method = "ms",
        max_iter = int(n/10),
        ms_scaling_factor = 0,
        osd_method = "osd0"
    )

    num_samples = 0
    num_logical_errors = 0
    while num_logical_errors < target_num_logical_errors:
        error = np.random.binomial(1, p, n)
        syndrome = (hz@error)%2
        error = (error + decoder.decode(syndrome)) % 2
        if ((code.lz@error)%2).any():
            num_logical_errors += 1
        num_samples += 1

    return num_samples

def window_decoding(
    code, 
    width, 
    offset, 
    phys_err_rate, 
    num_logical_errors, 
    decoding_method = "bp_osd"
):
    hz = code.hz
    p = phys_err_rate
    n = hz.shape[1]
    m = hz.shape[0]

    ideal_decoder = bposd_decoder(
        code.hz,
        error_rate = p,
        bp_method = "ms",
        ms_scaling_factor = 0,
        osd_method = "osd_cs",
        osd_order = 20
    )

    if decoding_method == "bp":
        decoder = first_min_bp_decoder(
            syndrome_history_map(hz, width),
            p
        )
    else:
        decoder = bposd_decoder(
            syndrome_history_map(hz, width),
            error_rate = p,
            bp_method = "ms",
            ms_scaling_factor = 0,
            osd_method = "osd_cs",
            osd_order = 20
        )
 

    hz = csr_matrix(hz)
    lifetime = 0
    for t in range(num_logical_errors):
        data_err = np.random.binomial(1, p, width*n)
        meas_err = np.random.binomial(1, p, width*m)

        while True:
            syndrome_full = meas_err.copy()
            syndrome_full[m:] += meas_err[:-m]
            for i in range(width):
                syndrome_full[i*m:(i+1)*m] += hz@data_err[i*n:(i+1)*n]
            syndrome_full %= 2

            correction = decoder.decode(syndrome_full)[:offset*n]
            residual_err = np.zeros(n)

            stopped = False
            for i in range(offset):
                lifetime += 1
                residual_err += data_err[i*n:(i+1)*n]
                residual_err += correction[i*n:(i+1)*n]
                residual_err %= 2
                syndrome_ideal = (hz@residual_err)%2
                if (code.lz@(residual_err+ideal_decoder.decode(syndrome_ideal))%2).any():
                    stopped = True
                    break
            if stopped:
                break

                    
            data_err[:-offset*n] = data_err[offset*n:]
            data_err[-offset*n:] = np.random.binomial(1, p, offset*n)
            data_err[:n] = (data_err[:n]+residual_err) % 2
            meas_err[:-offset*m] = meas_err[offset*m:]
            meas_err[-offset*m:] = np.random.binomial(1, p, offset*m)
        
    return lifetime