use crate::graph::TannerGraph;
use crate::matrix::{DenseMatrixGF2, SparseMatrixGF2};
use crate::{Gf2, LdpcError, Result};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LdpcParams {
    pub n: usize,
    pub k: usize,
    pub topology: CodeTopology,
    pub generation: GenerationMethod,
    pub seed: Option<u64>,
}

impl LdpcParams {
    pub fn rate(&self) -> f64 {
        self.k as f64 / self.n as f64
    }
    pub fn m(&self) -> usize {
        self.n - self.k
    }

    pub fn validate(&self) -> Result<()> {
        if self.k >= self.n {
            return Err(LdpcError::InvalidParameters("k doit être < n".into()));
        }
        if self.n < 4 {
            return Err(LdpcError::InvalidParameters("n trop petit".into()));
        }
        self.topology.validate(self)
    }
}

// Topologie
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CodeTopology {
    // Chaque varnode a degré wc, chaque checknode a degré wr
    // Condition nécessaire n * wc == m * wr
    Regular { wc: usize, wr: usize },
    // Potentielle implémentation Irregular (plus tard !)
}

impl CodeTopology {
    fn validate(&self, params: &LdpcParams) -> Result<()> {
        match self {
            CodeTopology::Regular { wc, wr } => {
                if params.n * wc != params.m() * wr {
                    return Err(LdpcError::InvalidParameters(format!(
                        "n*wc ({}) != m*wr ({}) pour LDPC régulier",
                        params.n * wc,
                        params.m() * wr
                    )));
                }
            }
        }
        Ok(())
    }
}

// Méthode de génération
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GenerationMethod {
    // H = [H1 | H2 | ... | Hwc]^T, H1 régulière, Hi = permutation de H1
    Gallager,
    // Ajout de colonnes de poids fixe, rejet si cycle4 créé
    MacKayNeal { max_attempts: usize },
}

// Forme systématique
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystematicForm {
    // G = [I_k | P], dense
    pub g: DenseMatrixGF2,
    // Permutation de colonnes appliquée à H => [A | I_m]
    pub col_perm: Vec<usize>,
    // Permutation inverse => reformer le mot de code
    pub col_perm_inv: Vec<usize>,
}

// Structure principale
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LdpcCode {
    pub params: LdpcParams,
    pub h: SparseMatrixGF2,
    #[serde(skip, default = "default_graph")]
    pub graph: TannerGraph,
    pub systematic_form: Option<SystematicForm>,
}

impl LdpcCode {
    pub fn new(params: LdpcParams) -> Result<Self> {
        params.validate()?;
        let mut rng = build_rng(params.seed);
        let h = match &params.generation {
            GenerationMethod::Gallager => generate_gallager(&params, &mut rng)?,
            GenerationMethod::MacKayNeal { max_attempts } => {
                generate_mackay_neal(&params, *max_attempts, &mut rng)?
            }
        };
        let graph = TannerGraph::from_matrix(&h);
        Ok(Self {
            params,
            h,
            graph,
            systematic_form: None,
        })
    }

    pub fn from_matrix(h: SparseMatrixGF2, k: usize) -> Result<Self> {
        let n = h.cols;
        let params = LdpcParams {
            n,
            k,
            topology: CodeTopology::Regular { wc: 0, wr: 0 }, // inconnu
            generation: GenerationMethod::Gallager,
            seed: None,
        };
        let graph = TannerGraph::from_matrix(&h);
        Ok(Self {
            params,
            h,
            graph,
            systematic_form: None,
        })
    }

    // Calcule G par Gauss-Jordan sur H (to cache)
    pub fn compute_systematic_form(&mut self) -> Result<()> {
        if self.systematic_form.is_some() {
            return Ok(());
        }

        let mut dense = DenseMatrixGF2::zeros(self.m(), self.n());
        let h_dense = self.h.to_dense();
        for r in 0..self.m() {
            for c in 0..self.n() {
                dense.set(r, c, h_dense[r][c]);
            }
        }

        let (col_perm, rank) = dense.systematize(self.k());

        // Possibilité d'avoir rg < m donc juste rejeter et recommencer
        if rank < self.m() {
            return Err(LdpcError::SingularMatrix);
        }

        // Création de G^T (taille n,k) pour encodage c = G^T * m
        let mut g_t = DenseMatrixGF2::zeros(self.n(), self.k());

        // I_k (haut)
        for i in 0..self.k() {
            g_t.set(i, i, 1);
        }

        // (bas) Matrice A (extraite des k premières colonnes de la matrice dense pivotée)
        for r in 0..self.m() {
            for c in 0..self.k() {
                g_t.set(self.k() + r, c, dense.get(r, c));
            }
        }

        // Création de la permutation inverse pour réordonner le mot de code
        let mut col_perm_inv = vec![0usize; self.n()];
        for (i, &p) in col_perm.iter().enumerate() {
            col_perm_inv[p] = i;
        }

        self.systematic_form = Some(SystematicForm {
            g: g_t,
            col_perm,
            col_perm_inv,
        });
        Ok(())
    }

