use cord_trig::ir::{NodeId, TrigOp};
use crate::parser::ExprParser;

impl<'a> ExprParser<'a> {
    pub(crate) fn build_saw(&mut self, x: NodeId) -> Result<NodeId, String> {
        let s = self.graph.push(TrigOp::Sin(x));
        let c = self.graph.push(TrigOp::Cos(x));
        let a = self.graph.push(TrigOp::Atan2(s, c));
        let inv_pi = self.graph.push(TrigOp::Const(1.0 / std::f64::consts::PI));
        Ok(self.graph.push(TrigOp::Mul(a, inv_pi)))
    }

    pub(crate) fn build_tri(&mut self, x: NodeId) -> Result<NodeId, String> {
        let saw = self.build_saw(x)?;
        let a = self.graph.push(TrigOp::Abs(saw));
        let two = self.graph.push(TrigOp::Const(2.0));
        let scaled = self.graph.push(TrigOp::Mul(two, a));
        let one = self.graph.push(TrigOp::Const(1.0));
        Ok(self.graph.push(TrigOp::Sub(scaled, one)))
    }

    pub(crate) fn build_square(&mut self, x: NodeId) -> Result<NodeId, String> {
        let s = self.graph.push(TrigOp::Sin(x));
        let k = self.graph.push(TrigOp::Const(1000.0));
        let raw = self.graph.push(TrigOp::Mul(s, k));
        let lo = self.graph.push(TrigOp::Const(-1.0));
        let hi = self.graph.push(TrigOp::Const(1.0));
        Ok(self.graph.push(TrigOp::Clamp { val: raw, lo, hi }))
    }

    pub(crate) fn build_mix(&mut self, a: NodeId, b: NodeId, t: NodeId) -> Result<NodeId, String> {
        let one = self.graph.push(TrigOp::Const(1.0));
        let omt = self.graph.push(TrigOp::Sub(one, t));
        let at = self.graph.push(TrigOp::Mul(a, omt));
        let bt = self.graph.push(TrigOp::Mul(b, t));
        Ok(self.graph.push(TrigOp::Add(at, bt)))
    }

    pub(crate) fn build_smoothstep(&mut self, lo: NodeId, hi: NodeId, x: NodeId) -> Result<NodeId, String> {
        let lo_val = match self.graph.nodes.get(lo as usize) {
            Some(TrigOp::Const(v)) => *v,
            _ => return Err("smoothstep: lo must be a constant".into()),
        };
        let hi_val = match self.graph.nodes.get(hi as usize) {
            Some(TrigOp::Const(v)) => *v,
            _ => return Err("smoothstep: hi must be a constant".into()),
        };
        let range = hi_val - lo_val;
        if range.abs() < 1e-15 {
            return Err("smoothstep: lo and hi must differ".into());
        }
        let inv_range = self.graph.push(TrigOp::Const(1.0 / range));
        let diff = self.graph.push(TrigOp::Sub(x, lo));
        let raw = self.graph.push(TrigOp::Mul(diff, inv_range));
        let zero = self.graph.push(TrigOp::Const(0.0));
        let one = self.graph.push(TrigOp::Const(1.0));
        let t = self.graph.push(TrigOp::Clamp { val: raw, lo: zero, hi: one });
        let t2 = self.graph.push(TrigOp::Mul(t, t));
        let two = self.graph.push(TrigOp::Const(2.0));
        let three = self.graph.push(TrigOp::Const(3.0));
        let two_t = self.graph.push(TrigOp::Mul(two, t));
        let coeff = self.graph.push(TrigOp::Sub(three, two_t));
        Ok(self.graph.push(TrigOp::Mul(t2, coeff)))
    }

    pub(crate) fn build_quantize(&mut self, x: NodeId, n: NodeId) -> Result<NodeId, String> {
        let n_val = match self.graph.nodes.get(n as usize) {
            Some(TrigOp::Const(v)) if *v >= 1.0 => *v,
            _ => return Err("quantize: n must be a constant >= 1".into()),
        };
        let xn = self.graph.push(TrigOp::Mul(x, n));
        let two_pi = self.graph.push(TrigOp::Const(2.0 * std::f64::consts::PI));
        let phase = self.graph.push(TrigOp::Mul(xn, two_pi));
        let s = self.graph.push(TrigOp::Sin(phase));
        let c = self.graph.push(TrigOp::Cos(phase));
        let a = self.graph.push(TrigOp::Atan2(s, c));
        let inv_two_pi = self.graph.push(TrigOp::Const(1.0 / (2.0 * std::f64::consts::PI)));
        let fract_centered = self.graph.push(TrigOp::Mul(a, inv_two_pi));
        let half = self.graph.push(TrigOp::Const(0.5));
        let fract = self.graph.push(TrigOp::Add(fract_centered, half));
        let floor = self.graph.push(TrigOp::Sub(xn, fract));
        let inv_n = self.graph.push(TrigOp::Const(1.0 / n_val));
        Ok(self.graph.push(TrigOp::Mul(floor, inv_n)))
    }

    pub(crate) fn build_am(&mut self, signal: NodeId, modulator: NodeId, depth: NodeId) -> Result<NodeId, String> {
        let dm = self.graph.push(TrigOp::Mul(depth, modulator));
        let one = self.graph.push(TrigOp::Const(1.0));
        let env = self.graph.push(TrigOp::Add(one, dm));
        Ok(self.graph.push(TrigOp::Mul(signal, env)))
    }

