from typing import List, Tuple, Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import namedtuple
import numpy as np

# 定義狀態數據結構
SpeculativeState = namedtuple('SpeculativeState', [
    'generated_tokens',           # 已生成token序列
    'draft_tokens',               # 草稿token序列
    'draft_probabilities',        # 草稿概率分佈
    'target_probabilities',       # 目標模型概率分佈
    'acceptance_flags',           # 接受標記
    'position',                   # 當前位置
    'remaining_budget'            # 剩餘生成長度
])

class MDPState:
    def __init__(self, generated_tokens: List[int], draft_tokens: List[int],
                 draft_probs: torch.Tensor, target_probs: torch.Tensor,
                 current_pos: int, max_length: int):
        self.generated_tokens = generated_tokens
        self.draft_tokens = draft_tokens
        self.draft_probabilities = draft_probs
        self.target_probabilities = target_probs
        self.current_position = current_pos
        self.max_length = max_length
        self.remaining_budget = max_length - current_pos
        
    def get_acceptance_rates(self) -> torch.Tensor:
        """計算每個候選token的接受概率"""
        # 基於目標模型概率和草稿模型概率計算接受率
        acceptance_probs = torch.min(
            torch.ones_like(self.target_probabilities),
            self.target_probabilities / (self.draft_probabilities + 1e-8)
        )
        return acceptance_probs
    
    def is_terminal(self) -> bool:
        """判斷是否為終止狀態"""
        return (self.current_position >= self.max_length or 
                len(self.generated_tokens) > 0 and self.generated_tokens[-1] == 2)  # EOS token
    
    def get_valid_actions(self) -> List[int]:
        """獲取有效的動作空間"""
        if self.is_terminal():
            return []
        # 動作空間：接受所有、部分接受或拒絕
        return list(range(len(self.draft_tokens) + 1))

class SpeculativePolicy:
    def __init__(self, gamma: float = 0.99, epsilon: float = 0.1):
        self.gamma = gamma  # 折扣因子
        self.epsilon = epsilon  # 探索率
        self.value_network = ValueNetwork()
        self.policy_network = PolicyNetwork()
        
    def select_action(self, state: MDPState) -> int:
        """基於當前狀態選擇動作"""
        if np.random.random() < self.epsilon:
            # 探索：隨機選擇動作
            valid_actions = state.get_valid_actions()
            return np.random.choice(valid_actions) if valid_actions else 0
        else:
            # 利用：選擇價值最大的動作
            return self._greedy_action(state)
    
    def _greedy_action(self, state: MDPState) -> int:
        """貪心策略選擇動作"""
        valid_actions = state.get_valid_actions()
        if not valid_actions:
            return 0
            
        action_values = []
        for action in valid_actions:
            value = self._evaluate_action(state, action)
            action_values.append(value)
        
        return valid_actions[np.argmax(action_values)]
    
    def _evaluate_action(self, state: MDPState, action: int) -> float:
        """評估動作的長期價值"""
        # 即時獎勵
        immediate_reward = self._calculate_reward(state, action)
        
        # 預測下一狀態價值
        next_state = self._predict_next_state(state, action)
        if next_state.is_terminal():
            future_value = 0.0
        else:
            future_value = self.value_network(next_state)
            
        return immediate_reward + self.gamma * future_value
    
    def _calculate_reward(self, state: MDPState, action: int) -> float:
        """計算即時獎勵"""
        if action == 0:  # 拒絕所有
            return -1.0  # 懲罰完全拒絕
        
        acceptance_probs = state.get_acceptance_rates()
        accepted_tokens = min(action, len(acceptance_probs))
        
        # 獎勵與接受的token數量和概率成正比
        reward = accepted_tokens * torch.mean(acceptance_probs[:accepted_tokens]).item()
        
        # 懲罰過度冒險
        if action > len(state.draft_tokens):
            reward -= 0.5
            
        return reward

