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
// BPSK mod 0 -> +1.0, 1 -> -1.0
// Signal recu y = x + n, n ~ N(0, sig^2)
// LLR optimal L(y) = 2y/sig^2
// sig^2 = 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 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^2))]
            sum += (1.0 + (-llr).exp()).log2();
        }
        1.0 - sum / n_samples as f64
    }
}
