import matplotlib.pyplot as plt
import numpy as np
from sympy import Matrix, Poly, div, symbols
from functools import reduce


def cyclotomic_cosets(order):
    cosets = []
    for i in range(order):
        coset = set()
        for j in range(order):
            element = (2**j * i) % order
            coset.add(element)
        if coset not in cosets:
            cosets.append(coset)
    return cosets[1:]


def minimal_polynomials(cosets):
    x = symbols("x")
    min_polys = []

    for coset in cosets:
        polys = [Poly(x - a, domain="GF(2)") for a in coset]
        min_poly = reduce(lambda poly1, poly2: poly1 * poly2, polys)
        min_poly += 1
        min_polys.append(min_poly)

    return min_polys


def generate_poly(cosets, min_polys, error_correction):
    order = max(max(coset) for coset in cosets) + 1
    selected_polys = [
        min_polys[i] for i in range(1, 2 * error_correction + 1) if i < order
    ]
    generator_poly = reduce(lambda a, b: a * b, selected_polys)
    generator_poly = Poly(generator_poly, domain="GF(2)")
    return generator_poly


def bch_params(generator_poly, order):
    n_k = generator_poly.degree()
    k = order - n_k
    t = sum(generator_poly.as_expr().as_coefficients_dict().values()) // 2
    return k, t


def encode_message(vector, order, info_bits, generator_poly):
    x = symbols("x")
    info_poly = Poly(vector, x, domain="GF(2)")
    info_poly = info_poly * x ** (order - info_bits)
    quotient, remainder = div(info_poly, generator_poly)
    encoded_message = info_poly + remainder
    encoded_message = np.array(encoded_message.all_coeffs()[::-1]) % 2
    return encoded_message


def decode_message(encoded_message, order, info_bits, generator_poly):
    x = symbols("x")
    encoded_message = Poly(encoded_message, x, domain="GF(2)")
    quotient, remainder = div(encoded_message, generator_poly)
    decoded_message = quotient * x ** (order - info_bits)
    decoded_message = np.array(decoded_message.all_coeffs()[::-1]) % 2
    return decoded_message


def noisy_channel(bits, snr):
    noise = np.random.binomial(n=1, p=1 - 1 / (1 + snr), size=len(bits))
    return bits ^ noise


def simulate_bch(order, info_bits, snr, num_messages):
    errors = 0
    total_bits = 0
    for _ in range(num_messages):
        message = np.random.randint(2, size=info_bits)
        encoded_message = encode_message(message, order, info_bits, generator_poly)
        received_message = noisy_channel(encoded_message, snr)
        decoded_message = decode_message(
            received_message, order, info_bits, generator_poly
        )
        errors += np.sum(message != decoded_message)
        total_bits += len(message)
    return errors / total_bits


order = 63
info_bits = 36
error_locations = [18, 5]

cosets = cyclotomic_cosets(order)

min_polys = minimal_polynomials(cosets)

generator_poly = generate_poly(cosets, min_polys, 3)

info_bits, error_correction = bch_params(generator_poly, order)

parity_check_matrix = Matrix.zeros(order - info_bits, order)
n_k = order - info_bits
coeffs = generator_poly.all_coeffs()[::-1]

for i in range(n_k):
    for j in range(order):
        parity_check_matrix[i, j] = coeffs[(i + j) % n_k]

message_vector = [0] * 33 + [0, 0, 1, 0, 1, 1, 1]
encoded_message = encode_message(message_vector, order, info_bits, generator_poly)
encoded_message[error_locations[0]] ^= 1
encoded_message[error_locations[1]] ^= 1
decoded_message = decode_message(encoded_message, order, info_bits, generator_poly)
snrs = np.linspace(0, 10, 11)
num_messages = 1000
error_probs_31_21 = [simulate_bch(31, 21, snr, num_messages) for snr in snrs]
error_probs_15_7 = [simulate_bch(15, 7, snr, num_messages) for snr in snrs]
plt.figure(figsize=(8, 6))
plt.semilogy(snrs, error_probs_31_21, label="BCH (31, 21)", color='red', marker='o')
plt.semilogy(snrs, error_probs_15_7, label="BCH (15, 7)", color='blue', marker='o')
plt.xlabel("SNR (dB)")
plt.ylabel("Вероятность ошибки на бит")
plt.title("Вероятность ошибки на бит относительно SNR")
plt.legend()
plt.grid(True)
plt.gca().set_facecolor('dimgray')
plt.show()