    pub(crate) fn build_fm(&mut self, carrier: NodeId, modulator: NodeId, depth: NodeId) -> Result<NodeId, String> {
        let dm = self.graph.push(TrigOp::Mul(depth, modulator));
        let phase = self.graph.push(TrigOp::Add(carrier, dm));
        Ok(self.graph.push(TrigOp::Sin(phase)))
    }

    pub(crate) fn build_lpf_node(&mut self, signal: NodeId, k: NodeId) -> NodeId {
        let a = self.graph.push(TrigOp::Atan2(signal, k));
        let ka = self.graph.push(TrigOp::Mul(k, a));
        let coeff = self.graph.push(TrigOp::Const(2.0 / std::f64::consts::PI));
        self.graph.push(TrigOp::Mul(ka, coeff))
    }

    pub(crate) fn build_lpf(&mut self, signal: NodeId, k: NodeId) -> Result<NodeId, String> {
        Ok(self.build_lpf_node(signal, k))
    }

    pub(crate) fn build_hpf(&mut self, signal: NodeId, k: NodeId) -> Result<NodeId, String> {
        let lp = self.build_lpf_node(signal, k);
        Ok(self.graph.push(TrigOp::Sub(signal, lp)))
    }

    pub(crate) fn build_bpf(&mut self, signal: NodeId, lo: NodeId, hi: NodeId) -> Result<NodeId, String> {
        let lp_hi = self.build_lpf_node(signal, hi);
        let lp_lo = self.build_lpf_node(signal, lo);
        Ok(self.graph.push(TrigOp::Sub(lp_hi, lp_lo)))
    }

    pub(crate) fn build_dft(&mut self, signal: NodeId, n: NodeId) -> Result<NodeId, String> {
        let n_val = match self.graph.nodes.get(n as usize) {
            Some(TrigOp::Const(v)) if *v >= 1.0 && *v <= 64.0 => *v as u32,
            _ => return Err("dft: n must be a constant between 1 and 64".into()),
        };
        let mut sum: Option<NodeId> = None;
        for k in 1..=n_val {
            let kf = self.graph.push(TrigOp::Const(k as f64));
            let ks = self.graph.push(TrigOp::Mul(kf, signal));
            let s = self.graph.push(TrigOp::Sin(ks));
            let inv_k = self.graph.push(TrigOp::Const(1.0 / k as f64));
            let term = self.graph.push(TrigOp::Mul(s, inv_k));
            sum = Some(match sum {
                None => term,
                Some(prev) => self.graph.push(TrigOp::Add(prev, term)),
            });
        }
        Ok(sum.unwrap())
    }
}

#[cfg(test)]
mod tests {
    use crate::parse_expr;
    use cord_trig::eval::evaluate;

    #[test]
    fn saw_wave() {
        let g = parse_expr("saw(x)").unwrap();
        assert!(evaluate(&g, 0.0, 0.0, 0.0).abs() < 1e-10);
        assert!((evaluate(&g, std::f64::consts::FRAC_PI_2, 0.0, 0.0) - 0.5).abs() < 1e-10);
    }

    #[test]
    fn tri_wave() {
        let g = parse_expr("tri(x)").unwrap();
        assert!((evaluate(&g, 0.0, 0.0, 0.0) - -1.0).abs() < 1e-10);
    }

    #[test]
    fn square_wave() {
        let g = parse_expr("square(x)").unwrap();
        assert!((evaluate(&g, std::f64::consts::FRAC_PI_4, 0.0, 0.0) - 1.0).abs() < 1e-6);
        assert!((evaluate(&g, -std::f64::consts::FRAC_PI_4, 0.0, 0.0) - -1.0).abs() < 1e-6);
    }

    #[test]
    fn mix_lerp() {
        let g = parse_expr("mix(x, y, 0.25)").unwrap();
        assert!((evaluate(&g, 10.0, 20.0, 0.0) - 12.5).abs() < 1e-10);
    }

    #[test]
    fn smoothstep_edges() {
        let g = parse_expr("smoothstep(0, 1, x)").unwrap();
        assert!(evaluate(&g, -1.0, 0.0, 0.0).abs() < 1e-10);
        assert!((evaluate(&g, 2.0, 0.0, 0.0) - 1.0).abs() < 1e-10);
        assert!((evaluate(&g, 0.5, 0.0, 0.0) - 0.5).abs() < 1e-10);
    }

    #[test]
    fn fm_synthesis() {
        let g = parse_expr("fm(x, y, 2)").unwrap();
        assert!(evaluate(&g, 0.0, 0.0, 0.0).abs() < 1e-10);
    }

    #[test]
    fn lpf_saturates() {
        let g = parse_expr("lpf(x, 1)").unwrap();
        assert!(evaluate(&g, 0.0, 0.0, 0.0).abs() < 1e-10);
        let val = evaluate(&g, 1000.0, 0.0, 0.0);
        assert!((val - 1.0).abs() < 0.01);
    }

    #[test]
    fn dft_single_harmonic() {
        let g = parse_expr("dft(x, 1)").unwrap();
        let val = evaluate(&g, std::f64::consts::FRAC_PI_2, 0.0, 0.0);
        assert!((val - 1.0).abs() < 1e-10);
    }

    #[test]
    fn envelope_and_phase() {
        let g_env = parse_expr("envelope(x)").unwrap();
        let g_phase = parse_expr("phase(x)").unwrap();
        assert!((evaluate(&g_env, 3.0, 0.0, 0.0) - 10.0_f64.sqrt()).abs() < 1e-10);
        assert!((evaluate(&g_phase, 1.0, 0.0, 0.0) - std::f64::consts::FRAC_PI_4).abs() < 1e-10);
    }
}
