use num_complex::Complex;

use crate::iq_reader::IqChunk;

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

impl Agc {
    fn new(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 {
            power_estimate,
            current_gain,
            target_power,
            alpha_attack,
            alpha_release,
            beta,
            min_gain,
            max_gain,
        }
    }

    fn process_chunk(&mut self, chunk: &mut IqChunk) {
        for z in chunk.samples.iter_mut() {
            let i = z.re;
            let q = z.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;

            *z = Complex::new(i * final_gain, q * final_gain);
        }
    }
}