// Limbic Affect WGSL Kernel
//
// One forward step of the limbic model's recurrent affect core, run on the GPU.
// Given the previous hidden state, previous affective state, and an experience
// embedding, it produces the next hidden state and a bounded affect *delta*.
//
// This mirrors the CPU reference in limbic/limbic_model.ts exactly (same math):
//   pre[j]  = Σ_i Win[j,i]·x[i] + Σ_k Ws[j,k]·s[k]
//   a[j]    = sigmoid(A[j])                       // per-channel recurrence gate
//   h'[j]   = a[j]·h[j] + (1-a[j])·tanh(pre[j])   // SSM-style leak/input
//   Δ[k]    = tanh( Σ_j Wout[k,j]·h'[j] + b[k] )  // bounded affect change
//
// hidden_dim is assumed ≤ 64 (the limbic head is tiny — the heavy lifting is in
// the hippocampus SSM that produces the experience embedding). Reward is NOT
// computed here: it is only needed during training, which runs the CPU backward
// path; the per-turn inference step only needs Δ and the next hidden state.

export const LIMBIC_AFFECT_WGSL: string = /* wgsl */ `

struct Dims {
    input_dim  : u32,
    hidden_dim : u32,
    state_dim  : u32,
    _pad       : u32,
};

@group(0) @binding(0)  var<uniform>             dims        : Dims;
@group(0) @binding(1)  var<storage, read>       win         : array<f32>;  // hidden*input
@group(0) @binding(2)  var<storage, read>       ws          : array<f32>;  // hidden*state
@group(0) @binding(3)  var<storage, read>       a_logit     : array<f32>;  // hidden
@group(0) @binding(4)  var<storage, read>       wout_state  : array<f32>;  // state*hidden
@group(0) @binding(5)  var<storage, read>       bout_state  : array<f32>;  // state
@group(0) @binding(6)  var<storage, read>       x_in        : array<f32>;  // input
@group(0) @binding(7)  var<storage, read>       h_prev      : array<f32>;  // hidden
@group(0) @binding(8)  var<storage, read>       s_prev      : array<f32>;  // state
@group(0) @binding(9)  var<storage, read_write> h_out       : array<f32>;  // hidden
@group(0) @binding(10) var<storage, read_write> delta_out   : array<f32>;  // state

var<workgroup> hbuf : array<f32, 64>;

// Single-workgroup dispatch: (1, 1, 1) with workgroup_size 64.
@compute @workgroup_size(64, 1, 1)
fn affect_step(
    @builtin(local_invocation_id) lid : vec3<u32>,
) {
    let j = lid.x;

    // Pass 1: recurrent hidden update (one thread per hidden channel).
    if (j < dims.hidden_dim) {
        var pre : f32 = 0.0;
        for (var i : u32 = 0u; i < dims.input_dim; i = i + 1u) {
            pre = pre + win[j * dims.input_dim + i] * x_in[i];
        }
        for (var k : u32 = 0u; k < dims.state_dim; k = k + 1u) {
            pre = pre + ws[j * dims.state_dim + k] * s_prev[k];
        }
        let a  = 1.0 / (1.0 + exp(-a_logit[j]));
        let hn = a * h_prev[j] + (1.0 - a) * tanh(pre);
        hbuf[j]  = hn;
        h_out[j] = hn;
    }

    workgroupBarrier();

    // Pass 2: bounded affect delta (one thread per state dim).
    if (j < dims.state_dim) {
        var acc : f32 = bout_state[j];
        for (var m : u32 = 0u; m < dims.hidden_dim; m = m + 1u) {
            acc = acc + wout_state[j * dims.hidden_dim + m] * hbuf[m];
        }
        delta_out[j] = tanh(acc);
    }
}
`;