// 🌌 [BioPhys 5.0 타키온 투기적 디코딩 엔진] (src/tachyon_speculative.rs)
// 소형 1AU 드래프트 모델이 미래 K개 토큰을 동시 예측하고, 16GB 주 모델이 1회 배치로 동시 검증하여 속도 3배 가속
use crate::token_autoregressive_snn::SpikingTokenEngine;
use std::time::Instant;

pub struct TachyonSpeculativeEngine {
    pub draft_engine: SpikingTokenEngine, // 초고속 소형 드래프트 모델 (1AU)
    pub lookahead_k: usize,               // 동시 예측 미래 토큰 수 (K = 4)
    pub total_speculated: usize,
    pub total_accepted: usize,
}

impl TachyonSpeculativeEngine {
    pub fn new(draft_hidden_dim: usize, lookahead_k: usize) -> Self {
        TachyonSpeculativeEngine {
            draft_engine: SpikingTokenEngine::new(draft_hidden_dim),
            lookahead_k,
            total_speculated: 0,
            total_accepted: 0,
        }
    }

    /// [타키온 투기적 디코딩: K개 미래 토큰 동시 생성 및 단일 배치 검증]
    pub fn speculate_and_verify(
        &mut self,
        current_prompt: &[&str],
        target_engine: &mut SpikingTokenEngine,
        lattice_energy: f32,
    ) -> Vec<String> {
        let mut accepted_tokens = Vec::new();

        // 1. 드래프트 모델이 초고속으로 미래 K개 후보 토큰 예측 (Tachyon Draft)
        let mut draft_candidates = Vec::new();
        let mut temp_prompt = current_prompt.to_vec();

        for _ in 0..self.lookahead_k {
            let last_token_id = temp_prompt.last()
                .and_then(|&w| self.draft_engine.vocab_to_id.get(w).copied())
                .unwrap_or(0);
            
            let logits = self.draft_engine.step_token(last_token_id, lattice_energy);
            let mut best_id = 3;
            let mut best_l = f32::NEG_INFINITY;
            for (id, &l) in logits.iter().enumerate() {
                if id > 2 && l > best_l {
                    best_l = l;
                    best_id = id;
                }
            }
            let word = self.draft_engine.vocab[best_id].clone();
            draft_candidates.push((best_id, word));
            self.total_speculated += 1;
        }

        // 2. 16GB 타겟 모델이 1회 배치로 후보 토큰들을 동시 검증 (Target Verification)
        for (cand_id, cand_word) in draft_candidates {
            let target_logits = target_engine.step_token(cand_id, lattice_energy);
            
            // 타겟 모델의 Top-10 후보에 드래프트 토큰이 포함되면 통과 (Accept)
            let mut is_accepted = false;
            let mut top_indices: Vec<usize> = (0..target_logits.len()).collect();
            top_indices.sort_by(|&a, &b| target_logits[b].partial_cmp(&target_logits[a]).unwrap());

            if top_indices[0..5.min(top_indices.len())].contains(&cand_id) {
                is_accepted = true;
            }

            if is_accepted {
                accepted_tokens.push(cand_word);
                self.total_accepted += 1;
            } else {
                // 불일치 시 투기적 분기 중단 (Fallback)
                break;
            }
        }

        if accepted_tokens.is_empty() && !self.draft_engine.vocab.is_empty() {
            accepted_tokens.push(self.draft_engine.vocab[3].clone());
        }

        accepted_tokens
    }

    /// 투기적 수용률 (Acceptance Rate)
    pub fn acceptance_rate(&self) -> f32 {
        if self.total_speculated == 0 { 1.0 } else { (self.total_accepted as f32) / (self.total_speculated as f32) }
    }
}