// silkie_dynamics.sy -- Visualize H-H gate dynamics during generation.
// See the refractory curve. See where she loops. See WHY.
// Then: biological temperature + biological repetition suppression.
// Aislinn McCubbins + Kimi, April 2026

fn main() {
    println("=== SILKIE DYNAMICS VISUALIZATION ===")

    let n = load_model("examples/silkie_full.safetensors")
    let ns = create_neuron_stack(1, 512, 64)
    let cs = create_cascade()
    let ts = load_tokenizer("examples/silkie_tokenizer.json")

    let embed_w = model_get("token_embedding.weight")

    // Prompt: "Hello Silkie" (pre-encoded)
    let generated = tensor_from_ints(3591, 1207)
    let generated = t_reshape(generated, 1, 2)
    let seq_len = 2

    // Biological temperature (O-U noise sigma)
    // Higher = more exploration, lower = more exploitation
    // Biological temperature: sigma=1.0 means "body temperature" (37C)
    // The kernel scales by kT/q = 0.0267 internally
    // sigma > 1.0 = fever (more noise), < 1.0 = hypothermia (less noise)
    let bio_temp = 1.0

    println("TOKEN | WORD          | m_avg  h_avg  gate   | ACTION")
    println("------+---------------+----------------------+---------")

    let prev_token = 0

    for step in 0..25 {
        // Forward pass
        let flat = t_reshape(generated, seq_len)
        let x = embedding_lookup(embed_w, flat)
        let x = t_reshape(x, 1, seq_len, 512)

        for layer_idx in 0..4 {
            let prefix = "layers." + layer_idx
            let x = transformer_block(
                x,
                model_get(prefix + ".attention.in_proj_weight"),
                model_get(prefix + ".attention.in_proj_bias"),
                model_get(prefix + ".attention.out_proj.weight"),
                model_get(prefix + ".attention.out_proj.bias"),
                model_get(prefix + ".ff.0.weight"),
                model_get(prefix + ".ff.0.bias"),
                model_get(prefix + ".ff.3.weight"),
                model_get(prefix + ".ff.3.bias"),
                model_get(prefix + ".norm1.weight"),
                model_get(prefix + ".norm1.bias"),
                model_get(prefix + ".norm2.weight"),
                model_get(prefix + ".norm2.bias"), 8)
        }

        // Apply biological temperature BEFORE neurons process
        neuron_stack_apply_noise(bio_temp)

        let x = neuron_stack_forward(x)
        let cascade_out = cascade_forward(x)
        let x = x + cascade_out

        for layer_idx in 4..8 {
            let prefix = "layers." + layer_idx
            let x = transformer_block(
                x,
                model_get(prefix + ".attention.in_proj_weight"),
                model_get(prefix + ".attention.in_proj_bias"),
                model_get(prefix + ".attention.out_proj.weight"),
                model_get(prefix + ".attention.out_proj.bias"),
                model_get(prefix + ".ff.0.weight"),
                model_get(prefix + ".ff.0.bias"),
                model_get(prefix + ".ff.3.weight"),
                model_get(prefix + ".ff.3.bias"),
                model_get(prefix + ".norm1.weight"),
                model_get(prefix + ".norm1.bias"),
                model_get(prefix + ".norm2.weight"),
                model_get(prefix + ".norm2.bias"), 8)
        }

        // Get neuron states AFTER processing
        let states = neuron_stack_dump_states()
        // Average across all 22 neurons
        let m_avg = t_mean(t_slice(states, -1, 1, 2))
        let h_avg = t_mean(t_slice(states, -1, 2, 3))
        let gate_avg = t_mean(t_slice(states, -1, 3, 4))

        // Output head
        let x = layer_norm(x, model_get("layers.0.norm1.weight"), model_get("layers.0.norm1.bias"))
        let logits = matmul(x, transpose(embed_w))
        let logits_2d = t_reshape(logits, seq_len, 4100)
        let last_logits = t_slice(logits_2d, 0, seq_len - 1, seq_len)

        // Sample (top-k = 50)
        let masked = topk_mask(last_logits, 50)
        let probs = softmax(masked)
        let next_token = t_multinomial(probs)
        let next_id = tensor_sum(next_token)

        // Decode token
        let token_text = detokenize(next_token)

        // Check for repeat: if same token as last, apply biological suppression
        let action = ""
        // We cant do if on floats easily, so just always apply mild suppression
        // No artificial suppression — pure biological dynamics only
        // The Q10 scaling + thermal noise IS the repetition mechanism

        // Print dynamics
        println("  " + step + "   | " + token_text + "	| " + tensor_sum(m_avg) + " " + tensor_sum(h_avg) + " " + tensor_sum(gate_avg) + " |")

        let prev_token = next_id
        let generated = t_cat(generated, t_reshape(next_token, 1, 1))
        let seq_len = seq_len + 1
    }

    println("")
    println("SILKIE SAYS:")
    let flat_out = t_reshape(generated, seq_len)
    println(detokenize(flat_out))
    println("")
    println("=== DYNAMICS COMPLETE ===")
}
