use egg::{Id, RecExpr};
use std::collections::HashMap;

use crate::ir::{
    BinaryOp, Builtin, Const, DType, DTypeKind, FloatEncoding, Node, NodeId, TypedDag, UnaryOp,
};
use crate::opt::OptError;
use crate::opt::lang::MathLang;

#[derive(Debug, Clone)]
pub struct TypeMap {
    /// One entry per [`MathLang`] node in the stripped [`RecExpr`] (same length as `expr.len()`).
    pub node_types: Vec<DType>,
    pub param_types: Vec<DType>,
}

fn float_dtype_for_complex_parts(c: DType) -> DType {
    match (c.kind, c.bits) {
        (DTypeKind::Complex(FloatEncoding::Ieee), 64) => DType::F32,
        (DTypeKind::Complex(FloatEncoding::Ieee), 128) => DType::F64,
        _ => DType::F64,
    }
}

pub fn strip_types(dag: &TypedDag) -> Result<(RecExpr<MathLang>, TypeMap), OptError> {
    let mut expr = RecExpr::<MathLang>::default();
    let mut id_map: HashMap<NodeId, Id> = HashMap::with_capacity(dag.nodes.len());
    let mut rec_dtypes: Vec<DType> = Vec::new();

    for (node_id, node) in dag.nodes.iter().enumerate() {
        let id = match node {
            Node::Param(param_id) => {
                let pid = expr.add(MathLang::Int(param_id.0 as i64));
                rec_dtypes.push(DType::I64);
                let id = expr.add(MathLang::Param([pid]));
                rec_dtypes.push(dag.dtypes[node_id]);
                id
            }
            Node::Const(Const::Bool(v)) => {
                let id = expr.add(MathLang::Bool(*v));
                rec_dtypes.push(dag.dtypes[node_id]);
                id
            }
            Node::Const(Const::Int(v)) => {
                let id = expr.add(MathLang::Int(*v));
                rec_dtypes.push(dag.dtypes[node_id]);
                id
            }
            Node::Const(Const::Float(v)) => {
                let id = expr.add(MathLang::Num((*v).into()));
                rec_dtypes.push(dag.dtypes[node_id]);
                id
            }
            Node::Const(Const::Complex(re, im)) => {
                let part_ty = float_dtype_for_complex_parts(dag.dtypes[node_id]);
                let re_id = expr.add(MathLang::Num((*re).into()));
                rec_dtypes.push(part_ty);
                let im_id = expr.add(MathLang::Num((*im).into()));
                rec_dtypes.push(part_ty);
                let id = expr.add(MathLang::Complex([re_id, im_id]));
                rec_dtypes.push(dag.dtypes[node_id]);
                id
            }
            Node::Unary(op, a) => {
                let a_id = *id_map.get(a).ok_or_else(|| OptError::InvalidExpr {
                    msg: format!("missing unary child mapping for node {node_id}"),
                })?;
                let id = match op {
                    UnaryOp::Neg => expr.add(MathLang::Neg([a_id])),
                    UnaryOp::Not => expr.add(MathLang::Not([a_id])),
                };
                rec_dtypes.push(dag.dtypes[node_id]);
                id
            }
            Node::Binary(op, a, b) => {
                let a_id = *id_map.get(a).ok_or_else(|| OptError::InvalidExpr {
                    msg: format!("missing binary lhs mapping for node {node_id}"),
                })?;
                let b_id = *id_map.get(b).ok_or_else(|| OptError::InvalidExpr {
                    msg: format!("missing binary rhs mapping for node {node_id}"),
                })?;
                let id = match op {
                    BinaryOp::Add => expr.add(MathLang::Add([a_id, b_id])),
                    BinaryOp::Sub => expr.add(MathLang::Sub([a_id, b_id])),
                    BinaryOp::Mul => expr.add(MathLang::Mul([a_id, b_id])),
                    BinaryOp::Div => expr.add(MathLang::Div([a_id, b_id])),
                    BinaryOp::FloorDiv => expr.add(MathLang::FloorDiv([a_id, b_id])),
                    BinaryOp::Rem => expr.add(MathLang::Rem([a_id, b_id])),
                    BinaryOp::Pow => expr.add(MathLang::Pow([a_id, b_id])),
                    BinaryOp::Eq => expr.add(MathLang::Eq([a_id, b_id])),
                    BinaryOp::Ne => expr.add(MathLang::Ne([a_id, b_id])),
                    BinaryOp::Lt => expr.add(MathLang::Lt([a_id, b_id])),
                    BinaryOp::Le => expr.add(MathLang::Le([a_id, b_id])),
                    BinaryOp::Gt => expr.add(MathLang::Gt([a_id, b_id])),
                    BinaryOp::Ge => expr.add(MathLang::Ge([a_id, b_id])),
                    BinaryOp::And => expr.add(MathLang::And([a_id, b_id])),
                    BinaryOp::Or => expr.add(MathLang::Or([a_id, b_id])),
                };
                rec_dtypes.push(dag.dtypes[node_id]);
                id
            }
            Node::IfElse(c, t, e) => {
                let c_id = *id_map.get(c).ok_or_else(|| OptError::InvalidExpr {
                    msg: format!("missing if condition mapping for node {node_id}"),
                })?;
                let t_id = *id_map.get(t).ok_or_else(|| OptError::InvalidExpr {
                    msg: format!("missing if true branch mapping for node {node_id}"),
                })?;
                let e_id = *id_map.get(e).ok_or_else(|| OptError::InvalidExpr {
                    msg: format!("missing if false branch mapping for node {node_id}"),
                })?;
                let id = expr.add(MathLang::IfElse([c_id, t_id, e_id]));
                rec_dtypes.push(dag.dtypes[node_id]);
                id
            }
            Node::Call(builtin, args) => {
                if args.len() != 1 {
                    return Err(OptError::UnsupportedCall {
                        name: builtin.to_string(),
                    });
                }
                let arg_id = *id_map.get(&args[0]).ok_or_else(|| OptError::InvalidExpr {
                    msg: format!("missing call arg mapping for node {node_id}"),
                })?;
                let id = match builtin {
                    Builtin::Exp => expr.add(MathLang::Exp([arg_id])),
                    Builtin::Log => expr.add(MathLang::Log([arg_id])),
                    Builtin::Sqrt => expr.add(MathLang::Sqrt([arg_id])),
                    Builtin::Abs => expr.add(MathLang::Abs([arg_id])),
                    Builtin::Sin => expr.add(MathLang::Sin([arg_id])),
                    Builtin::Cos => expr.add(MathLang::Cos([arg_id])),
                    Builtin::Tanh => expr.add(MathLang::Tanh([arg_id])),
                    Builtin::Sigmoid => expr.add(MathLang::Sigmoid([arg_id])),
                    Builtin::Relu => expr.add(MathLang::Relu([arg_id])),
                    Builtin::Softplus => expr.add(MathLang::Softplus([arg_id])),
                    Builtin::Floor => expr.add(MathLang::Floor([arg_id])),
                    Builtin::Ceil => expr.add(MathLang::Ceil([arg_id])),
                    Builtin::ExpSigmoid => expr.add(MathLang::FusedExpSigmoid([arg_id])),
                    Builtin::SigmoidLog => expr.add(MathLang::FusedSigmoidLog([arg_id])),
                };
                rec_dtypes.push(dag.dtypes[node_id]);
                id
            }
            Node::Cast(dtype, child, _) => {
                let token = expr.add(MathLang::TypeToken(dtype.to_str().into()));
                // `TypeMap::node_types` is indexed by `RecExpr` position, not restored DAG node id.
                // Keep a dtype slot for the token so later expr indices stay aligned during restore.
                rec_dtypes.push(*dtype);
                let child_id = *id_map.get(child).ok_or_else(|| OptError::InvalidExpr {
                    msg: format!("missing cast child mapping for node {node_id}"),
                })?;
                let id = expr.add(MathLang::Cast([token, child_id]));
                rec_dtypes.push(*dtype);
                id
            }
        };
        id_map.insert(NodeId(node_id), id);
    }

    debug_assert_eq!(rec_dtypes.len(), expr.as_ref().len());

    let type_map = TypeMap {
        node_types: rec_dtypes,
        param_types: dag.params.iter().map(|(_, dt)| *dt).collect(),
    };
    Ok((expr, type_map))
}

