use crate::iq_reader::IqSample;

// Automatic Gain Control
pub struct Agc<I> {
    inner: I,
    // Previous power estimate
    pub power_estimate: f32,
    // Previous gain
    pub current_gain: f32,
    pub target_power: f32,
    pub alpha_attack: f32,
    pub alpha_release: f32,
    pub beta: f32,
    pub min_gain: f32,
    pub max_gain: f32,
}

impl<I> Agc<I> {
    pub fn new(
        inner: I,
        sample_rate: f32,
        target_power: f32,
        min_gain: f32,
        max_gain: f32,
    ) -> Self {
        // Target attack time 5 ms
        let tau_attack = 0.005;

        // Target release time 50 ms
        let tau_release = 0.05;

        let alpha_attack = 1.0 - f32::exp(-1.0 / (sample_rate * tau_attack));
        let alpha_release = 1.0 - f32::exp(-1.0 / (sample_rate * tau_release));

        let beta = 0.999;

        let power_estimate = 0.0;
        let current_gain = 1.0;

        Self {
            inner,
            power_estimate,
            current_gain,
            target_power,
            alpha_attack,
            alpha_release,
            beta,
            min_gain,
            max_gain,
        }
    }

    pub fn process_sample(&mut self, sample: IqSample) -> IqSample {
        let i = sample.re;
        let q = sample.im;

        // Instant Power
        let inst_power = i * i + q * q;

        let alpha = if inst_power > self.power_estimate {
            self.alpha_attack
        } else {
            self.alpha_release
        };

        // IIR filter
        let power_estimate = alpha * inst_power + (1.0 - alpha) * self.power_estimate;

        // Update Power
        self.power_estimate = power_estimate.max(1e-10);

        let raw_gain = (self.target_power / self.power_estimate).sqrt();

        // Gain in [min_gain ; max_gain]
        let raw_gain = match raw_gain {
            g if g < self.min_gain => self.min_gain,
            g if g > self.max_gain => self.max_gain,
            _ => raw_gain,
        };

        let final_gain = self.beta * self.current_gain + (1.0 - self.beta) * raw_gain;

        self.current_gain = final_gain;

        IqSample::new(i * final_gain, q * final_gain)
    }
}

impl<I, E> Iterator for Agc<I>
where
    I: Iterator<Item = Result<IqSample, E>>,
{
    type Item = Result<IqSample, E>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.inner.next()? {
            Ok(sample) => Some(Ok(self.process_sample(sample))),
            Err(e) => Some(Err(e)),
        }
    }
}
