// Native decoder composition fixture. Dimensions/data are deliberately small;
// the layers and gradients are the same reusable operations for larger models.
import aner.tensor
import aner.nn
import aner.dev

class DecoderBlock {
    var qkv: Tensor
    var projection: Tensor
    var expand: Tensor
    var contract: Tensor
    var gain1: Tensor
    var bias1: Tensor
    var gain2: Tensor
    var bias2: Tensor

    fn forward(x: Tensor, mask: Tensor) -> Tensor {
        let batch = tensor.dim(x, 0)
        let length = tensor.dim(x, 1)
        let normalized = nn.layer_norm(x, gain1, bias1, epsilon: 0.00001)
        let projected = tensor.batched_matmul(normalized, qkv)
        // Width 4, two heads of width 2. Each axis remains explicit.
        let query = tensor.permute(tensor.reshape(tensor.slice(projected,
            axis: 2, start: 0, length: 4), List(batch, length, 2, 2)), List(0, 2, 1, 3))
        let key = tensor.permute(tensor.reshape(tensor.slice(projected,
            axis: 2, start: 4, length: 4), List(batch, length, 2, 2)), List(0, 2, 1, 3))
        let value = tensor.permute(tensor.reshape(tensor.slice(projected,
            axis: 2, start: 8, length: 4), List(batch, length, 2, 2)), List(0, 2, 1, 3))
        let heads = nn.attention(query, key, value, mask, causal: true, offset: 0)
        let joined = tensor.reshape(tensor.permute(heads, List(0, 2, 1, 3)), List(batch, length, 4))
        let residual = nn.add(x, tensor.batched_matmul(joined, projection))
        let hidden = nn.gelu(tensor.batched_matmul(
            nn.layer_norm(residual, gain2, bias2, epsilon: 0.00001), expand))
        return nn.add(residual, tensor.batched_matmul(hidden, contract))
    }

    fn update(optimizer: Optimizer) {
        // Stable names own optimizer slots. Every distinct parameter is updated once.
        qkv = nn.update(optimizer, qkv, "block.qkv")
        projection = nn.update(optimizer, projection, "block.projection")
        expand = nn.update(optimizer, expand, "block.expand")
        contract = nn.update(optimizer, contract, "block.contract")
        gain1 = nn.update(optimizer, gain1, "block.norm1.gain")
        bias1 = nn.update(optimizer, bias1, "block.norm1.bias")
        gain2 = nn.update(optimizer, gain2, "block.norm2.gain")
        bias2 = nn.update(optimizer, bias2, "block.norm2.bias")
    }
}

let block = DecoderBlock(
    qkv: nn.parameter(nn.xavier(4, 12, 41)),
    projection: nn.parameter(nn.xavier(4, 4, 42)),
    expand: nn.parameter(nn.xavier(4, 16, 43)),
    contract: nn.parameter(nn.xavier(16, 4, 44)),
    gain1: nn.parameter([[1.0, 1.0, 1.0, 1.0]]),
    bias1: nn.parameter(nn.zeros(1, 4)),
    gain2: nn.parameter([[1.0, 1.0, 1.0, 1.0]]),
    bias2: nn.parameter(nn.zeros(1, 4)))
var tokens = nn.parameter(nn.random(8, 4, 45, 0.2))
var positions = nn.parameter(nn.random(4, 4, 46, 0.1))
let optimizer = nn.optimizer("adam", rate: 0.02)
let ids = List(0, 1, 2, 3, 1, 2, 3, 0)
let targets = List(1, 2, 3, 4, 2, 3, 4, -1)
let position_ids = List(0, 1, 2, 3)
// Separate padding policy from the causal policy. Padding is excluded from loss.
let mask = tensor.reshape([[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0]], List(2, 1, 1, 4))
var initial = 0.0
var final = 0.0
var step = 0
while step < 40 {
    let embedded = nn.embedding(tokens, ids, shape: List(2, 4))
    let positioned = nn.add(embedded, tensor.broadcast_to(
        nn.embedding(positions, position_ids, shape: List(1, 4)), List(2, 4, 4)))
    let hidden = block.forward(positioned, mask)
    // Tied vocabulary projection: both uses accumulate into the same token table.
    let logits = tensor.batched_matmul(hidden, nn.transpose(tokens))
    let loss = nn.cross_entropy_ids(logits, targets, ignore_index: -1)
    if step == 0 { initial = nn.item(loss) }
    final = nn.item(loss)
    nn.backward(loss)
    if step == 0 {
        dev.capture("First sequence logits", logits)
        dev.capture("Token table gradients", nn.grad(tokens))
    }
    tokens = nn.update(optimizer, tokens, "token.embedding")
    positions = nn.update(optimizer, positions, "position.embedding")
    block.update(optimizer)
    step = step + 1
}
print("Initial token loss")
print(initial)
print("Final token loss")
print(final)