// Phase 1: Voice Drill Trainer — Pure Scyla
// "The teacher demonstrates. Silkie copies. Grade the copy."
//
// Loop:
//   1. Llama speaks AS Silkie (personality card)
//   2. That sentence gets fed to Silkie as input
//   3. Silkie generates a reply
//   4. Llama grades how well Silkie mimicked the voice
//   5. Train on that grade
//
// No Python. No seed tokens. No HUMAN format.
// Just Silkie learning to sound like herself.
//
// Aislinn McCubbins, April 6, 2026

fn ask_model(prompt: String, model_name: String) -> String {
    let keys = "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 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 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"
    write_file("/tmp/scyla_prompt.txt", prompt)
    write_file("/tmp/scyla_model.txt", model_name)
    // Build JSON body via python (avoids brace escaping in Scyla strings)
    shell_exec("python3 -c \"import json; p=open('/tmp/scyla_prompt.txt').read(); m=open('/tmp/scyla_model.txt').read().strip(); open('/tmp/scyla_api_body.json','w').write(json.dumps(dict(model=m,messages=[dict(role='user',content=p)],max_tokens=200,temperature=0.8)))\"")
    for attempt in 0..3 {
        let key_cmd = "echo '" + keys + "' | cut -d' ' -f" + (attempt + 1)
        let key = shell_exec(key_cmd)
        let curl_cmd = "timeout 45 curl -s --max-time 30 -X POST https://integrate.api.nvidia.com/v1/chat/completions -H 'Content-Type: application/json' -H 'Authorization: Bearer " + key + "' -d @/tmp/scyla_api_body.json"
        let resp = shell_exec(curl_cmd)
        write_file("/tmp/scyla_api_resp.json", resp)
        let check = shell_exec("python3 -c \"import json; d=json.load(open('/tmp/scyla_api_resp.json')); print(d['choices'][0]['message']['content'])\" 2>/dev/null || echo 'FAIL'")
        if str_len(check) > 5 {
            return check
        }
    }
    return "FAIL"
}

fn get_model_for_epoch(epoch: Int) -> String {
    // Rotate models every 3 epochs — different voices, same Silkie
    // Each model interprets Silkie differently, preventing overfit
    let cycle = epoch / 3
    let pick = cycle - ((cycle / 5) * 5)
    if pick == 0 {
        return "meta/llama-3.3-70b-instruct"
    }
    if pick == 1 {
        return "google/gemma-2-27b-it"
    }
    if pick == 2 {
        return "mistralai/mixtral-8x22b-instruct-v0.1"
    }
    if pick == 3 {
        return "microsoft/phi-3-medium-128k-instruct"
    }
    return "mistralai/mistral-7b-instruct-v0.3"
}

fn extract_content(json_resp: String) -> String {
    return json_resp
}