class ValueNetwork(nn.Module):
    """狀態價值網絡"""
    def __init__(self, hidden_size: int = 128):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(512, hidden_size),  # 假設狀態特徵維度為512
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 1)
        )
    
    def forward(self, state: MDPState) -> torch.Tensor:
        # 將狀態轉換為特徵向量
        state_features = self._extract_features(state)
        return self.network(state_features)
    
    def _extract_features(self, state: MDPState) -> torch.Tensor:
        """提取狀態特徵"""
        features = []
        
        # 接受率特徵
        acceptance_rates = state.get_acceptance_rates()
        features.extend([
            acceptance_rates.mean().item(),
            acceptance_rates.std().item(),
            acceptance_rates.max().item()
        ])
        
        # 位置特徵
        features.extend([
            state.current_position / state.max_length,
            state.remaining_budget / state.max_length
        ])
        
        # 概率分佈特徵
        target_entropy = -torch.sum(
            state.target_probabilities * torch.log(state.target_probabilities + 1e-8)
        ).item()
        draft_entropy = -torch.sum(
            state.draft_probabilities * torch.log(state.draft_probabilities + 1e-8)
        ).item()
        
        features.extend([target_entropy, draft_entropy])
        
        return torch.tensor(features, dtype=torch.float32).unsqueeze(0)

class PolicyNetwork(nn.Module):
    """策略網絡"""
    def __init__(self, action_dim: int = 10, hidden_size: int = 128):
        super().__init__()
        self.action_dim = action_dim
        self.network = nn.Sequential(
            nn.Linear(512, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, action_dim)
        )
    
    def forward(self, state: MDPState) -> torch.Tensor:
        state_features = self._extract_features(state)
        action_logits = self.network(state_features)
        return F.softmax(action_logits, dim=-1)
    
    def _extract_features(self, state: MDPState) -> torch.Tensor:
        # 簡化特徵提取，實際應用需要更復雜的特徵工程
        return ValueNetwork._extract_features(state, state)

