#!/usr/bin/env python3
"""
Unified Generator with TinyGSM/GSM8K Support
完全按照成功版本的逻辑 - 使用 simple_math_problem 函数名
包含实时准确率显示功能
"""

import re
import torch
import ast
from pathlib import Path
from tqdm import tqdm
from typing import List, Dict, Optional, Union

# Default prompts
DEFAULT_CHAT_PROMPT = """[USER] Hello, how are you today? [SYSTEM] I am"""


class Generator:
    """
    统一生成器，支持：
    - chat: 标准对话生成
    - tinygsm dataset: 自动在 GSM8K test 上评估（执行生成的 Python 代码）
    """
    
    def __init__(self, config, model, tokenizer, device="cuda",output_dir=None):
        self.config = config.generation
        self.model = model.eval().to(device)
        self.tokenizer = tokenizer
        self.device = device
        
        # 将 output_dir 保存为实例属性，供保存结果使用
        self.output_dir = output_dir or Path('output')
        
        # 检测 dataset 类型
        self.dataset_name = getattr(config, 'dataset', None)
        if hasattr(self.dataset_name, 'dataset_name'):
            self.dataset_name = self.dataset_name.dataset_name
        
        self.is_tinygsm = self.dataset_name == "tinygsm"
        
        print(f"[Generator] Initialized")
        print(f"  Dataset: {self.dataset_name}")
        print(f"  Device: {device}")
        if self.is_tinygsm:
            print(f"  Mode: TinyGSM - Will evaluate on GSM8K test set with CODE EXECUTION")

    def generate(self, prompt: Optional[Union[str, List[str]]] = None):
        """
        主生成入口
        
        - 如果是 tinygsm dataset: 自动加载 GSM8K test 并评估
        - 否则: 标准 chat 生成
        """
        if self.is_tinygsm:
            return self._generate_tinygsm()
        else:
            return self._generate_chat(prompt)
    
    def _generate_chat(self, prompt: Optional[str] = None):
        """标准 chat 模式生成"""
        
        if prompt is None:
            prompt_text = getattr(self.config, "prompt", DEFAULT_CHAT_PROMPT)
        else:
            prompt_text = prompt
            
        max_len = getattr(self.config, "max_new_tokens", 512)
        temperature = getattr(self.config, "temperature", 0.7)
        top_p = getattr(self.config, "top_p", 0.9)
        return_gen_only = getattr(self.config, "return_generation_only", True)

        input_ids = torch.tensor(
            [self.tokenizer.encode(prompt_text)],
            dtype=torch.long,
            device=self.device
        )

        with torch.no_grad():
            outputs = self.model.generate(
                input_ids,
                max_generation_length=max_len,
                tokenizer=self.tokenizer,
                temperature=temperature,
                top_p=top_p,
                return_generation_only=return_gen_only
            )

        generated_texts = []
        for seq in outputs:
            text = self.tokenizer.decode(seq.tolist())
            generated_texts.append(text)

        if getattr(self.config, "save_to_file", False):
            self._save_generations(generated_texts, "generations.txt")

        for text in generated_texts:
            print(f"{prompt_text} [{text}]")
        
        return generated_texts
    
    def _generate_tinygsm(self):
        """TinyGSM mode: 在 GSM8K test set 上生成并验证"""
        print("\n" + "="*60)
        print("TinyGSM Dataset Detected")
        print("Evaluating on GSM8K Test Set")
        print("="*60 + "\n")
        
        print("Step 1: Loading GSM8K Test Set...")
        gsm8k_data = self._load_gsm8k_test()
        
        if not gsm8k_data or not gsm8k_data['questions']:
            print("❌ Failed to load GSM8K test set")
            return None
        
        questions = gsm8k_data['questions']
        ground_truth = gsm8k_data['answers']
        
        max_samples = getattr(self.config, 'max_samples', None)
        # max_samples = 100 # 如果需要限制数量，取消注释
        if max_samples is not None and len(questions) > max_samples:
            print(f"Limiting to {max_samples} samples")
            questions = questions[:max_samples]
            ground_truth = ground_truth[:max_samples]
        
        print(f"✓ Loaded {len(questions)} questions\n")
        
        print("Step 2: Generating Solutions...")
        # 【修改点 1/2】: 传入 ground_truth 到 _generate_solutions
        results = self._generate_solutions(questions, ground_truth)
        
        print("\nStep 3: Verifying Answers...")
        # _verify_answers 现在主要用于汇总和格式化结果
        verification = self._verify_answers(results, ground_truth)
        results['verification'] = verification
        results['ground_truth'] = ground_truth
        
        if getattr(self.config, 'save_to_file', False):
            self._save_tinygsm_results(results)
        
        self._print_verification_summary(verification)
        
        return results

    def _load_gsm8k_test(self) -> Dict[str, List]:
        """加载 GSM8K test set"""
        try:
            from datasets import load_dataset
            
            print("Loading from HuggingFace...")
            ds = load_dataset("openai/gsm8k", "main", split="test")
            
            questions = []
            answers = []
            
            with tqdm(total=len(ds), desc="Loading", unit="sample") as pbar:
                for sample in ds:
                    q = (sample.get("question", "") or "").strip()
                    a = (sample.get("answer", "") or "").strip()
                    
                    if not q or not a:
                        pbar.update(1)
                        continue
                    
                    answer_num = self._extract_gsm8k_answer(a)
                    
                    questions.append(q)
                    answers.append(answer_num)
                    
                    pbar.update(1)
            
            return {'questions': questions, 'answers': answers}
            
        except Exception as e:
            print(f"❌ Error loading GSM8K: {e}")
            return {'questions': [], 'answers': []}

    def _extract_gsm8k_answer(self, answer_text: str) -> str:
        """从 GSM8K 答案中提取数值"""
        match = re.search(r'####\s*([+-]?[\d,]+\.?\d*)', answer_text)
        if match:
            return match.group(1).replace(',', '').strip()
        
        match = re.search(r'result\s*=\s*([+-]?[\d,]+\.?\d*)', answer_text, re.IGNORECASE)
        if match:
            return match.group(1).replace(',', '').strip()
        
        numbers = re.findall(r'([+-]?[\d,]+\.?\d*)', answer_text)
        if numbers:
            return numbers[-1].replace(',', '').strip()
        
        return "0"

    def _create_math_prompt(self, question: str) -> str:
        """创建结构化的数学问题提示"""
        return f"""<|bos|>Question: {question}

Solution:
def simple_math_problem() -> int:
    # {question}
    """

    def _generate_solutions(self, questions: List[str], ground_truth: List[str]) -> Dict:
        """
        为每个问题生成代码解决方案
        【修改点 2/2】：在循环中加入实时准确率计算和进度条更新
        """
        max_len = 512
        temperature = getattr(self.config, 'temperature', 0.7)
        top_p = getattr(self.config, 'top_p', 0.9)
        
        results = {
            'questions': [],
            'prompts': [],
            'generated_solutions': [],
            'extracted_code': [],
            'executed_answers': [],
            'full_outputs': []
        }
        
        correct_count = 0
        total_processed = 0
        
        # 将问题和标准答案打包，并在 tqdm 中遍历
        pbar = tqdm(zip(questions, ground_truth), total=len(questions), desc="Generating", unit="q")
        
        for idx, (question, truth_val) in enumerate(pbar):
            prompt = self._create_math_prompt(question)
            
            input_ids = torch.tensor(
                [self.tokenizer.encode(prompt)],
                dtype=torch.long,
                device=self.device
            )

            with torch.no_grad():
                outputs = self.model.generate(
                    input_ids,
                    max_generation_length=max_len,
                    tokenizer=self.tokenizer,
                    temperature=temperature,
                    top_p=top_p,
                    return_generation_only=False
                )

            full_text = self.tokenizer.decode(outputs[0].tolist())
            
            # 调试打印 (减少频率)
            if idx == 0:
                print(f"\n{'='*60}")
                print(f"🔍 DEBUG Sample {idx + 1}: Raw generated text (first 500 chars):")
                print(f"{'='*60}")
                print(full_text[:500])
                print(f"{'='*60}\n")
            
            if full_text.startswith(prompt):
                solution = full_text[len(prompt):]
            else:
                solution = full_text
            
            solution = solution.replace("<|eos|>", "").replace("<|bos|>", "").strip()
            
            code = self._extract_function_code(full_text)
            
            if code:
                answer = self._execute_math_code(code)
            else:
                answer = None
                
            # --- 实时验证和进度条更新 ---
            is_correct = False
            if truth_val is not None:
                is_correct = self._compare_answers(answer, truth_val)
                if is_correct:
                    correct_count += 1
                total_processed += 1
                
                # 更新进度条后缀显示准确率
                current_acc = correct_count / total_processed
                pbar.set_postfix({"Acc": f"{current_acc:.2%}"})
            # ------------------------------
            
            results['questions'].append(question)
            results['prompts'].append(prompt)
            results['generated_solutions'].append(solution)
            results['extracted_code'].append(code if code else "No code extracted")
            results['executed_answers'].append(answer)
            results['full_outputs'].append({
                'question': question,
                'solution': solution,
                'code': code,
                'answer': answer
            })
            
            if idx < 3 or idx == len(questions) - 1:
                print(f"\n{'─'*60}")
                print(f"Sample {idx + 1}")
                print(f"{'─'*60}")
                print(f"Q: {question[:100]}...")
                if code:
                    print(f"Code:\n{code[:200]}...")
                print(f"Answer: {answer if answer is not None else '[None]'}")
                if truth_val:
                    status = "✓" if is_correct else "✗"
                    print(f"Truth:  {truth_val} {status}")
                print(f"{'─'*60}")
        
        return results

    def _extract_function_code(self, generated_text):
        """Extract and clean function code from generated text"""
        func_start = generated_text.find("def simple_math_problem")
        if func_start == -1:
            return None
            
        func_text = generated_text[func_start:]
        
        lines = func_text.split('\n')
        cleaned_lines = []
        
        for i, line in enumerate(lines):
            if i == 0:
                cleaned_lines.append("def simple_math_problem() -> int:")
            else:
                if not cleaned_lines and not line.strip():
                    continue
                    
                # Stop at next function or end markers
                stripped_line = line.strip()
                if stripped_line and not line.startswith(' ') and not line.startswith('\t') and i > 0:
                    if 'def ' in line or stripped_line == 'result = simple_math_problem()':
                        break
                
                # Stop at markdown code block markers or answer markers
                if stripped_line in ['```', '```python', '```py'] or stripped_line.startswith('Answer:') or '<|eos|>' in stripped_line:
                    break
                
                clean_line = line.rstrip()
                original_line = clean_line
                
                if clean_line.strip():
                    # Check for end markers (even if indented)
                    stripped = clean_line.strip()
                    if stripped in ['```', '```python', '```py'] or stripped.startswith('Answer:') or '<|eos|>' in stripped:
                        break
                    
                    clean_line = clean_line.replace(''', "'").replace(''', "'")
                    clean_line = clean_line.replace('"', '"').replace('"', '"')
                    
                    stripped = clean_line.strip()
                    
                    # 检查是否要跳过
                    if stripped.startswith('"""') or stripped.startswith("'''") or stripped.startswith('#'):
                        continue
                    
                    if clean_line.strip() and not clean_line.startswith('    '):
                        clean_line = '    ' + clean_line.strip()
                    
                    cleaned_lines.append(clean_line)
                else:
                    cleaned_lines.append('')
        
        while cleaned_lines and not cleaned_lines[-1].strip():
            cleaned_lines.pop()
            
        has_return = any('return' in line for line in cleaned_lines)
        if not has_return:
            cleaned_lines.append('    return result')
        
        final_code = '\n'.join(cleaned_lines)
        return final_code

    def _execute_math_code(self, code):
        """Safely execute mathematical code"""
        try:
            if not code:
                return None
            
            if 'return' not in code:
                code += '\n    return result'
            
            full_code = code + '\n\nresult = simple_math_problem()'
            
            try:
                ast.parse(full_code)
            except SyntaxError as e:
                return None
            
            safe_globals = {
                "__builtins__": {
                    "abs": abs, "round": round, "min": min, "max": max, "sum": sum,
                    "int": int, "float": float, "str": str, "len": len, "range": range,
                    "list": list, "dict": dict, "set": set, "tuple": tuple,
                    "True": True, "False": False, "None": None
                }
            }
            safe_locals = {}
            
            exec(full_code, safe_globals, safe_locals)
            
            if 'result' in safe_locals:
                result_value = float(safe_locals['result'])
                return result_value
            else:
                return None
            
        except Exception as e:
            return None

    def _verify_answers(self, results: Dict, ground_truth: List[str]) -> Dict:
        """验证答案"""
        predictions = results['executed_answers']
        
        verification = {
            'total': len(ground_truth),
            'correct': 0,
            'incorrect': 0,
            'no_answer': 0,
            'accuracy': 0.0,
            'details': []
        }
        
        for idx, (pred, truth) in enumerate(zip(predictions, ground_truth)):
            if pred is None:
                verification['no_answer'] += 1
                is_correct = False
            else:
                is_correct = self._compare_answers(pred, truth)
                
                if is_correct:
                    verification['correct'] += 1
                else:
                    verification['incorrect'] += 1
            
            verification['details'].append({
                'index': idx,
                'question': results['questions'][idx],
                'predicted': pred,
                'ground_truth': truth,
                'correct': is_correct
            })
        
        verification['accuracy'] = (
            verification['correct'] / verification['total']
            if verification['total'] > 0 else 0.0
        )
        
        return verification

    def _compare_answers(self, pred: str, truth: str) -> bool:
        """数值比较"""
        pred_str = str(pred).strip().replace(',', '')
        truth_str = str(truth).strip().replace(',', '')
        
        if pred_str == truth_str:
            return True
        
        try:
            pred_num = float(pred_str)
            truth_num = float(truth_str)
            return abs(pred_num - truth_num) < 1e-6
        except (ValueError, TypeError):
            pass
        
        return False

    def _print_verification_summary(self, verification: Dict):
        """打印验证总结"""
        print(f"\n{'='*60}")
        print("GSM8K Evaluation Results")
        print(f"{'='*60}")
        print(f"Total:      {verification['total']}")
        print(f"Correct:    {verification['correct']} ({verification['accuracy']*100:.2f}%)")
        print(f"Incorrect:  {verification['incorrect']}")
        print(f"No Answer:  {verification['no_answer']}")
        print(f"{'='*60}")
        
        incorrect = [d for d in verification['details'] if not d['correct']][:3]
        if incorrect:
            print("\nSample Incorrect Predictions:")
            for detail in incorrect:
                print(f"  Q: {detail['question'][:60]}...")
                print(f"  Predicted: {detail['predicted']}, Truth: {detail['ground_truth']}\n")

    def _save_generations(self, texts: List[str], filename: str):
        """保存 chat 生成结果"""
        output_dir = Path(getattr(self.config, 'output_dir', self.output_dir))
        output_dir.mkdir(parents=True, exist_ok=True)
        output_path = output_dir / filename
        
        with open(output_path, 'w', encoding='utf-8') as f:
            for i, text in enumerate(texts, 1):
                f.write(f"{'='*60}\n")
                f.write(f"Generation {i}\n")
                f.write(f"{'='*60}\n")
                f.write(text + "\n\n")
        
        print(f"[Generator] Saved to {output_path}")

    def _save_tinygsm_results(self, results: Dict):
        """保存 TinyGSM/GSM8K 评估结果"""
        # 使用实例属性 self.output_dir
        output_dir = Path(self.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        
        full_path = output_dir / "gsm8k_results.txt"
        with open(full_path, 'w', encoding='utf-8') as f:
            for i, item in enumerate(results['full_outputs'], 1):
                f.write(f"{'='*60}\n")
                f.write(f"Problem {i}\n")
                f.write(f"{'='*60}\n")
                f.write(f"Question:\n{item['question']}\n\n")
                f.write(f"Solution:\n{item['solution']}\n\n")
                
                if item.get('code'):
                    f.write(f"Extracted Code:\n{item['code']}\n\n")
                
                f.write(f"Answer: {item['answer']}\n")
                
                if 'ground_truth' in results and i <= len(results['ground_truth']):
                    truth = results['ground_truth'][i-1]
                    # 由于我们在 _generate_solutions 中已经比较了，这里的比较用于最终文件的标记
                    match = "✓" if self._compare_answers(item['answer'], truth) else "✗"
                    f.write(f"Truth: {truth} {match}\n")
                
                f.write(f"{'='*60}\n\n")
        
        print(f"[Generator] Saved results to {full_path}")
        
        if 'verification' in results:
            verify_path = output_dir / "verification.txt"
            with open(verify_path, 'w', encoding='utf-8') as f:
                v = results['verification']
                f.write(f"{'='*60}\n")
                f.write("GSM8K Verification Results\n")
                f.write(f"{'='*60}\n")
                f.write(f"Total: {v['total']}\n")
                f.write(f"Correct: {v['correct']}\n")
                f.write(f"Incorrect: {v['incorrect']}\n")
                f.write(f"No Answer: {v['no_answer']}\n")
                f.write(f"Accuracy: {v['accuracy']*100:.2f}%\n")
                f.write(f"{'='*60}\n\n")
                
                for detail in v['details']:
                    status = "✓" if detail['correct'] else "✗"
                    f.write(f"{detail['index']+1}. {status} ")
                    f.write(f"Pred: {detail['predicted']}, Truth: {detail['ground_truth']}\n")
            
            print(f"[Generator] Saved verification to {verify_path}")