use crate::{Gf2, LdpcError, Llr, Result};

// Trait Channel

pub trait Channel: Send + Sync {
    fn transmit(&self, codeword: &[Gf2], rng: &mut impl rand::Rng) -> Vec<Llr>;
    fn capacity(&self) -> f64;
}

// Canal AWGN
// Modulation BPSK : 0 -> +1.0, 1 -> -1.0
// Signal reçu : y = x + n, n eq N(0, sig²)
// LLR optimal : L(y) = 2y/sig²
// sig² = N_0/2 = 1/(2 R SNR_lin)

#[derive(Debug, Clone)]
pub struct AwgnChannel {
    pub snr_db: f64,
    sigma: f64,
}

impl AwgnChannel {
    pub fn new(snr_db: f64, code_rate: f64) -> Result<Self> {
        if !(0.0..1.0).contains(&code_rate) {
            return Err(LdpcError::OutOfRange("code_rate ∈ ]0, 1[".into()));
        }
        let snr_lin = 10.0_f64.powf(snr_db / 10.0);
        let sigma = (1.0 / (2.0 * code_rate * snr_lin)).sqrt();
        Ok(Self { snr_db, sigma })
    }

    pub fn sigma(&self) -> f64 {
        self.sigma
    }
    pub fn snr_linear(&self) -> f64 {
        10.0_f64.powf(self.snr_db / 10.0)
    }

    #[inline]
    pub fn llr_from_received(y: f64, sigma: f64) -> Llr {
        2.0 * y / (sigma * sigma)
    }
}

impl Channel for AwgnChannel {
    fn transmit(&self, codeword: &[Gf2], rng: &mut impl rand::Rng) -> Vec<Llr> {
        use rand_distr::{Distribution, Normal};
        let normal = Normal::new(0.0, self.sigma).unwrap();
        codeword
            .iter()
            .map(|&b| {
                let x = if b == 0 { 1.0_f64 } else { -1.0_f64 };
                let y = x + normal.sample(rng);
                Self::llr_from_received(y, self.sigma)
            })
            .collect()
    }

    fn capacity(&self) -> f64 {
        // Capacité BPSK-AWGN par Monte-Carlo
        use rand_distr::{Distribution, Normal};
        let mut rng = rand::thread_rng();
        let normal = Normal::new(0.0, self.sigma).unwrap();
        let n_samples = 10_000usize;
        let mut sum = 0.0f64;
        for _ in 0..n_samples {
            let n: f64 = normal.sample(&mut rng);
            let y = 1.0 + n; // bit 0 transmis (x=+1)
            let llr = Self::llr_from_received(y, self.sigma);
            // I = 1 - E[log2(1 + exp(-2y/sig²))]
            sum += (1.0 + (-llr).exp()).log2();
        }
        1.0 - sum / n_samples as f64
    }
}

// Canal BSC
// Chaque bit inversé avec probabilité p
// LLR : +-log((1-p)/p) selon le bit reçu

#[derive(Debug, Clone)]
pub struct BscChannel {
    pub crossover_prob: f64,
    llr_magnitude: Llr,
}

impl BscChannel {
    pub fn new(crossover_prob: f64) -> Result<Self> {
        if crossover_prob <= 0.0 || crossover_prob >= 0.5 {
            return Err(LdpcError::OutOfRange("p ∈ ]0, 0.5[".into()));
        }
        let llr_magnitude = ((1.0 - crossover_prob) / crossover_prob).ln();
        Ok(Self {
            crossover_prob,
            llr_magnitude,
        })
    }
}

impl Channel for BscChannel {
    fn transmit(&self, codeword: &[Gf2], rng: &mut impl rand::Rng) -> Vec<Llr> {
        codeword
            .iter()
            .map(|&b| {
                let rcv = if rng.gen::<f64>() < self.crossover_prob {
                    b ^ 1
                } else {
                    b
                };
                if rcv == 0 {
                    self.llr_magnitude
                } else {
                    -self.llr_magnitude
                }
            })
            .collect()
    }

    // C_BSC(p) = 1 - Hb(p) avec Hb = entropie binaire
    fn capacity(&self) -> f64 {
        let p = self.crossover_prob;
        let hb = -p * p.log2() - (1.0 - p) * (1.0 - p).log2();
        1.0 - hb
    }
}

// Canal BEC
// Bit effacé avec probabilité ε -> LLR = 0 (incertitude totale)
// Bit reçu correctement → LLR = +-CERTAIN_LLR (grand mais fini)

#[derive(Debug, Clone)]
pub struct BecChannel {
    pub erasure_prob: f64,
}

impl BecChannel {
    pub fn new(erasure_prob: f64) -> Result<Self> {
        if erasure_prob <= 0.0 || erasure_prob >= 1.0 {
            return Err(LdpcError::OutOfRange("ε ∈ ]0, 1[".into()));
        }
        Ok(Self { erasure_prob })
    }

    // Grand mais fini pour stabilité numérique
    const CERTAIN_LLR: Llr = 100.0;
}

impl Channel for BecChannel {
    fn transmit(&self, codeword: &[Gf2], rng: &mut impl rand::Rng) -> Vec<Llr> {
        codeword
            .iter()
            .map(|&b| {
                if rng.gen::<f64>() < self.erasure_prob {
                    0.0 // Effacement
                } else if b == 0 {
                    Self::CERTAIN_LLR
                } else {
                    -Self::CERTAIN_LLR
                }
            })
            .collect()
    }

    /// C_BEC(ε) = 1 - ε
    fn capacity(&self) -> f64 {
        1.0 - self.erasure_prob
    }
}

// #[cfg(test)]
// mod tests {
//     use super::*;
//
//     #[test]
//     fn test_awgn_llr_formula() {
//         assert!((AwgnChannel::llr_from_received(1.0, 1.0) - 2.0).abs() < 1e-12);
//         assert!((AwgnChannel::llr_from_received(-1.0, 1.0) + 2.0).abs() < 1e-12);
//     }
//
//     #[test]
//     fn test_bec_capacity() {
//         let bec = BecChannel::new(0.3).unwrap();
//         assert!((bec.capacity() - 0.7).abs() < 1e-12);
//     }
//
//     #[test]
//     fn test_bsc_capacity_bounds() {
//         let bsc = BscChannel::new(0.1).unwrap();
//         let cap = bsc.capacity();
//         assert!(cap > 0.0 && cap < 1.0);
//     }
// }