class SpeculativeDecoder:
    def __init__(self, target_model: nn.Module, draft_model: nn.Module,
                 max_draft_tokens: int = 5, policy: Optional[SpeculativePolicy] = None):
        self.target_model = target_model
        self.draft_model = draft_model
        self.max_draft_tokens = max_draft_tokens
        self.policy = policy or SpeculativePolicy()
        
        # 性能統計
        self.stats = {
            'total_tokens': 0,
            'accepted_tokens': 0,
            'target_calls': 0,
            'draft_calls': 0
        }
    
    def generate(self, input_ids: torch.Tensor, max_length: int,
                temperature: float = 1.0) -> List[int]:
        """使用投機解碼生成序列"""
        generated_tokens = input_ids.tolist()
        current_position = len(generated_tokens)
        
        while current_position < max_length and not self._is_eos(generated_tokens):
            # 草稿階段：生成候選token序列
            draft_tokens, draft_probs = self._draft_stage(
                generated_tokens, current_position, temperature
            )
            
            # 驗證階段：目標模型驗證候選token
            target_probs = self._verification_stage(
                generated_tokens, draft_tokens, current_position, temperature
            )
            
            # MDP決策：決定接受多少個候選token
            state = MDPState(
                generated_tokens=generated_tokens,
                draft_tokens=draft_tokens,
                draft_probs=draft_probs,
                target_probs=target_probs,
                current_pos=current_position,
                max_length=max_length
            )
            
            accept_count = self.policy.select_action(state)
            
            # 執行動作，更新狀態
            new_tokens = self._execute_decision(
                draft_tokens, target_probs, accept_count
            )
            
            # 更新生成序列
            generated_tokens.extend(new_tokens)
            current_position += len(new_tokens)
            
            # 更新統計信息
            self._update_stats(len(new_tokens), accept_count, len(draft_tokens))
            
            # 如果拒絕所有候選，回退到傳統解碼
            if accept_count == 0:
                next_token = self._traditional_step(generated_tokens, temperature)
                generated_tokens.append(next_token)
                current_position += 1
        
        return generated_tokens
    
    def _draft_stage(self, generated_tokens: List[int], current_pos: int,
                    temperature: float) -> Tuple[List[int], torch.Tensor]:
        """草稿模型生成候選序列"""
        self.stats['draft_calls'] += 1
        
        draft_tokens = []
        draft_probs = []
        
        # 使用草稿模型並行生成多個候選token
        draft_input = torch.tensor(generated_tokens, dtype=torch.long).unsqueeze(0)
        
        for i in range(self.max_draft_tokens):
            with torch.no_grad():
                draft_output = self.draft_model(draft_input)
                next_token_logits = draft_output[0, -1, :] / temperature
                next_token_probs = F.softmax(next_token_logits, dim=-1)
                
                # 採樣下一個token
                next_token = torch.multinomial(next_token_probs, 1).item()
                draft_tokens.append(next_token)
                draft_probs.append(next_token_probs)
                
                # 更新輸入
                draft_input = torch.cat([
                    draft_input, 
                    torch.tensor([[next_token]], dtype=torch.long)
                ], dim=1)
                
                # 如果生成EOS token，提前終止
                if next_token == 2:  # EOS
                    break
        
        draft_probs_tensor = torch.stack(draft_probs)
        return draft_tokens, draft_probs_tensor
    
    def _verification_stage(self, generated_tokens: List[int], 
                          draft_tokens: List[int], current_pos: int,
                          temperature: float) -> torch.Tensor:
        """目標模型驗證候選序列"""
        self.stats['target_calls'] += 1
        
        # 構建包含候選序列的完整輸入
        verification_input = generated_tokens + draft_tokens
        input_tensor = torch.tensor(verification_input, dtype=torch.long).unsqueeze(0)
        
        with torch.no_grad():
            target_output = self.target_model(input_tensor)
            target_logits = target_output[0, len(generated_tokens):, :] / temperature
            target_probs = F.softmax(target_logits, dim=-1)
        
        return target_probs
    
    def _execute_decision(self, draft_tokens: List[int], 
                         target_probs: torch.Tensor, 
                         accept_count: int) -> List[int]:
        """執行接受決策"""
        accepted_tokens = []
        
        for i in range(accept_count):
            if i < len(draft_tokens):
                # 計算接受概率
                acceptance_prob = torch.min(
                    torch.tensor(1.0),
                    target_probs[i, draft_tokens[i]] / (target_probs[i, draft_tokens[i]] + 1e-8)
                )
                
                # 根據接受概率決定是否接受該token
                if torch.rand(1) < acceptance_prob:
                    accepted_tokens.append(draft_tokens[i])
                else:
                    # 拒絕當前token，從目標分佈中重新採樣
                    new_token = torch.multinomial(target_probs[i], 1).item()
                    accepted_tokens.append(new_token)
                    break
            else:
                break
        
        return accepted_tokens
    
    def _traditional_step(self, generated_tokens: List[int], 
                         temperature: float) -> int:
        """傳統自迴歸解碼單步"""
        input_tensor = torch.tensor(generated_tokens, dtype=torch.long).unsqueeze(0)
        
        with torch.no_grad():
            output = self.target_model(input_tensor)
            next_token_logits = output[0, -1, :] / temperature
            next_token_probs = F.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(next_token_probs, 1).item()
        
        return next_token
    
    def _is_eos(self, tokens: List[int]) -> bool:
        """檢查是否生成EOS token"""
        return len(tokens) > 0 and tokens[-1] == 2
    
    def _update_stats(self, new_tokens_count: int, accept_count: int, 
                     draft_length: int):
        """更新性能統計"""
        self.stats['total_tokens'] += new_tokens_count
        self.stats['accepted_tokens'] += accept_count
    
    def get_efficiency_metrics(self) -> Dict[str, float]:
        """獲取效率指標"""
        acceptance_rate = (self.stats['accepted_tokens'] / 
                         self.stats['total_tokens'] if self.stats['total_tokens'] > 0 else 0)
        
        speedup_factor = (self.stats['total_tokens'] / 
                         self.stats['target_calls'] if self.stats['target_calls'] > 0 else 1)
        
        return {
            'acceptance_rate': acceptance_rate,
            'speedup_factor': speedup_factor,
            'target_calls': self.stats['target_calls'],
            'draft_calls': self.stats['draft_calls'],
            'total_tokens': self.stats['total_tokens']
        }

