use burn::module::Module;
use burn::nn::{Linear, LinearConfig};
use burn::tensor::backend::Backend;
use burn::tensor::activation::silu;
use burn::tensor::{Float, Tensor};

#[derive(Module, Debug)]
pub struct FeedForward<B: Backend> {
    pub gate_proj: Linear<B>,
    pub up_proj: Linear<B>,
    pub down_proj: Linear<B>,
}

impl<B: Backend> FeedForward<B> {
    pub fn new(hidden_size: usize, intermediate_size: usize, device: &B::Device) -> Self {
        let gate_proj = LinearConfig::new(hidden_size, intermediate_size).init(device);
        let up_proj = LinearConfig::new(hidden_size, intermediate_size).init(device);
        let down_proj = LinearConfig::new(intermediate_size, hidden_size).init(device);

        Self {
            gate_proj,
            up_proj,
            down_proj,
        }
    }

    pub fn forward(&self, x: Tensor<B, 3, Float>) -> Tensor<B, 3, Float> {
        let gate = self.gate_proj.forward(x.clone());
        let up = self.up_proj.forward(x);
        let activated = silu(gate) * up;
        self.down_proj.forward(activated)
    }
}