use crate::code::{self, LdpcCode};
use crate::matrix::MatrixGF2;

pub struct DenseEncoder {
    pub g_matrix: MatrixGF2,
    pub h_reduced: MatrixGF2, // A enlever c'est juste pour l'affichage...
    pub k: usize,
    pub n: usize,
    pub col_swaps: Vec<(usize, usize)>,
    pub rank: usize,
}

impl DenseEncoder {
    pub fn new(code: &LdpcCode) -> Self {
        let mut h_reduced = code.h_matrix.clone();
        let (rank, col_swaps) = h_reduced.gauss_jordan_swap_cols();

        let n = h_reduced.cols;
        let k = n - rank;

        let mut g_data = vec![0; k * n];

        for r in 0..k {
            for c in 0..rank {
                g_data[r * n + c] = h_reduced.get(c, rank + r);
            }
        }
        for i in 0..k {
            g_data[i * n + (rank + i)] = 1;
        }
        Self {
            g_matrix: MatrixGF2::new(k, n, g_data),
            h_reduced,
            k,
            n,
            col_swaps,
            rank,
        }
    }

    pub fn encode(&self, message: &[u8]) -> Vec<u8> {
        assert_eq!(message.len(), self.k, "Taille du message non valide");

        let mut codeword: Vec<u8> = (0..self.n)
            .map(|c| {
                message.iter().enumerate().fold(0, |acc, (r, &msg_bit)| {
                    acc ^ (msg_bit & self.g_matrix.get(r, c))
                })
            })
            .collect();

        // Reverse les changement de colonnes
        // for &(c1, c2) in self.col_swaps.iter().rev() {
        //     codeword.swap(c1, c2);
        // }
        self.g_matrix
            .apply_inverse_permutation(&mut codeword, &self.col_swaps);

        codeword
    }

    pub fn extract_message(&self, decoded_codeword: &[u8]) -> Vec<u8> {
        let mut working_copy = decoded_codeword.to_vec();

        self.g_matrix
            .apply_permutation(&mut working_copy, &self.col_swaps);

        let mut message = Vec::with_capacity(self.k);
        for i in 0..self.k {
            message.push(working_copy[self.rank + i]);
        }
        message
    }
}