class EnhancedSpeculativeDecoder(SpeculativeDecoder):
    def __init__(self, target_model: nn.Module, draft_model: nn.Module,
                 max_draft_tokens: int = 5, policy: Optional[SpeculativePolicy] = None,
                 adaptive_draft: bool = True, lookahead_window: int = 3):
        super().__init__(target_model, draft_model, max_draft_tokens, policy)
        self.adaptive_draft = adaptive_draft
        self.lookahead_window = lookahead_window
        self.confidence_threshold = 0.8
        
        # 自適應參數
        self.draft_length_history = []
        self.acceptance_history = []
    
    def _adaptive_draft_length(self) -> int:
        """自適應調整草稿生成長度"""
        if not self.adaptive_draft or not self.acceptance_history:
            return self.max_draft_tokens
        
        # 基於歷史接受率調整生成長度
        recent_acceptance = np.mean(self.acceptance_history[-10:]) if self.acceptance_history else 0.5
        recent_draft_length = np.mean(self.draft_length_history[-5:]) if self.draft_length_history else self.max_draft_tokens
        
        if recent_acceptance > 0.8:
            # 高接受率時增加生成長度
            adaptive_length = min(self.max_draft_tokens + 1, 10)
        elif recent_acceptance < 0.3:
            # 低接受率時減少生成長度
            adaptive_length = max(1, self.max_draft_tokens - 1)
        else:
            adaptive_length = self.max_draft_tokens
            
        return adaptive_length
    
    def _lookahead_verification(self, generated_tokens: List[int],
                              draft_tokens: List[int], current_pos: int,
                              temperature: float) -> torch.Tensor:
        """前瞻性驗證，考慮後續token的依賴關係"""
        target_probs = super()._verification_stage(
            generated_tokens, draft_tokens, current_pos, temperature
        )
        
        if self.lookahead_window > 0 and len(draft_tokens) > 1:
            # 對每個位置，考慮後續窗口內的概率分佈
            enhanced_probs = []
            for i in range(len(draft_tokens)):
                lookahead_end = min(i + self.lookahead_window, len(draft_tokens))
                
                # 計算當前位置在考慮後續token時的調整概率
                adjusted_probs = self._adjust_probs_with_lookahead(
                    target_probs[i:lookahead_end], draft_tokens[i:lookahead_end]
                )
                enhanced_probs.append(adjusted_probs)
            
            target_probs = torch.stack(enhanced_probs)
        
        return target_probs
    
    def _adjust_probs_with_lookahead(self, probs_window: torch.Tensor,
                                   tokens_window: List[int]) -> torch.Tensor:
        """使用前瞻窗口調整概率分佈"""
        base_probs = probs_window[0]
        
        if len(probs_window) == 1:
            return base_probs
        
        # 考慮後續token的連貫性調整當前概率
        coherence_scores = []
        for token_idx in range(base_probs.shape[0]):
            # 計算選擇當前token時後續序列的連貫性得分
            coherence_score = self._calculate_coherence_score(
                token_idx, tokens_window, probs_window
            )
            coherence_scores.append(coherence_score)
        
        coherence_tensor = torch.tensor(coherence_scores, dtype=torch.float32)
        adjusted_probs = base_probs * coherence_tensor
        adjusted_probs = adjusted_probs / adjusted_probs.sum()
        
        return adjusted_probs
    
    def _calculate_coherence_score(self, current_token: int,
                                 tokens_window: List[int],
                                 probs_window: torch.Tensor) -> float:
        """計算連貫性得分"""
        score = 1.0
        
        # 簡化實現：檢查當前token與後續token的兼容性
        for i in range(1, len(probs_window)):
            # 基於語言模型的轉移概率估計連貫性
            transition_prob = probs_window[i, tokens_window[i]]
            score *= transition_prob.item()
            
        return score
    
    def generate(self, input_ids: torch.Tensor, max_length: int,
                temperature: float = 1.0) -> List[int]:
        """增強的生成方法"""
        # 自適應調整草稿長度
        adaptive_max_draft = self._adaptive_draft_length()
        
        generated_tokens = input_ids.tolist()
        current_position = len(generated_tokens)
        
        while current_position < max_length and not self._is_eos(generated_tokens):
            # 使用自適應草稿長度
            draft_tokens, draft_probs = self._draft_stage(
                generated_tokens, current_position, temperature
            )
            
            # 使用前瞻驗證
            target_probs = self._lookahead_verification(
                generated_tokens, draft_tokens, current_position, temperature
            )
            
            state = MDPState(
                generated_tokens=generated_tokens,
                draft_tokens=draft_tokens,
                draft_probs=draft_probs,
                target_probs=target_probs,
                current_pos=current_position,
                max_length=max_length
            )
            
            accept_count = self.policy.select_action(state)
            
            new_tokens = self._execute_decision(
                draft_tokens, target_probs, accept_count
            )
            
            generated_tokens.extend(new_tokens)
            current_position += len(new_tokens)
            
            # 更新歷史記錄用於自適應調整
            self.acceptance_history.append(
                accept_count / len(draft_tokens) if draft_tokens else 0
            )
            self.draft_length_history.append(len(draft_tokens))
            
            self._update_stats(len(new_tokens), accept_count, len(draft_tokens))
            
            if accept_count == 0:
                next_token = self._traditional_step(generated_tokens, temperature)
                generated_tokens.append(next_token)
                current_position += 1
        
        return generated_tokens

