use crate::code::{LdpcCode, SystematicForm};
use crate::matrix::{DenseMatrixGF2, SparseMatrixGF2};
use crate::{BitVec, Gf2, LdpcError, Result};

// Trait Encoder

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

    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(())
    }
}

// Méthode d'encodage

#[derive(Debug, Clone)]
pub enum EncodingMethod {
    // Via G = [I_k | P]. Complexité O(n^2). Point de départ simple
    Systematic,
    // Richardson-Urbanke. Complexité O(n + g^2), g = gap << n
    RichardsonUrbanke,
}

// Encodeur systématique
// c = m * G = [m | m*P]
// Les bits de parité p = m * P sont calculés par multiplication dense
// Le mot de code est ensuite réordonné selon la permutation inverse

pub struct SystematicEncoder {
    k: usize,
    n: usize,
    g: DenseMatrixGF2,
    perm_inv: Vec<usize>,
}

impl SystematicEncoder {
    pub fn new(code: &mut LdpcCode) -> Result<Self> {
        code.compute_systematic_form()?;
        let sf = code.systematic_form.as_ref().unwrap();
        Ok(Self {
            k: code.k(),
            n: code.n(),
            g: sf.g.clone(),
            perm_inv: sf.col_perm_inv.clone(),
        })
    }
}

impl Encoder for SystematicEncoder {
    fn encode(&self, message: &[Gf2]) -> Result<BitVec> {
        self.check_input(message)?;
        let c_perm = self.g.multiply_vec(message);
        // Réordonner les bits selon la permutation inverse
        let mut c = vec![0u8; self.n];
        for (i, &ci) in c_perm.iter().enumerate() {
            c[self.perm_inv[i]] = ci;
        }
        Ok(c)
    }

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

// Encodeur Richardson-Urbanke
// H est réarrangée par permutations en :
//
//      ┌         ┐
//      │ A  B  T │
//  H = │         │    T = triangulaire inférieure, φ = -ET^{-1}B + D
//      │ C  D  E │
//      └         ┘
//
// Encodage en deux étapes :
//  1. p₂ = φ^{-1} * (-ET^{-1}*As - Cs)       [dense g×g, g = gap]
//  2. p₁ = T^{-1} * (As + Bp₂)              [substitution arrière]
// Complexité totale : O(n + g²)

pub struct RichardsonUrbankeEncoder {
    k: usize,
    n: usize,
    a: SparseMatrixGF2,
    b: SparseMatrixGF2,
    c: SparseMatrixGF2,
    d: SparseMatrixGF2,
    e: SparseMatrixGF2,
    // T^{-1} : résolution par substitution arrière (T triangulaire inférieure)
    t_inv: DenseMatrixGF2,
    // φ^{-1} : petit (g×g), calculé une fois offline
    phi_inv: DenseMatrixGF2,
    col_perm_inv: Vec<usize>,
}

impl RichardsonUrbankeEncoder {
    pub fn new(code: &LdpcCode) -> Result<Self> {
        // TODO: implémenter le pré-traitement de H par permutations de lignes/colonnes
        // pour identifier les blocs A, B, C, D, E, T et calculer phi^{-1}.
        todo!("Pré-traitement Richardson-Urbanke")
    }

    // Substitution arrière sur T triangulaire inférieure : résout T*x = b en GF(2)
    fn backward_substitution(t_inv: &DenseMatrixGF2, b: &[Gf2]) -> Vec<Gf2> {
        let n = b.len();
        let mut x = vec![0u8; n];
        for i in 0..n {
            let mut s = b[i];
            for j in 0..i {
                s ^= t_inv.get(i, j) & x[j];
            }
            x[i] = s;
        }
        x
    }
}

impl Encoder for RichardsonUrbankeEncoder {
    fn encode(&self, message: &[Gf2]) -> Result<BitVec> {
        self.check_input(message)?;
        // Étape 1 : As
        let a_s = self.a.multiply_vec(message);
        // Étape 2 : p₂ = φ^{-1} * (E*T^{-1}*As + Cs)
        let t_inv_as = Self::backward_substitution(&self.t_inv, &a_s);
        let e_t_inv_as = self.e.multiply_vec(&t_inv_as);
        let c_s = self.c.multiply_vec(message);
        let rhs: Vec<Gf2> = e_t_inv_as
            .iter()
            .zip(c_s.iter())
            .map(|(&a, &b)| a ^ b)
            .collect();
        let p2 = self.phi_inv.multiply_vec(&rhs);
        // Étape 3 : p₁ = T^{-1} * (As + Bp₂)
        let b_p2 = self.b.multiply_vec(&p2);
        let sum: Vec<Gf2> = a_s.iter().zip(b_p2.iter()).map(|(&a, &b)| a ^ b).collect();
        let p1 = Self::backward_substitution(&self.t_inv, &sum);
        // Assemblage et dépermutation
        let mut c_perm: Vec<Gf2> = message
            .iter()
            .chain(p1.iter())
            .chain(p2.iter())
            .cloned()
            .collect();
        let mut c = vec![0u8; self.n];
        for (i, &ci) in c_perm.iter().enumerate() {
            c[self.col_perm_inv[i]] = ci;
        }
        Ok(c)
    }

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

// Factory

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

// #[cfg(test)]
// mod tests {
//     use super::*;
//     use crate::code::{CodeTopology, GenerationMethod, LdpcCode, LdpcParams};
//
//     fn make_code() -> LdpcCode {
//         LdpcCode::new(LdpcParams {
//             n: 12,
//             k: 6,
//             topology: CodeTopology::Regular { wc: 2, wr: 4 },
//             generation: GenerationMethod::Gallager,
//             seed: Some(42),
//         })
//         .unwrap()
//     }
//
//     #[test]
//     fn test_encoder_wrong_input_length_errors() {
//         let mut code = make_code();
//         let enc = SystematicEncoder::new(&mut code).unwrap();
//         let result = enc.encode(&vec![0u8; 5]); // devrait être 6
//         assert!(result.is_err());
//     }
// }