use cord_expr::ExprInfo;
use cord_trig::ir::{NodeId, TrigGraph, TrigOp};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Op {
    Union,
    Intersection,
    Difference,
    SmoothUnion,
    Add,
    Subtract,
    Multiply,
    Divide,
    Power,
    InvPower,
}

impl Op {
    pub const ALL: &[Op] = &[
        Op::Union,
        Op::Intersection,
        Op::Difference,
        Op::SmoothUnion,
        Op::Add,
        Op::Subtract,
        Op::Multiply,
        Op::Divide,
        Op::Power,
        Op::InvPower,
    ];

    pub fn label(self) -> &'static str {
        match self {
            Op::Union => "union(A, B)",
            Op::Intersection => "intersect(A, B)",
            Op::Difference => "diff(A, B)",
            Op::SmoothUnion => "smooth(A, B, k)",
            Op::Add => "A + B",
            Op::Subtract => "A - B",
            Op::Multiply => "A \u{00d7} B",
            Op::Divide => "A / B",
            Op::Power => "A ^ B",
            Op::InvPower => "A ^ -B",
        }
    }

    pub fn check(self, a: &ExprInfo, b: &ExprInfo) -> Option<&'static str> {
        match self {
            Op::Power | Op::InvPower => Some("general power not yet in IR"),

            Op::Union | Op::Intersection | Op::Difference | Op::SmoothUnion => {
                if a.dimensions == 0 && b.dimensions == 0 {
                    return Some("both expressions are constants");
                }
                if a.dimensions != b.dimensions && a.dimensions != 0 && b.dimensions != 0 {
                    return Some("dimension mismatch");
                }
                None
            }

            Op::Divide => None,
            Op::Add | Op::Subtract | Op::Multiply => None,
        }
    }
}

impl std::fmt::Display for Op {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.label())
    }
}

pub fn combine(a: &TrigGraph, b: &TrigGraph, op: Op, smooth_k: f64) -> TrigGraph {
    let mut graph = TrigGraph::new();

    let a_map = remap_nodes(&mut graph, a);
    let a_out = a_map[a.output as usize];

    let b_map = remap_nodes(&mut graph, b);
    let b_out = b_map[b.output as usize];

    let result = match op {
        Op::Add => graph.push(TrigOp::Add(a_out, b_out)),
        Op::Subtract => graph.push(TrigOp::Sub(a_out, b_out)),
        Op::Multiply => graph.push(TrigOp::Mul(a_out, b_out)),
        Op::Divide => graph.push(TrigOp::Div(a_out, b_out)),
        Op::Power => graph.push(TrigOp::Mul(a_out, a_out)),
        Op::InvPower => {
            let neg_b = graph.push(TrigOp::Neg(b_out));
            graph.push(TrigOp::Mul(a_out, neg_b))
        }
        Op::Union => graph.push(TrigOp::Min(a_out, b_out)),
        Op::Intersection => graph.push(TrigOp::Max(a_out, b_out)),
        Op::Difference => {
            let neg_b = graph.push(TrigOp::Neg(b_out));
            graph.push(TrigOp::Max(a_out, neg_b))
        }
        Op::SmoothUnion => build_smooth_union(&mut graph, a_out, b_out, smooth_k),
    };

    graph.set_output(result);
    graph
}