class EfficiencyBenchmark:
    def __init__(self, decoder: SpeculativeDecoder, test_dataset: List[str]):
        self.decoder = decoder
        self.test_dataset = test_dataset
        self.results = []
    
    def run_benchmark(self, num_samples: int = 100) -> Dict[str, float]:
        """運行性能基準測試"""
        import time
        from tqdm import tqdm
        
        samples = self.test_dataset[:num_samples]
        total_time = 0
        total_tokens = 0
        
        for sample in tqdm(samples, desc="Running Benchmark"):
            input_ids = self._text_to_ids(sample)
            
            start_time = time.time()
            output_tokens = self.decoder.generate(
                input_ids, max_length=len(input_ids) + 50
            )
            end_time = time.time()
            
            generation_time = end_time - start_time
            generated_tokens = len(output_tokens) - len(input_ids)
            
            total_time += generation_time
            total_tokens += generated_tokens
            
            # 記錄每次生成的結果
            metrics = self.decoder.get_efficiency_metrics()
            self.results.append({
                'time': generation_time,
                'tokens': generated_tokens,
                'speed': generated_tokens / generation_time,
                **metrics
            })
        
        # 計算總體統計
        avg_speed = total_tokens / total_time
        avg_acceptance = np.mean([r['acceptance_rate'] for r in self.results])
        avg_speedup = np.mean([r['speedup_factor'] for r in self.results])
        
        return {
            'average_speed': avg_speed,
            'average_acceptance_rate': avg_acceptance,
            'average_speedup_factor': avg_speedup,
            'total_time': total_time,
            'total_tokens': total_tokens
        }
    
    def compare_with_baseline(self, baseline_decoder: SpeculativeDecoder) -> Dict[str, float]:
        """與基線方法比較"""
        baseline_benchmark = EfficiencyBenchmark(baseline_decoder, self.test_dataset)
        baseline_results = baseline_benchmark.run_benchmark()
        our_results = self.run_benchmark()
        
        comparison = {
            'speed_improvement': our_results['average_speed'] / baseline_results['average_speed'],
            'acceptance_improvement': (our_results['average_acceptance_rate'] - 
                                     baseline_results['average_acceptance_rate']),
            'speedup_improvement': (our_results['average_speedup_factor'] - 
                                  baseline_results['average_speedup_factor']),
            'efficiency_gain': our_results['total_tokens'] / our_results['total_time'] - 
                             baseline_results['total_tokens'] / baseline_results['total_time']
        }
        
        return comparison
    
    def _text_to_ids(self, text: str) -> torch.Tensor:
        """文本轉換為token ID（簡化實現）"""
        # 實際應用中應使用對應的tokenizer
        return torch.tensor([ord(c) for c in text[:100]], dtype=torch.long)

