//! Activation operator code generation.

use proc_macro2::TokenStream;
use quote::quote;

use super::sanitize_name;
use crate::error::Result;
use crate::ir::Node;

/// Generate softmax operation.
pub fn generate_softmax(node: &Node) -> Result<TokenStream> {
    let input = &node.inputs[0];
    let output = &node.outputs[0];

    let in_ident = quote::format_ident!("{}", sanitize_name(input));
    let out_ident = quote::format_ident!("{}", sanitize_name(output));

    // Get axis attribute (default -1)
    let axis = node.get_int("axis").unwrap_or(-1);

    Ok(quote! {
        let #out_ident = #in_ident.softmax(#axis)?;
    })
}

/// Generate leaky ReLU operation.
pub fn generate_leaky_relu(node: &Node) -> Result<TokenStream> {
    let input = &node.inputs[0];
    let output = &node.outputs[0];

    let in_ident = quote::format_ident!("{}", sanitize_name(input));
    let out_ident = quote::format_ident!("{}", sanitize_name(output));

    // Get alpha attribute (default 0.01)
    let alpha = node.get_float("alpha").unwrap_or(0.01);

    Ok(quote! {
        let #out_ident = #in_ident.leaky_relu(#alpha)?;
    })
}