pub(crate) const TENSOR_SYMBOL_PREFIX: &str = "get_tensor_";
pub(crate) const CONSTANT_SYMBOL_PREFIX: &str = "get_constant_";

macro_rules! for_each_external_symbol {
    ($callback:ident) => {
        $callback! {
            "barrier_init" => barrier_init,
            "set_constants" => set_constants,
            "set_u0" => set_u0,
            "reset" => reset,
            "reset_grad" => reset_grad,
            "reset_rgrad" => reset_rgrad,
            "reset_sgrad" => reset_sgrad,
            "reset_srgrad" => reset_srgrad,
            "rhs" => rhs,
            "rhs_grad" => rhs_grad,
            "rhs_rgrad" => rhs_rgrad,
            "rhs_sgrad" => rhs_sgrad,
            "rhs_srgrad" => rhs_srgrad,
            "mass" => mass,
            "mass_rgrad" => mass_rgrad,
            "set_u0_grad" => set_u0_grad,
            "set_u0_rgrad" => set_u0_rgrad,
            "set_u0_sgrad" => set_u0_sgrad,
            "calc_out" => calc_out,
            "calc_out_grad" => calc_out_grad,
            "calc_out_rgrad" => calc_out_rgrad,
            "calc_out_sgrad" => calc_out_sgrad,
            "calc_out_srgrad" => calc_out_srgrad,
            "calc_stop" => calc_stop,
            "calc_stop_grad" => calc_stop_grad,
            "calc_stop_rgrad" => calc_stop_rgrad,
            "calc_stop_sgrad" => calc_stop_sgrad,
            "calc_stop_srgrad" => calc_stop_srgrad,
            "set_id" => set_id,
            "get_dims" => get_dims,
            "set_inputs" => set_inputs,
            "get_inputs" => get_inputs,
            "set_inputs_grad" => set_inputs_grad,
            "set_inputs_rgrad" => set_inputs_rgrad,
        }
    };
}

#[allow(unused_imports)]
pub(crate) use for_each_external_symbol;

#[allow(unused_macros)]
macro_rules! insert_external_symbols {
    ($symbols:expr, $sym:path) => {{
        use $sym as sym;
        $symbols.insert("barrier_init".to_string(), sym::barrier_init as *const u8);
        $symbols.insert("set_constants".to_string(), sym::set_constants as *const u8);
        $symbols.insert("set_u0".to_string(), sym::set_u0 as *const u8);
        $symbols.insert("reset".to_string(), sym::reset as *const u8);
        $symbols.insert("reset_grad".to_string(), sym::reset_grad as *const u8);
        $symbols.insert("reset_rgrad".to_string(), sym::reset_rgrad as *const u8);
        $symbols.insert("reset_sgrad".to_string(), sym::reset_sgrad as *const u8);
        $symbols.insert("reset_srgrad".to_string(), sym::reset_srgrad as *const u8);
        $symbols.insert("rhs".to_string(), sym::rhs as *const u8);
        $symbols.insert("rhs_grad".to_string(), sym::rhs_grad as *const u8);
        $symbols.insert("rhs_rgrad".to_string(), sym::rhs_rgrad as *const u8);
        $symbols.insert("rhs_sgrad".to_string(), sym::rhs_sgrad as *const u8);
        $symbols.insert("rhs_srgrad".to_string(), sym::rhs_srgrad as *const u8);
        $symbols.insert("mass".to_string(), sym::mass as *const u8);
        $symbols.insert("mass_rgrad".to_string(), sym::mass_rgrad as *const u8);
        $symbols.insert("set_u0_grad".to_string(), sym::set_u0_grad as *const u8);
        $symbols.insert("set_u0_rgrad".to_string(), sym::set_u0_rgrad as *const u8);
        $symbols.insert("set_u0_sgrad".to_string(), sym::set_u0_sgrad as *const u8);
        $symbols.insert("calc_out".to_string(), sym::calc_out as *const u8);
        $symbols.insert("calc_out_grad".to_string(), sym::calc_out_grad as *const u8);
        $symbols.insert(
            "calc_out_rgrad".to_string(),
            sym::calc_out_rgrad as *const u8,
        );
        $symbols.insert(
            "calc_out_sgrad".to_string(),
            sym::calc_out_sgrad as *const u8,
        );
        $symbols.insert(
            "calc_out_srgrad".to_string(),
            sym::calc_out_srgrad as *const u8,
        );
        $symbols.insert("calc_stop".to_string(), sym::calc_stop as *const u8);
        $symbols.insert(
            "calc_stop_grad".to_string(),
            sym::calc_stop_grad as *const u8,
        );
        $symbols.insert(
            "calc_stop_rgrad".to_string(),
            sym::calc_stop_rgrad as *const u8,
        );
        $symbols.insert(
            "calc_stop_sgrad".to_string(),
            sym::calc_stop_sgrad as *const u8,
        );
        $symbols.insert(
            "calc_stop_srgrad".to_string(),
            sym::calc_stop_srgrad as *const u8,
        );
        $symbols.insert("set_id".to_string(), sym::set_id as *const u8);
        $symbols.insert("get_dims".to_string(), sym::get_dims as *const u8);
        $symbols.insert("set_inputs".to_string(), sym::set_inputs as *const u8);
        $symbols.insert("get_inputs".to_string(), sym::get_inputs as *const u8);
        $symbols.insert(
            "set_inputs_grad".to_string(),
            sym::set_inputs_grad as *const u8,
        );
        $symbols.insert(
            "set_inputs_rgrad".to_string(),
            sym::set_inputs_rgrad as *const u8,
        );
    }};
}

#[allow(unused_imports)]
pub(crate) use insert_external_symbols;

macro_rules! collect_external_symbol_names {
    ($($name:literal => $func:ident,)+) => {
        pub(crate) const EXTERNAL_SYMBOL_NAMES: &[&str] = &[$($name),+];
    };
}

for_each_external_symbol!(collect_external_symbol_names);

pub(crate) const REQUIRED_EXTERNAL_SYMBOL_NAMES: &[&str] = &[
    "set_constants",
    "set_u0",
    "reset",
    "reset_grad",
    "reset_rgrad",
    "reset_sgrad",
    "reset_srgrad",
    "rhs",
    "rhs_grad",
    "rhs_rgrad",
    "rhs_sgrad",
    "rhs_srgrad",
    "mass",
    "mass_rgrad",
    "set_u0_grad",
    "set_u0_rgrad",
    "set_u0_sgrad",
    "calc_out",
    "calc_out_grad",
    "calc_out_rgrad",
    "calc_out_sgrad",
    "calc_out_srgrad",
    "calc_stop",
    "calc_stop_grad",
    "calc_stop_rgrad",
    "calc_stop_sgrad",
    "calc_stop_srgrad",
    "set_id",
    "get_dims",
    "set_inputs",
    "get_inputs",
    "set_inputs_grad",
    "set_inputs_rgrad",
];

pub(crate) fn normalize_symbol_name(name: &str) -> &str {
    name.strip_prefix('_').unwrap_or(name)
}

pub(crate) fn is_external_symbol_name(name: &str) -> bool {
    let name = normalize_symbol_name(name);
    EXTERNAL_SYMBOL_NAMES.contains(&name)
        || name.starts_with(TENSOR_SYMBOL_PREFIX)
        || name.starts_with(CONSTANT_SYMBOL_PREFIX)
}
