// Native inference cache correctness on new seeded reference weights.
// This checks mechanics; it is not a trained language-quality demonstration.
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
}

nn.configure("cpu", "Float64", "preserve")
let codec = token.byte_tokenizer()
let model = gen.decoder(codec, context: 24, width: 8,
    heads: 2, layers: 2, seed: 63, dropout: 0.1)
let identity = model.fingerprint()
let prompt = codec.encode("ab", bos: true, eos: false)

// Omitting the trailing option selects a new cache local to this generate call.
let automatic = model.generate(prompt, max_new_tokens: 8,
    temperature: 0.0, top_k: 0, top_p: 1.0, seed: 7)
let cached = model.generate(prompt, max_new_tokens: 8,
    temperature: 0.0, top_k: 0, top_p: 1.0, seed: 7, use_cache: true)
let recomputed = model.generate(prompt, max_new_tokens: 8,
    temperature: 0.0, top_k: 0, top_p: 1.0, seed: 7, use_cache: false)
verify("cached_default", same_ids(automatic, cached))
verify("greedy_reference", same_ids(cached, recomputed))

let sampled = model.generate(prompt, max_new_tokens: 8,
    temperature: 0.8, top_k: 20, top_p: 0.95, seed: 71)
let reference_sample = model.generate(prompt, max_new_tokens: 8,
    temperature: 0.8, top_k: 20, top_p: 0.95, seed: 71, use_cache: false)
verify("seeded_reference", same_ids(sampled, reference_sample))
let another = codec.encode("xy", bos: true, eos: false)
let independent = model.generate(another, 8, 0.0, 0, 1.0, 7)
let independent_reference = model.generate(another, 8, 0.0, 0, 1.0, 7, false)
verify("independent_calls", same_ids(independent, independent_reference))
verify("zero_output", model.generate(prompt, 0, 0.0, 0, 1.0, 7).len() == 0)
verify("model_unchanged", model.fingerprint() == identity && model.steps() == 0)
print("cached_greedy_hex")
print(codec.decode_bytes(cached, skip_special: true).hex())
print("cache_checks_passed")