// Native packed CUDA KV generation on bounded seeded research weights.
// This demonstrates inference mechanisms, not trained language quality.
import aner.token
import aner.gen
import aner.nn

fn verify(label: String, condition: Bool) -> Unit {
    print("check_" + label)
    if !condition {
        let failure: Result<Bool, String> = Result.err(label)
        print(failure.unwrap())
    }
    print(condition)
}
fn same_ids(left: List, right: List) -> Bool {
    if left.len() != right.len() { return false }
    var index = 0
    while index < left.len() {
        if cast<Int64>(left.get(index)) != cast<Int64>(right.get(index)) { return false }
        index = index + 1
    }
    return true
}

// CPU Float64 and Metal Float32 provide a portable reference with this API.
nn.configure("cuda", "Float32", "preserve")
let codec = token.byte_tokenizer()
let model = gen.decoder(codec, context: 24, width: 16,
    heads: 4, layers: 2, seed: 63, dropout: 0.0)
let identity = model.fingerprint()
let first = codec.encode("ab", bos: true, eos: false)
let second = codec.encode("science", bos: true, eos: false)
let prompts = List(first, List(), second, first)
let budgets = List(4, 2, 0, 1)
let seeds = List(71, 19, 23, 71)
let packed = model.generate_batch(prompts, max_new_tokens: budgets,
    temperature: 0.8, top_k: 20, top_p: 0.95, seeds: seeds)
let reference = model.generate_batch(prompts, budgets, 0.8, 20, 0.95, seeds, false)
verify("ragged_reference", same_ids(cast<List>(packed.get(0)), cast<List>(reference.get(0))))
verify("empty_prompt_bos", same_ids(cast<List>(packed.get(1)),
    model.generate(List(), 2, 0.8, 20, 0.95, 19)))
verify("zero_budget", cast<List>(packed.get(2)).len() == 0)
verify("independent_budget", same_ids(cast<List>(packed.get(3)),
    model.generate(first, 1, 0.8, 20, 0.95, 71)))
let reordered = model.generate_batch(List(List(), first), List(2, 4),
    0.8, 20, 0.95, List(19, 71), true)
verify("row_order_rng", same_ids(cast<List>(packed.get(0)), cast<List>(reordered.get(1))))
let shared = model.generate_batch(List(first, first), 2, 0.0, 0, 1.0, List(7, 99))
verify("shared_budget", same_ids(cast<List>(shared.get(0)), cast<List>(shared.get(1))))
let zero = model.generate_batch(prompts, List(0, 0, 0, 0), 0.0, 0, 1.0, seeds)
verify("all_zero", zero.len() == 4 && cast<List>(zero.get(0)).len() == 0)
verify("model_unchanged", model.fingerprint() == identity && model.steps() == 0)
print("batch_checks_passed")