#[cfg(test)]
mod tests {
    use super::strip_types;
    use crate::ir::{BinaryOp, DType, Node, NodeId, ParamId, TypedDag};
    use crate::opt::lang::MathLang;

    #[test]
    fn cast_typetoken_keeps_typemap_aligned_with_expr_indices() {
        let dag = TypedDag::try_new(
            vec![
                Node::Param(ParamId(0)),
                Node::Cast(DType::F32, NodeId(0), false),
                Node::Binary(BinaryOp::Add, NodeId(1), NodeId(1)),
            ],
            vec![DType::F64, DType::F32, DType::F32],
            vec![("x".to_string(), DType::F64)],
            vec![NodeId(2)],
        )
        .expect("well-formed test dag");

        let (expr, type_map) = strip_types(&dag).expect("strip should succeed");
        assert_eq!(type_map.node_types.len(), expr.as_ref().len());

        assert!(matches!(expr.as_ref()[0], MathLang::Int(0)));
        assert!(matches!(expr.as_ref()[1], MathLang::Param(_)));
        assert!(matches!(expr.as_ref()[2], MathLang::TypeToken(_)));
        assert!(matches!(expr.as_ref()[3], MathLang::Cast(_)));
        assert!(matches!(expr.as_ref()[4], MathLang::Add(_)));

        assert_eq!(type_map.node_types[0], DType::I64);
        assert_eq!(type_map.node_types[1], DType::F64);
        assert_eq!(type_map.node_types[2], DType::F32);
        assert_eq!(type_map.node_types[3], DType::F32);
        assert_eq!(type_map.node_types[4], DType::F32);
    }
}
