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

#[derive(Debug, Clone)]
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)]
pub enum CodeTopology {
    // Chaque var-node a degré wc, chaque check-node a degré wr
    // Condition nécessaire : n * wc == m * wr.
    Regular { wc: usize, wr: usize },
    // Degrés définis par des polynômes de distribution (density evolution)
    // lambda[i] = fraction d'arêtes sur des var-nodes de degré i+2
    // rho[i]   = fraction d'arêtes sur des check-nodes de degré i+2
    Irregular { lambda: Vec<f64>, rho: Vec<f64> },
}

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 ({}) doit égaler m*wr ({}) pour un code régulier",
                        params.n * wc,
                        params.m() * wr
                    )));
                }
            }
            CodeTopology::Irregular { lambda, rho } => {
                let sl: f64 = lambda.iter().sum();
                let sr: f64 = rho.iter().sum();
                if (sl - 1.0).abs() > 1e-6 || (sr - 1.0).abs() > 1e-6 {
                    return Err(LdpcError::InvalidParameters(
                        "Les polynômes lambda et rho doivent sommer à 1".into(),
                    ));
                }
            }
        }
        Ok(())
    }
}

// Méthode de génération

#[derive(Debug, Clone)]
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 cycle-4 créé
    MacKayNeal { max_attempts: usize },
    // Progressive Edge Growth : maximise le girth local arête par arête
    Peg,
}

// Forme systématique

#[derive(Debug, Clone)]
pub struct SystematicForm {
    // G = [I_k | P], dense car P est généralement pleine
    pub g: DenseMatrixGF2,
    // Permutation de colonnes appliquée à H pour obtenir [A | I_m]
    pub col_perm: Vec<usize>,
    // Permutation inverse pour réordonner le mot de code
    pub col_perm_inv: Vec<usize>,
}

// Structure principale

#[derive(Debug, Clone)]
pub struct LdpcCode {
    pub params: LdpcParams,
    pub h: SparseMatrixGF2,
    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)?
            }
            GenerationMethod::Peg => generate_peg(&params)?,
        };
        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)
    // Appelé automatiquement par SystematicEncoder si nécessaire
    pub fn compute_systematic_form(&mut self) -> Result<()> {
        if self.systematic_form.is_some() {
            return Ok(());
        }

        let dense = self.h.to_dense();
        // On travaille sur [H | I_m] pour tracker les opérations de lignes
        let m = self.params.m();
        let n = self.params.n;

        // Construire la matrice augmentée de taille m (n + m)
        let mut aug = DenseMatrixGF2::zeros(m, n + m);
        for r in 0..m {
            for c in 0..n {
                aug.set(r, c, dense[r][c]);
            }
            aug.set(r, n + r, 1); // partie identité
        }

        // Gauss-Jordan sur les n premières colonnes
        let (col_perm, rank) = aug.gauss_jordan_gf2();
        if rank < m {
            return Err(LdpcError::SingularMatrix);
        }

        // Extraire P (colonnes n..n+m de aug) -> partie redondante de G
        // G = [I_k | P^T] après transposition correcte
        // TODO: extraire proprement P et construire G = [I_k | P]
        let g = DenseMatrixGF2::identity(self.params.k); // placeholder

        let col_perm_inv = {
            let mut inv = vec![0usize; col_perm.len()];
            for (i, &p) in col_perm.iter().enumerate() {
                inv[p] = i;
            }
            inv
        };

        self.systematic_form = Some(SystematicForm {
            g,
            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)
    }
}

// RNG

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 est une matrice régulière (chaque ligne contient exactement wr uns)
// H2..Hwc sont des 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();

    // H1 : blocs réguliers de wr uns par ligne
    for r in 0..rows_per_block {
        for j in 0..wr {
            ones.push((r, r * wr + j));
        }
    }
    // H2..Hwc : permutations aléatoires de colonnes
    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]));
            }
        }
    }
    Ok(SparseMatrixGF2::from_positions(m, n, ones))
}

// MacKay-Neal
// Ajoute les colonnes une à une avec poids wc
// Rejette toute colonne créant un cycle-4 (deux colonnes ne peuvent
// partager qu'un seul 1 commun). Relance si max_attempts est dépassé

fn generate_mackay_neal(
    params: &LdpcParams,
    max_attempts: usize,
    rng: &mut impl rand::Rng,
) -> Result<SparseMatrixGF2> {
    let CodeTopology::Regular { wc, .. } = params.topology else {
        return Err(LdpcError::InvalidParameters(
            "MacKayNeal nécessite un code régulier".into(),
        ));
    };
    let n = params.n;
    let m = params.m();
    let mut ones: Vec<(usize, usize)> = Vec::new();

    use rand::seq::SliceRandom;
    for col in 0..n {
        let mut placed = false;
        for _attempt in 0..max_attempts {
            // Tirer wc lignes sans remise
            let mut rows: Vec<usize> = (0..m).collect();
            rows.shuffle(rng);
            let candidate: Vec<usize> = rows[..wc].to_vec();

            // Vérifier cycle-4 : cette colonne ne partage pas ≥2 lignes avec une colonne existante
            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));
                }
                placed = true;
                break;
            }
        }
        if !placed {
            return Err(LdpcError::GenerationFailed {
                attempts: max_attempts,
            });
        }
    }
    Ok(SparseMatrixGF2::from_positions(m, n, ones))
}

