//! Feed-forward network: SwiGLU variant used by Llama.
//!
//! Llama FFN: output = down_proj(silu(gate_proj(x)) * up_proj(x))

use crate::tensor::{matmul_2d, silu_inplace, Tensor};

/// Compute FFN for one layer.
///
/// Weights must be pre-transposed: [in, out] layout for direct matmul.
/// x: [seq_len, hidden_size]
/// gate_w: [hidden_size, intermediate_size] (pre-transposed)
/// up_w: [hidden_size, intermediate_size] (pre-transposed)
/// down_w: [intermediate_size, hidden_size] (pre-transposed)
pub fn ffn_forward(
    x: &Tensor,
    gate_w: &Tensor,
    up_w: &Tensor,
    down_w: &Tensor,
) -> Tensor {
    let mut gate = matmul_2d(x, gate_w).expect("gate matmul");
    silu_inplace(&mut gate);

    let up = matmul_2d(x, up_w).expect("up matmul");

    gate.mul_inplace(&up).expect("gate * up");

    matmul_2d(&gate, down_w).expect("down matmul")
}

/// Transpose a 2D tensor [M, N] -> [N, M].
pub fn transpose_2d(t: &Tensor) -> Tensor {
    assert_eq!(t.shape.len(), 2);
    let (m, n) = (t.shape[0], t.shape[1]);
    let mut out = vec![0.0f32; m * n];
    for i in 0..m {
        for j in 0..n {
            out[j * m + i] = t.data[i * n + j];
        }
    }
    Tensor {
        data: out,
        shape: vec![n, m],
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_transpose_2d() {
        let t = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
        let tt = transpose_2d(&t);
        assert_eq!(tt.shape, vec![3, 2]);
        assert_eq!(tt.data, vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
    }

    #[test]
    fn test_ffn_output_shape() {
        let hidden = 4;
        let inter = 8;
        let seq = 2;
        let x = Tensor::from_data(vec![0.1; seq * hidden], vec![seq, hidden]).unwrap();
        // Weights are pre-transposed: [hidden, inter] and [inter, hidden]
        let gate_w = Tensor::from_data(vec![0.1; hidden * inter], vec![hidden, inter]).unwrap();
        let up_w = Tensor::from_data(vec![0.1; hidden * inter], vec![hidden, inter]).unwrap();
        let down_w = Tensor::from_data(vec![0.1; inter * hidden], vec![inter, hidden]).unwrap();
        let out = ffn_forward(&x, &gate_w, &up_w, &down_w);
        assert_eq!(out.shape, vec![seq, hidden]);
    }
}