    pub fn rate(&self) -> f64 {
        self.params.rate()
    }
    pub fn n(&self) -> usize {
        self.params.n
    }
    pub fn k(&self) -> usize {
        self.params.k
    }
    pub fn m(&self) -> usize {
        self.params.m()
    }
    pub fn girth(&self) -> usize {
        self.graph.girth()
    }

    pub fn is_codeword(&self, c: &[Gf2]) -> bool {
        self.h.multiply_vec(c).iter().all(|&s| s == 0)
    }
}

fn build_rng(seed: Option<u64>) -> impl rand::Rng {
    use rand::SeedableRng;
    rand::rngs::StdRng::seed_from_u64(seed.unwrap_or_else(rand::random))
}

// Gallager
// H divisée en wc sous-matrices de taille (m / wc) * n
// H1 = matrice régulière (ligne contient wr 1)
// H2..Hwc = permutations aléatoires de colonnes de H1

fn generate_gallager(params: &LdpcParams, rng: &mut impl rand::Rng) -> Result<SparseMatrixGF2> {
    let CodeTopology::Regular { wc, wr } = params.topology else {
        return Err(LdpcError::InvalidParameters(
            "Gallager nécessite un code régulier".into(),
        ));
    };
    let n = params.n;
    let m = params.m();
    if m % wc != 0 {
        return Err(LdpcError::InvalidParameters(
            "m doit être divisible par wc".into(),
        ));
    }
    let rows_per_block = m / wc;
    let mut ones: Vec<(usize, usize)> = Vec::new();

    for r in 0..rows_per_block {
        for j in 0..wr {
            ones.push((r, r * wr + j));
        }
    }
    use rand::seq::SliceRandom;
    for block in 1..wc {
        let mut perm: Vec<usize> = (0..n).collect();
        perm.shuffle(rng);
        for r in 0..rows_per_block {
            for j in 0..wr {
                ones.push((block * rows_per_block + r, perm[r * wr + j]));
            }
        }
    }

    for b in 1..wc {
        let r_source = b * rows_per_block;

        let row_0_cols: Vec<usize> = ones
            .iter()
            .filter(|&&(r, _)| r == 0)
            .map(|&(_, c)| c)
            .collect();

        if let Some(idx) = ones
            .iter()
            .position(|&(r, c)| r == r_source && !row_0_cols.contains(&c))
        {
            ones[idx].0 = 0;
        }
    }

    Ok(SparseMatrixGF2::from_positions(m, n, ones))
}

// MacKay-Neal
// Ajoute les colonnes une à une avec poids wc
// Rejette colonne créant cycle4 (2 colonnes n'ont qu'un 1 en commun)
fn generate_mackay_neal(
    params: &LdpcParams,
    max_attempts: usize,
    rng: &mut impl rand::Rng,
) -> Result<SparseMatrixGF2> {
    let CodeTopology::Regular { wc, wr } = params.topology else {
        return Err(LdpcError::InvalidParameters(
            "MacKayNeal nécessite régulier".into(),
        ));
    };
    let n = params.n;
    let m = params.m();
    let mut ones: Vec<(usize, usize)> = Vec::new();

    // Suivi du poids de chaque ligne
    let mut row_weights = vec![0usize; m];

    use rand::seq::SliceRandom;
    for col in 0..n {
        let mut placed = false;
        for _attempt in 0..max_attempts {
            let mut available_rows: Vec<usize> = (0..m).filter(|&r| row_weights[r] < wr).collect();

            if available_rows.len() < wc {
                break; // Plus de lignes dispo
            }

            available_rows.shuffle(rng);
            let candidate = available_rows[..wc].to_vec();

            // Vérifier cycle4
            let mut ok = true;
            let mut c2 = 0;
            while c2 < col && ok {
                let existing: Vec<usize> = ones
                    .iter()
                    .filter(|&&(_, c)| c == c2)
                    .map(|&(r, _)| r)
                    .collect();
                let shared = candidate.iter().filter(|r| existing.contains(r)).count();
                if shared >= 2 {
                    ok = false;
                }
                c2 += 1;
            }

            if ok {
                for &r in &candidate {
                    ones.push((r, col));
                    row_weights[r] += 1;
                }
                placed = true;
                break;
            }
        }
        if !placed {
            return Err(LdpcError::GenerationFailed {
                attempts: max_attempts,
            });
        }
    }
    Ok(SparseMatrixGF2::from_positions(m, n, ones))
}

fn default_graph() -> TannerGraph {
    TannerGraph::from_matrix(&SparseMatrixGF2::zeros(1, 1))
}