fn remap_nodes(target: &mut TrigGraph, source: &TrigGraph) -> Vec<NodeId> {
    let mut map = Vec::with_capacity(source.nodes.len());

    for op in &source.nodes {
        let new_op = match op {
            TrigOp::InputX => TrigOp::InputX,
            TrigOp::InputY => TrigOp::InputY,
            TrigOp::InputZ => TrigOp::InputZ,
            TrigOp::Const(c) => TrigOp::Const(*c),
            TrigOp::Add(a, b) => TrigOp::Add(map[*a as usize], map[*b as usize]),
            TrigOp::Sub(a, b) => TrigOp::Sub(map[*a as usize], map[*b as usize]),
            TrigOp::Mul(a, b) => TrigOp::Mul(map[*a as usize], map[*b as usize]),
            TrigOp::Div(a, b) => TrigOp::Div(map[*a as usize], map[*b as usize]),
            TrigOp::Neg(a) => TrigOp::Neg(map[*a as usize]),
            TrigOp::Abs(a) => TrigOp::Abs(map[*a as usize]),
            TrigOp::Sin(a) => TrigOp::Sin(map[*a as usize]),
            TrigOp::Cos(a) => TrigOp::Cos(map[*a as usize]),
            TrigOp::Tan(a) => TrigOp::Tan(map[*a as usize]),
            TrigOp::Asin(a) => TrigOp::Asin(map[*a as usize]),
            TrigOp::Acos(a) => TrigOp::Acos(map[*a as usize]),
            TrigOp::Atan(a) => TrigOp::Atan(map[*a as usize]),
            TrigOp::Sinh(a) => TrigOp::Sinh(map[*a as usize]),
            TrigOp::Cosh(a) => TrigOp::Cosh(map[*a as usize]),
            TrigOp::Tanh(a) => TrigOp::Tanh(map[*a as usize]),
            TrigOp::Asinh(a) => TrigOp::Asinh(map[*a as usize]),
            TrigOp::Acosh(a) => TrigOp::Acosh(map[*a as usize]),
            TrigOp::Atanh(a) => TrigOp::Atanh(map[*a as usize]),
            TrigOp::Sqrt(a) => TrigOp::Sqrt(map[*a as usize]),
            TrigOp::Exp(a) => TrigOp::Exp(map[*a as usize]),
            TrigOp::Ln(a) => TrigOp::Ln(map[*a as usize]),
            TrigOp::Hypot(a, b) => TrigOp::Hypot(map[*a as usize], map[*b as usize]),
            TrigOp::Atan2(a, b) => TrigOp::Atan2(map[*a as usize], map[*b as usize]),
            TrigOp::Min(a, b) => TrigOp::Min(map[*a as usize], map[*b as usize]),
            TrigOp::Max(a, b) => TrigOp::Max(map[*a as usize], map[*b as usize]),
            TrigOp::Clamp { val, lo, hi } => TrigOp::Clamp {
                val: map[*val as usize],
                lo: map[*lo as usize],
                hi: map[*hi as usize],
            },
        };
        map.push(target.push(new_op));
    }

    map
}

/// smooth_min(a, b, k) = min(a,b) - h^2 * k * 0.25
/// where h = clamp(0.5 + 0.5*(b-a)/k, 0, 1)
fn build_smooth_union(g: &mut TrigGraph, a: NodeId, b: NodeId, k: f64) -> NodeId {
    let k_node = g.push(TrigOp::Const(k));
    let half = g.push(TrigOp::Const(0.5));
    let quarter = g.push(TrigOp::Const(0.25));
    let zero = g.push(TrigOp::Const(0.0));
    let one = g.push(TrigOp::Const(1.0));

    let diff = g.push(TrigOp::Sub(b, a));
    let inv_k = g.push(TrigOp::Const(1.0 / k));
    let ratio = g.push(TrigOp::Mul(diff, inv_k));
    let scaled = g.push(TrigOp::Mul(half, ratio));
    let h_raw = g.push(TrigOp::Add(half, scaled));
    let h = g.push(TrigOp::Clamp { val: h_raw, lo: zero, hi: one });
    let h2 = g.push(TrigOp::Mul(h, h));
    let h2k = g.push(TrigOp::Mul(h2, k_node));
    let correction = g.push(TrigOp::Mul(h2k, quarter));
    let m = g.push(TrigOp::Min(a, b));
    g.push(TrigOp::Sub(m, correction))
}

#[cfg(test)]
mod tests {
    use super::*;
    use cord_expr::{classify, parse_expr};
    use cord_trig::eval::evaluate;

    #[test]
    fn combine_add() {
        let a = parse_expr("x").unwrap();
        let b = parse_expr("y").unwrap();
        let c = combine(&a, &b, Op::Add, 0.0);
        assert!((evaluate(&c, 3.0, 4.0, 0.0) - 7.0).abs() < 1e-10);
    }

    #[test]
    fn combine_union() {
        let a = parse_expr("x").unwrap();
        let b = parse_expr("y").unwrap();
        let c = combine(&a, &b, Op::Union, 0.0);
        assert!((evaluate(&c, 3.0, 4.0, 0.0) - 3.0).abs() < 1e-10);
    }

    #[test]
    fn validity_checks() {
        let a_info = classify(&parse_expr("sin(x) + y").unwrap());
        let b_info = classify(&parse_expr("z").unwrap());
        assert!(Op::Union.check(&a_info, &b_info).is_some());
        assert!(Op::Add.check(&a_info, &b_info).is_none());
    }
}
