use crate::code::LdpcCode;
use crate::{BitVec, Gf2, LdpcError, Result};

pub trait Encoder: Send + Sync {
    fn encode(&self, message: &[Gf2]) -> Result<BitVec>;
    fn message_len(&self) -> usize;
    fn codeword_len(&self) -> usize;
    fn extract_message(&self, codeword: &[Gf2]) -> Vec<Gf2>;

    fn check_input(&self, msg: &[Gf2]) -> Result<()> {
        if msg.len() != self.message_len() {
            return Err(LdpcError::DimensionMismatch {
                expected: self.message_len(),
                got: msg.len(),
            });
        }
        Ok(())
    }
}

#[derive(Debug, Clone)]
pub enum EncodingMethod {
    Systematic,
}

pub struct SystematicEncoder {
    k: usize,
    n: usize,
    // g_t: DenseMatrixGF2,
    packed_g_cols: Vec<Vec<u64>>,
    perm_inv: Vec<usize>,
    col_perm: Vec<usize>,
}

impl SystematicEncoder {
    pub fn new(code: &mut LdpcCode) -> Result<Self> {
        code.compute_systematic_form()?;
        let sf = code.systematic_form.as_ref().unwrap();

        let k = code.k();
        let n = code.n();
        let g_t = &sf.g;

        let num_blocks = (n + 63) / 64;

        // Bitpacking
        let mut packed_g_cols = vec![vec![0u64; num_blocks]; k];

        for j in 0..k {
            for i in 0..n {
                if g_t.get(i, j) == 1 {
                    let block_idx = i / 64;
                    let bit_idx = i % 64;
                    packed_g_cols[j][block_idx] |= 1 << bit_idx;
                }
            }
        }

        Ok(Self {
            k,
            n,
            packed_g_cols,
            perm_inv: sf.col_perm_inv.clone(),
            col_perm: sf.col_perm.clone(),
        })
    }
}

impl Encoder for SystematicEncoder {
    // fn encode(&self, message: &[Gf2]) -> Result<BitVec> {
    //     self.check_input(message)?;
    //
    //     let c_perm = self.g_t.multiply_vec(message);
    //
    //     // Retablir l'ordre initial des bits selon la permutation de H
    //     let mut c = vec![0u8; self.n];
    //     // for (i, &ci) in c_perm.iter().enumerate() {
    //     //     c[self.perm_inv[i]] = ci;
    //     // }
    //
    //     for i in 0..self.n {
    //         c[i] = c_perm[self.perm_inv[i]];
    //     }
    //
    //     Ok(c)
    // }
    fn encode(&self, message: &[Gf2]) -> Result<BitVec> {
        self.check_input(message)?;

        let num_blocks = (self.n + 63) / 64;
        let mut accum = vec![0u64; num_blocks];

        for (j, &bit) in message.iter().enumerate() {
            if bit == 1 {
                for b in 0..num_blocks {
                    accum[b] ^= self.packed_g_cols[j][b];
                }
            }
        }

        let mut c = vec![0u8; self.n];
        for i in 0..self.n {
            let block_idx = i / 64;
            let bit_idx = i % 64;
            let val = ((accum[block_idx] >> bit_idx) & 1) as u8;

            c[self.col_perm[i]] = val;
        }

        Ok(c)
    }

    fn extract_message(&self, codeword: &[Gf2]) -> Vec<Gf2> {
        let mut msg = vec![0u8; self.k];
        for j in 0..self.k {
            msg[j] = codeword[self.col_perm[j]];
        }
        msg
    }

    fn message_len(&self) -> usize {
        self.k
    }

    fn codeword_len(&self) -> usize {
        self.n
    }
}

pub fn build_encoder(code: &mut LdpcCode, method: EncodingMethod) -> Result<Box<dyn Encoder>> {
    match method {
        EncodingMethod::Systematic => Ok(Box::new(SystematicEncoder::new(code)?)),
    }
}