class ProductionSpeculativeDecoder(EnhancedSpeculativeDecoder):
    def __init__(self, target_model: nn.Module, draft_model: nn.Module,
                 max_draft_tokens: int = 5, policy: Optional[SpeculativePolicy] = None,
                 batch_size: int = 1, use_quantization: bool = True):
        super().__init__(target_model, draft_model, max_draft_tokens, policy)
        
        self.batch_size = batch_size
        self.use_quantization = use_quantization
        
        # 生產環境優化
        if use_quantization:
            self.draft_model = self._quantize_model(self.draft_model)
    
    def _quantize_model(self, model: nn.Module) -> nn.Module:
        """模型量化以加速推理"""
        try:
            model.eval（)
            quantized_model = torch.quantization.quantize_dynamic(
                model, {nn.Linear}, dtype=torch.qint8
            )
            return quantized_model
        except Exception as e:
            print(f"Quantization failed: {e}, using original model")
            return model
    
    def batch_generate(self, input_batch: List[torch.Tensor], 
                      max_length: int) -> List[List[int]]:
        """批量生成以提高GPU利用率"""
        results = []
        
        for i in range(0, len(input_batch), self.batch_size):
            batch_inputs = input_batch[i:i + self.batch_size]
            batch_results = []
            
            for input_ids in batch_inputs:
                result = self.generate(input_ids, max_length)
                batch_results.append(result)
            
            results.extend(batch_results)
        
        return results
    
    def warmup(self, warmup_sequences: int = 10):
        """預熱運行以確保穩定性能"""
        dummy_input = torch.randint(0, 1000, (1, 10))
        
        for _ in range(warmup_sequences):
            _ = self.generate(dummy_input, max_length=20)

class AdaptiveMonitoringSystem:
    def __init__(self, decoder: ProductionSpeculativeDecoder):
        self.decoder = decoder
        self.performance_history = []
        self.adaptive_config = {
            'min_draft_length': 1,
            'max_draft_length': 8,
            'target_acceptance': 0.7,
            'adjustment_step': 1
        }
    
    def monitor_and_adjust(self):
        """監控性能並自適應調整參數"""
        current_metrics = self.decoder.get_efficiency_metrics()
        self.performance_history.append(current_metrics)
        
        if len(self.performance_history) < 5:
            return
        
        # 分析趨勢並調整參數
        recent_acceptance = np.mean([
            m['acceptance_rate'] for m in self.performance_history[-5:]
        ])
        
        current_draft_length = self.decoder.max_draft_tokens
        
        if recent_acceptance > self.adaptive_config['target_acceptance'] + 0.1:
            # 接受率過高，增加草稿長度以追求更高加速比
            new_length = min(
                current_draft_length + self.adaptive_config['adjustment_step'],
                self.adaptive_config['max_draft_length']
            )
            self.decoder.max_draft_tokens = new_length
        elif recent_acceptance < self.adaptive_config['target_acceptance'] - 0.1:
            # 接受率過低，減少草稿長度保證效率
            new_length = max(
                current_draft_length - self.adaptive_config['adjustment_step'],
                self.adaptive_config['min_draft_length']
            )
            self.decoder.max_draft_tokens = new_length