// PEG
// Progressive Edge Growth
// Pour chaque arête à ajouter, choisit le check-node qui maximise
// le girth local du graphe courant (BFS depuis le var-node courant)

fn generate_peg(params: &LdpcParams) -> Result<SparseMatrixGF2> {
    let CodeTopology::Regular { wc, .. } = params.topology else {
        return Err(LdpcError::InvalidParameters(
            "PEG nécessite un code régulier".into(),
        ));
    };
    let n = params.n;
    let m = params.m();
    let mut var_to_chk: Vec<Vec<usize>> = vec![vec![]; n];
    let mut chk_to_var: Vec<Vec<usize>> = vec![vec![]; m];
    let mut ones: Vec<(usize, usize)> = Vec::new();

    for v in 0..n {
        for _edge in 0..wc {
            // BFS depuis v dans le graphe courant pour trouver le check-node
            // non-voisin de v le plus éloigné (maximise le girth local)
            let best_chk = peg_find_best_check(v, &var_to_chk, &chk_to_var, m);
            var_to_chk[v].push(best_chk);
            chk_to_var[best_chk].push(v);
            ones.push((best_chk, v));
        }
    }
    Ok(SparseMatrixGF2::from_positions(m, n, ones))
}

// BFS depuis le var-node v, retourne le check-node non-voisin de v
// qui est le plus éloigné (ou le moins chargé en cas d'égalité)
fn peg_find_best_check(
    v: usize,
    var_to_chk: &[Vec<usize>],
    chk_to_var: &[Vec<usize>],
    n_chk: usize,
) -> usize {
    use std::collections::VecDeque;
    let current_neighbors: &[usize] = &var_to_chk[v];
    let mut dist_chk = vec![usize::MAX; n_chk];
    let mut dist_var = vec![usize::MAX; var_to_chk.len()];
    dist_var[v] = 0;
    let mut queue: VecDeque<(bool, usize, usize)> = VecDeque::new();
    queue.push_back((true, v, 0));
    let mut max_dist = 0;
    let mut reachable_non_neighbors: Vec<(usize, usize)> = Vec::new(); // (dist, chk)

    while let Some((is_var, node, dist)) = queue.pop_front() {
        if is_var {
            for &c in &var_to_chk[node] {
                if dist_chk[c] == usize::MAX {
                    dist_chk[c] = dist + 1;
                    max_dist = max_dist.max(dist + 1);
                    queue.push_back((false, c, dist + 1));
                }
            }
        } else {
            for &u in &chk_to_var[node] {
                if dist_var[u] == usize::MAX {
                    dist_var[u] = dist + 1;
                    queue.push_back((true, u, dist + 1));
                }
            }
        }
    }

    // Parmi les check-nodes non-voisins de v, choisir le plus éloigné
    // (ou le moins chargé en cas d'égalité)
    let mut best = (0usize, 0usize, usize::MAX); // (dist, chk_id, charge)
    for c in 0..n_chk {
        if current_neighbors.contains(&c) {
            continue;
        }
        let d = dist_chk[c];
        let load = chk_to_var[c].len();
        if d > best.0 || (d == best.0 && load < best.2) {
            best = (d, c, load);
        }
    }
    // Fallback : check le moins chargé si aucun atteignable
    if best.0 == 0 {
        (0..n_chk).min_by_key(|&c| chk_to_var[c].len()).unwrap_or(0)
    } else {
        best.1
    }
}

// #[cfg(test)]
// mod tests {
//     use super::*;
//
//     #[test]
//     fn test_gallager_regular_degrees() {
//         let params = LdpcParams {
//             n: 12,
//             k: 6,
//             topology: CodeTopology::Regular { wc: 2, wr: 4 },
//             generation: GenerationMethod::Gallager,
//             seed: Some(42),
//         };
//         let code = LdpcCode::new(params).unwrap();
//         for col in 0..code.n() {
//             assert_eq!(code.h.col_weight(col), 2);
//         }
//     }
//
//     #[test]
//     fn test_mackay_neal_no_4cycle() {
//         let params = LdpcParams {
//             n: 20,
//             k: 10,
//             topology: CodeTopology::Regular { wc: 3, wr: 6 },
//             generation: GenerationMethod::MacKayNeal { max_attempts: 1000 },
//             seed: Some(0),
//         };
//         let code = LdpcCode::new(params).unwrap();
//         assert!(!code.graph.has_4_cycle());
//     }
//
//     #[test]
//     fn test_peg_girth_at_least_6() {
//         let params = LdpcParams {
//             n: 30,
//             k: 15,
//             topology: CodeTopology::Regular { wc: 3, wr: 6 },
//             generation: GenerationMethod::Peg,
//             seed: None,
//         };
//         let code = LdpcCode::new(params).unwrap();
//         assert!(code.girth() >= 6);
//     }
// }