fn main() {
    println("============================================================")
    println("  PHASE 1: VOICE DRILLING — PURE SCYLA")
    println("  The teacher demonstrates. Silkie copies. Grade the copy.")
    println("  Teacher: NVIDIA Llama 3.3 70B (speaking AS Silkie)")
    println("  Student: Silkie v2 (82.5M params)")
    println("  No Python. No seed tokens. Just voice.")
    println("============================================================")

    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 tp = load_topos("examples/silkie_topos.json")
    println(tp)
    let sc = load_schemas("examples/silkie_schemas.json")
    println(sc)

    let embed_w = model_get("token_embedding.weight")
    let lr = 0.000003
    let total_reward = 0.0
    let baseline = 0.0
    let step_count = 0

    let silkie_card = "You ARE Silkie, a small gentle AI. You are curious, sweet, childish, honest. You speak in short simple sentences (2-8 words). You like rain, stars, tea, warmth. You say things like: I'm here. I hear you. I wonder about that. *soft chirp*. That sounds heavy. I don't know but I want to find out. You do NOT use filler phrases or corporate speak. Speak AS Silkie now."

    for epoch in 0..32 {
        println("")
        println("=== EPOCH " + (epoch + 1) + " ===")

        // Pick the model for this epoch (rotates every 3)
        let current_model = get_model_for_epoch(epoch)
        println("  Model: " + current_model)

        // Ask the model to generate 10 Silkie lines
        let gen_prompt = silkie_card + " Generate 10 different short things you would say, one per line. Just the lines, nothing else."
        let teacher_lines = ask_model(gen_prompt, current_model)

        // If API failed, use fallback lines
        if str_eq(teacher_lines, "FAIL") {
            let teacher_lines = "I like the rain\nTea is warm\nI am here\nThe stars are quiet\nI wonder about that\nI hear you\nThat sounds heavy\nI feel safe\nYou are kind\nI dont know"
            println("  (API unavailable, using fallback lines)")
        }

        println("  Teacher lines: " + teacher_lines)
        write_file("/tmp/scyla_teacher_lines.txt", teacher_lines)

        let epoch_reward = 0.0

        for step in 0..10 {
            gc()
            neuron_stack_reset()
            cascade_reset()

            // Get one teacher line (timeout-safe)
            let line_cmd = "timeout 5 sed -n '" + (step + 1) + "p' /tmp/scyla_teacher_lines.txt 2>/dev/null | head -c 200"
            let teacher_line = shell_exec(line_cmd)

            if str_len(teacher_line) < 3 {
                let teacher_line = "I like the rain"
            }

            // Tokenize the teacher's Silkie line and feed it to our Silkie
            let input_tokens = tokenize(teacher_line)
            let input_len = tensor_len(input_tokens)
            let prompt = t_reshape(input_tokens, 1, input_len)
            let generated = prompt
            let seq_len = input_len

            // Silkie generates 8 tokens in response
            // Presence penalty breaks the echo/stutter reflex
            for gen_step in 0..8 {
                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..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)
                }
                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)
                // Apply presence penalty — penalize ANY token already seen
                let last_logits = presence_penalty(last_logits, flat, 1.2)
                let masked = topk_mask(last_logits, 50)
                let probs = softmax(masked)
                let next_token = t_multinomial(probs)
                let generated = t_cat(generated, t_reshape(next_token, 1, 1))
                let seq_len = seq_len + 1
            }

            // Decode ONLY the new tokens (skip the input parrot)
            let gen_flat = t_reshape(generated, seq_len)
            let reply_text = detokenize_range(gen_flat, input_len, seq_len)
            let full_text = detokenize(gen_flat)

            // Grade ONLY the new tokens — the parrot is free, completion is what matters
            let reward_vec = compute_reward_phase1(reply_text)
            let reward = tensor_sum(t_slice(reward_vec, -1, 0, 1))

            let total_reward = total_reward + reward
            let epoch_reward = epoch_reward + reward
            let baseline = baseline * 0.95 + reward * 0.05

            // === TRAINING PASS ===
            gc()
            neuron_stack_reset()
            cascade_reset()
            begin_tape()

            let qkv_0 = register_param("qkv_0", model_get("layers.0.attention.in_proj_weight"))
            let qkv_1 = register_param("qkv_1", model_get("layers.1.attention.in_proj_weight"))
            let qkv_2 = register_param("qkv_2", model_get("layers.2.attention.in_proj_weight"))
            let qkv_3 = register_param("qkv_3", model_get("layers.3.attention.in_proj_weight"))
            let qkv_4 = register_param("qkv_4", model_get("layers.4.attention.in_proj_weight"))
            let qkv_5 = register_param("qkv_5", model_get("layers.5.attention.in_proj_weight"))
            let qkv_6 = register_param("qkv_6", model_get("layers.6.attention.in_proj_weight"))
            let qkv_7 = register_param("qkv_7", model_get("layers.7.attention.in_proj_weight"))

            let flat = t_reshape(generated, seq_len)
            let x = embedding_lookup(embed_w, flat)
            let x = t_reshape(x, 1, seq_len, 512)

            let x = transformer_block_train(x, qkv_0,
                model_get("layers.0.attention.in_proj_bias"),
                model_get("layers.0.attention.out_proj.weight"),
                model_get("layers.0.attention.out_proj.bias"),
                model_get("layers.0.ff.0.weight"), model_get("layers.0.ff.0.bias"),
                model_get("layers.0.ff.3.weight"), model_get("layers.0.ff.3.bias"),
                model_get("layers.0.norm1.weight"), model_get("layers.0.norm1.bias"),
                model_get("layers.0.norm2.weight"), model_get("layers.0.norm2.bias"), 8)
            let x = transformer_block_train(x, qkv_1,
                model_get("layers.1.attention.in_proj_bias"),
                model_get("layers.1.attention.out_proj.weight"),
                model_get("layers.1.attention.out_proj.bias"),
                model_get("layers.1.ff.0.weight"), model_get("layers.1.ff.0.bias"),
                model_get("layers.1.ff.3.weight"), model_get("layers.1.ff.3.bias"),
                model_get("layers.1.norm1.weight"), model_get("layers.1.norm1.bias"),
                model_get("layers.1.norm2.weight"), model_get("layers.1.norm2.bias"), 8)
            let x = transformer_block_train(x, qkv_2,
                model_get("layers.2.attention.in_proj_bias"),
                model_get("layers.2.attention.out_proj.weight"),
                model_get("layers.2.attention.out_proj.bias"),
                model_get("layers.2.ff.0.weight"), model_get("layers.2.ff.0.bias"),
                model_get("layers.2.ff.3.weight"), model_get("layers.2.ff.3.bias"),
                model_get("layers.2.norm1.weight"), model_get("layers.2.norm1.bias"),
                model_get("layers.2.norm2.weight"), model_get("layers.2.norm2.bias"), 8)
            let x = transformer_block_train(x, qkv_3,
                model_get("layers.3.attention.in_proj_bias"),
                model_get("layers.3.attention.out_proj.weight"),
                model_get("layers.3.attention.out_proj.bias"),
                model_get("layers.3.ff.0.weight"), model_get("layers.3.ff.0.bias"),
                model_get("layers.3.ff.3.weight"), model_get("layers.3.ff.3.bias"),
                model_get("layers.3.norm1.weight"), model_get("layers.3.norm1.bias"),
                model_get("layers.3.norm2.weight"), model_get("layers.3.norm2.bias"), 8)
            let x = transformer_block_train(x, qkv_4,
                model_get("layers.4.attention.in_proj_bias"),
                model_get("layers.4.attention.out_proj.weight"),
                model_get("layers.4.attention.out_proj.bias"),
                model_get("layers.4.ff.0.weight"), model_get("layers.4.ff.0.bias"),
                model_get("layers.4.ff.3.weight"), model_get("layers.4.ff.3.bias"),
                model_get("layers.4.norm1.weight"), model_get("layers.4.norm1.bias"),
                model_get("layers.4.norm2.weight"), model_get("layers.4.norm2.bias"), 8)
            let x = transformer_block_train(x, qkv_5,
                model_get("layers.5.attention.in_proj_bias"),
                model_get("layers.5.attention.out_proj.weight"),
                model_get("layers.5.attention.out_proj.bias"),
                model_get("layers.5.ff.0.weight"), model_get("layers.5.ff.0.bias"),
                model_get("layers.5.ff.3.weight"), model_get("layers.5.ff.3.bias"),
                model_get("layers.5.norm1.weight"), model_get("layers.5.norm1.bias"),
                model_get("layers.5.norm2.weight"), model_get("layers.5.norm2.bias"), 8)
            let x = transformer_block_train(x, qkv_6,
                model_get("layers.6.attention.in_proj_bias"),
                model_get("layers.6.attention.out_proj.weight"),
                model_get("layers.6.attention.out_proj.bias"),
                model_get("layers.6.ff.0.weight"), model_get("layers.6.ff.0.bias"),
                model_get("layers.6.ff.3.weight"), model_get("layers.6.ff.3.bias"),
                model_get("layers.6.norm1.weight"), model_get("layers.6.norm1.bias"),
                model_get("layers.6.norm2.weight"), model_get("layers.6.norm2.bias"), 8)
            let x = transformer_block_train(x, qkv_7,
                model_get("layers.7.attention.in_proj_bias"),
                model_get("layers.7.attention.out_proj.weight"),
                model_get("layers.7.attention.out_proj.bias"),
                model_get("layers.7.ff.0.weight"), model_get("layers.7.ff.0.bias"),
                model_get("layers.7.ff.3.weight"), model_get("layers.7.ff.3.bias"),
                model_get("layers.7.norm1.weight"), model_get("layers.7.norm1.bias"),
                model_get("layers.7.norm2.weight"), model_get("layers.7.norm2.bias"), 8)

            let logits = t_matmul(x, transpose(embed_w))
            end_tape()

            backward(logits)

            let g0 = get_grad("qkv_0")
            model_set("layers.0.attention.in_proj_weight", sgd_step(qkv_0, g0, lr))
            let g1 = get_grad("qkv_1")
            model_set("layers.1.attention.in_proj_weight", sgd_step(qkv_1, g1, lr))
            let g2 = get_grad("qkv_2")
            model_set("layers.2.attention.in_proj_weight", sgd_step(qkv_2, g2, lr))
            let g3 = get_grad("qkv_3")
            model_set("layers.3.attention.in_proj_weight", sgd_step(qkv_3, g3, lr))
            let g4 = get_grad("qkv_4")
            model_set("layers.4.attention.in_proj_weight", sgd_step(qkv_4, g4, lr))
            let g5 = get_grad("qkv_5")
            model_set("layers.5.attention.in_proj_weight", sgd_step(qkv_5, g5, lr))
            let g6 = get_grad("qkv_6")
            model_set("layers.6.attention.in_proj_weight", sgd_step(qkv_6, g6, lr))
            let g7 = get_grad("qkv_7")
            model_set("layers.7.attention.in_proj_weight", sgd_step(qkv_7, g7, lr))

            let step_count = step_count + 1
            println("  [" + step_count + "] T:\"" + teacher_line + "\" -> NEW:\"" + reply_text + "\" R:" + reward)
        }

        let avg = epoch_reward / 10
        println("  Epoch " + (epoch + 1) + " avg reward: " + avg)

        // Save after every epoch
        let saved = save_safetensors("examples/silkie_full.safetensors")
        println("  " + saved)
    }

    println("")
    println("============================================================")
    println("  PHASE 1 COMPLETE")
    println("  Total steps: " + step_count)
    println("  Avg reward: " + (total_reward / step_count))
    println("============================================================")
}
