"""
Baseline Benchmark Framework

Evaluates LLM reasoning using four methods:
  - IO       (Input-Output, direct answer)
  - CoT      (Chain-of-Thought)
  - CoT-SC   (CoT with Self-Consistency, majority vote)
  - ToT      (Tree-of-Thought, beam search)

Each method produces accuracy scores, token usage, and timing.
Results are saved to data/results/ as JSON files.
"""

import json
import os
import re
import sys
import time
from pathlib import Path
from typing import Callable, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed

import numpy as np
from tqdm import tqdm

# Add project root to path
_project_root = Path(__file__).resolve().parent.parent.parent
if str(_project_root) not in sys.path:
    sys.path.insert(0, str(_project_root))

from src.llm_client import LLMClient, make_client

# ── Default Prompts (English) ───────────────────────────────────────

IO_SYSTEM = "You are a helpful assistant. Answer the user's question directly and concisely. Provide only the answer, with no additional explanation."

COT_SYSTEM = "You are a math reasoning assistant. Work through the problem step by step, then provide the final answer at the end in the format: \\boxed{answer}"

COT_PROMPT_TEMPLATE = "Solve the following problem step by step. At the end, put your final answer in the format \\\\boxed{{answer}}.\n\nProblem: {question}\n\nLet's think step by step:"

TOT_SYSTEM = "You are a reasoning assistant. Think systematically and evaluate multiple approaches."

TOT_GENERATE_PROMPT = "Given the current state of reasoning about:\n{context}\n\nGenerate {n_branches} distinct next steps. Each step should advance the reasoning toward solving:\n{question}\n\nNumber them 1 through {n_branches}, one per line."

TOT_EVALUATE_PROMPT = "Given the problem:\n{question}\n\nAnd the current reasoning step:\n{step}\n\nRate this step's promise on a scale of 1-10. Consider:\n- Does it make progress toward the answer?\n- Is it logically sound?\n- How likely is it to lead to the correct answer?\n\nReturn only a number 1-10."

TOT_DECIDE_PROMPT = "Based on the following reasoning traces, determine the final answer to:\n\n{question}\n\nReasoning traces considered:\n{traces}\n\nFinal answer in the format: \\\\boxed{{answer}}"

# ── Chinese Prompts (中文提示) ──────────────────────────────────────
# Used when the dataset name starts with "c_" or is "cmath"/"c_aime".

IO_SYSTEM_ZH = "你是一个乐于助人的数学助手。请直接回答用户的问题，只给出答案，不需要额外的解释。"

COT_SYSTEM_ZH = "你是一个数学推理助手。请逐步推理，最后把答案放在 \\boxed{} 中。"

COT_PROMPT_TEMPLATE_ZH = "请逐步解答以下问题，最后把答案放在 \\\\boxed{{}} 中。\n\n问题：{question}\n\n让我们逐步思考："

TOT_SYSTEM_ZH = "你是一个推理助手。请系统性地思考，评估多种解题路径。"

TOT_GENERATE_PROMPT_ZH = "当前推理状态：\n{context}\n\n请生成 {n_branches} 个不同的下一步推理方向，每个方向都应该推进对以下问题的解答：\n{question}\n\n请用编号列表列出："

TOT_EVALUATE_PROMPT_ZH = "问题：{question}\n\n当前推理步骤：{step}\n\n请给这个推理步骤打分（1-10分），考虑：\n- 它是否朝着正确答案推进？\n- 逻辑是否合理？\n- 有多大可能导向正确答案？\n\n只返回一个数字1-10。"

TOT_DECIDE_PROMPT_ZH = "基于以下推理线索，确定问题的最终答案：\n\n问题：{question}\n\n推理线索：\n{traces}\n\n最终答案请放在 \\\\boxed{{}} 中。"

# ── Language Detection ──────────────────────────────────────────────

def is_chinese_task(task_name: str) -> bool:
    """Detect whether the task/dataset uses Chinese prompts."""
    return task_name.lower().startswith("c_") or task_name.lower() in ("cmath", "c_aime")

def get_prompts(task_name: str) -> dict:
    """Return the appropriate prompt set based on task language."""
    if is_chinese_task(task_name):
        return {
            "io_system": IO_SYSTEM_ZH,
            "cot_system": COT_SYSTEM_ZH,
            "cot_template": COT_PROMPT_TEMPLATE_ZH,
            "tot_system": TOT_SYSTEM_ZH,
            "tot_generate": TOT_GENERATE_PROMPT_ZH,
            "tot_evaluate": TOT_EVALUATE_PROMPT_ZH,
            "tot_decide": TOT_DECIDE_PROMPT_ZH,
        }
    return {
        "io_system": IO_SYSTEM,
        "cot_system": COT_SYSTEM,
        "cot_template": COT_PROMPT_TEMPLATE,
        "tot_system": TOT_SYSTEM,
        "tot_generate": TOT_GENERATE_PROMPT,
        "tot_evaluate": TOT_EVALUATE_PROMPT,
        "tot_decide": TOT_DECIDE_PROMPT,
    }


# ── Answer Extraction ────────────────────────────────────────────────

def extract_answer(text: str) -> Optional[str]:
    """Extract answer from model output. Tries \\boxed{} first, then number patterns."""
    # Try \boxed{...} pattern with nested brace support
    boxed_match = re.search(r'\\boxed\{', text)
    if boxed_match:
        start = boxed_match.end()
        depth = 1
        i = start
        while i < len(text) and depth > 0:
            if text[i] == '{':
                depth += 1
            elif text[i] == '}':
                depth -= 1
            i += 1
        if depth == 0:
            inner = text[start:i-1].strip()
            if inner:
                return inner

    # Try "Answer: ..." pattern
    ans = re.search(r'(?:answer|Answer|final answer|Final Answer)[:\s]+(.+?)(?:\.|$)', text)
    if ans:
        return ans.group(1).strip()

    # Try last number in text
    numbers = re.findall(r'-?\d+(?:\.\d+)?(?:/\d+)?', text)
    if numbers:
        return numbers[-1].strip()

    return None


def normalize_answer(ans: Optional[str]) -> str:
    """Normalize an answer for comparison."""
    if ans is None:
        return ""
    ans = ans.strip().lower()
    # Remove LaTeX commands
    ans = re.sub(r'\\[a-z]+', '', ans)
    # Handle LaTeX fractions: {a}{b} -> a/b
    ans = re.sub(r'\}\s*\{', '/', ans)
    ans = re.sub(r'[{}]', '', ans)
    ans = ans.replace(' ', '')
    return ans


def is_correct(predicted: Optional[str], expected: str) -> bool:
    """Check if predicted answer matches expected."""
    if predicted is None:
        return False
    pred_norm = normalize_answer(predicted)
    exp_norm = normalize_answer(expected)
    if pred_norm == exp_norm:
        return True
    # Try numeric comparison
    try:
        return abs(float(pred_norm) - float(exp_norm)) < 0.01
    except (ValueError, TypeError):
        return pred_norm == exp_norm


# ── Baseline Runner ──────────────────────────────────────────────────

class BaselineRunner:
    """Run baseline evaluations on a given model."""

    def __init__(
        self,
        model: str = "qwen2.5:32b",
        temperature: float = 0.0,
        seed: int = 42,
        results_dir: str = "data/results",
    ):
        self.model_name = model
        self.temperature = temperature
        self.seed = seed
        self.results_dir = Path(results_dir)
        self.results_dir.mkdir(parents=True, exist_ok=True)

    # ── IO ────────────────────────────────────────────────────────

    def evaluate_io(self, dataset: list, task_name: str = "gsm8k") -> dict:
        """Input-Output: direct answer, no reasoning chain."""
        client = make_client(self.model_name, temperature=self.temperature)
        prompts = get_prompts(task_name)
        results = []
        correct = 0

        print(f"\n{'='*60}")
        print(f"IO Evaluation — {self.model_name} on {task_name}")
        print(f"{'='*60}")

        for item in tqdm(dataset, desc="IO"):
            question = item["question"]
            expected = item["answer"]

            start = time.time()
            response = client.generate(
                prompt=question,
                system=prompts["io_system"],
            )
            elapsed = time.time() - start

            predicted = extract_answer(response)
            c = is_correct(predicted, expected)
            if c:
                correct += 1

            results.append({
                "question": question,
                "expected": expected,
                "predicted": predicted,
                "full_response": response,
                "correct": c,
                "time_s": round(elapsed, 2),
                "tokens": client.count_tokens(response),
            })

            if len(results) <= 2:
                status = "✓" if c else "✗"
                print(f"  [{status}] Q: {question[:60]}...")
                print(f"        Pred: {predicted} | Exp: {expected}")

        accuracy = correct / len(dataset) if dataset else 0
        summary = {
            "method": "IO",
            "model": self.model_name,
            "task": task_name,
            "samples": len(dataset),
            "correct": correct,
            "accuracy": round(accuracy, 4),
            "avg_time_s": round(np.mean([r["time_s"] for r in results]), 2),
        }

        return self._save_results(task_name, "IO", summary, results)

    # ── CoT ───────────────────────────────────────────────────────

    def evaluate_cot(self, dataset: list, task_name: str = "gsm8k") -> dict:
        """Chain-of-Thought with step-by-step reasoning."""
        client = make_client(self.model_name, temperature=self.temperature)
        prompts = get_prompts(task_name)
        results = []
        correct = 0

        print(f"\n{'='*60}")
        print(f"CoT Evaluation — {self.model_name} on {task_name}")
        print(f"{'='*60}")

        for item in tqdm(dataset, desc="CoT"):
            question = item["question"]
            expected = item["answer"]
            prompt = prompts["cot_template"].format(question=question)

            start = time.time()
            response = client.generate(
                prompt=prompt,
                system=prompts["cot_system"],
            )
            elapsed = time.time() - start

            predicted = extract_answer(response)
            c = is_correct(predicted, expected)
            if c:
                correct += 1

            results.append({
                "question": question,
                "expected": expected,
                "predicted": predicted,
                "full_response": response,
                "correct": c,
                "time_s": round(elapsed, 2),
                "tokens": client.count_tokens(response),
            })

            if len(results) <= 2:
                status = "✓" if c else "✗"
                print(f"  [{status}] Q: {question[:60]}...")
                print(f"        Pred: {predicted} | Exp: {expected}")

        accuracy = correct / len(dataset) if dataset else 0
        summary = {
            "method": "CoT",
            "model": self.model_name,
            "task": task_name,
            "samples": len(dataset),
            "correct": correct,
            "accuracy": round(accuracy, 4),
            "avg_time_s": round(np.mean([r["time_s"] for r in results]), 2),
        }

        return self._save_results(task_name, "CoT", summary, results)

    # ── CoT-SC ────────────────────────────────────────────────────

    def evaluate_cot_sc(
        self,
        dataset: list,
        task_name: str = "gsm8k",
        n_chains: int = 5,
        max_workers: int = 3,
    ) -> dict:
        """CoT with Self-Consistency: multiple chains + majority vote."""
        client = make_client(self.model_name, temperature=0.7)
        prompts = get_prompts(task_name)
        results = []
        correct = 0

        print(f"\n{'='*60}")
        print(f"CoT-SC Evaluation ({n_chains} chains) — {self.model_name} on {task_name}")
        print(f"{'='*60}")

        for item in tqdm(dataset, desc="CoT-SC"):
            question = item["question"]
            expected = item["answer"]
            prompt = prompts["cot_template"].format(question=question)

            # Generate multiple chains (in parallel)
            chains = []
            start = time.time()
            with ThreadPoolExecutor(max_workers=max_workers) as pool:
                futures = [
                    pool.submit(client.generate, prompt, prompts["cot_system"])
                    for _ in range(n_chains)
                ]
                for f in as_completed(futures):
                    chains.append(f.result())
            elapsed = time.time() - start

            # Extract answers and vote
            answers = [extract_answer(c) for c in chains]
            # Majority vote (pick most common non-None answer)
            valid_answers = [a for a in answers if a is not None]
            if valid_answers:
                from collections import Counter
                answer_counts = Counter(valid_answers)
                predicted = answer_counts.most_common(1)[0][0]
            else:
                predicted = None

            c = is_correct(predicted, expected)
            if c:
                correct += 1

            results.append({
                "question": question,
                "expected": expected,
                "predicted": predicted,
                "chains": chains,
                "answers": answers,
                "correct": c,
                "time_s": round(elapsed, 2),
                "tokens": sum(client.count_tokens(ch) for ch in chains),
            })

        accuracy = correct / len(dataset) if dataset else 0
        summary = {
            "method": f"Cot-SC (n={n_chains})",
            "model": self.model_name,
            "task": task_name,
            "samples": len(dataset),
            "correct": correct,
            "accuracy": round(accuracy, 4),
            "avg_time_s": round(np.mean([r["time_s"] for r in results]), 2),
        }

        return self._save_results(task_name, f"Cot-SC_{n_chains}", summary, results)

    # ── ToT ───────────────────────────────────────────────────────

    def evaluate_tot(
        self,
        dataset: list,
        task_name: str = "gsm8k",
        n_branches: int = 3,
        max_depth: int = 2,
    ) -> dict:
        """Tree-of-Thought with beam search."""
        client = make_client(self.model_name, temperature=0.7, max_tokens=1024)
        prompts = get_prompts(task_name)
        results = []
        correct = 0

        print(f"\n{'='*60}")
        print(f"ToT Evaluation (b={n_branches}, d={max_depth}) — {self.model_name} on {task_name}")
        print(f"{'='*60}")

        for item in tqdm(dataset, desc="ToT"):
            question = item["question"]
            expected = item["answer"]

            start = time.time()
            try:
                predicted = self._tot_search(client, prompts, question, n_branches, max_depth)
            except Exception as e:
                print(f"  ToT error: {e}")
                predicted = None
            elapsed = time.time() - start

            c = is_correct(predicted, expected)
            if c:
                correct += 1

            results.append({
                "question": question,
                "expected": expected,
                "predicted": predicted,
                "correct": c,
                "time_s": round(elapsed, 2),
            })

        accuracy = correct / len(dataset) if dataset else 0
        summary = {
            "method": f"ToT (b={n_branches}, d={max_depth})",
            "model": self.model_name,
            "task": task_name,
            "samples": len(dataset),
            "correct": correct,
            "accuracy": round(accuracy, 4),
            "avg_time_s": round(np.mean([r["time_s"] for r in results]), 2),
        }

        return self._save_results(task_name, f"ToT_b{n_branches}_d{max_depth}", summary, results)

    def _tot_search(
        self, client: LLMClient, prompts: dict, question: str, n_branches: int, max_depth: int
    ) -> Optional[str]:
        """Internal ToT beam search."""
        # Level 0: Generate initial thoughts
        beam = self._tot_generate(client, prompts, question, "", n_branches)

        for depth in range(1, max_depth + 1):
            new_beam = []
            for thought in beam:
                context = "\n".join(f"- {t}" for t in beam[:3])
                next_steps = self._tot_generate(client, prompts, question, context, n_branches)
                for step in next_steps:
                    score = self._tot_evaluate(client, prompts, question, step)
                    new_beam.append((step, score))
            # Keep top n_branches
            new_beam.sort(key=lambda x: x[1], reverse=True)
            beam = [s for s, _ in new_beam[:n_branches]]

        # Decide final answer from best traces
        traces = "\n".join(f"Trace {i+1}: {t}" for i, t in enumerate(beam))
        decide_prompt = prompts["tot_decide"].format(question=question, traces=traces)
        decision = client.generate(decide_prompt, system=prompts["tot_system"])
        return extract_answer(decision)

    def _tot_generate(self, client: LLMClient, prompts: dict, question: str, context: str, n: int) -> list:
        prompt = prompts["tot_generate"].format(
            question=question, context=context or "No prior reasoning yet.", n_branches=n
        )
        response = client.generate(prompt, system=prompts["tot_system"])
        # Parse numbered lines
        steps = []
        for line in response.split("\n"):
            line = line.strip()
            if line and (line[0].isdigit() and ". " in line[:4]):
                steps.append(line.split(". ", 1)[1])
        return steps[:n]

    def _tot_evaluate(self, client: LLMClient, prompts: dict, question: str, step: str) -> float:
        prompt = prompts["tot_evaluate"].format(question=question, step=step)
        response = client.generate(prompt, system=prompts["tot_system"])
        try:
            return float(extract_answer(response) or response.strip()[:2])
        except (ValueError, TypeError):
            return 5.0

    # ── Results Saving ────────────────────────────────────────────

    def _save_results(self, task: str, method: str, summary: dict, details: list) -> dict:
        """Save results to JSON and return the summary."""
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        safe_method = method.replace(" ", "_").replace("(", "").replace(")", "").replace(",", "")
        filename = f"{task}_{safe_method}_{self.model_name.replace(':', '_')}_{timestamp}.json"
        filepath = self.results_dir / filename

        output = {
            "summary": summary,
            "details": details,
        }

        with open(filepath, "w") as f:
            json.dump(output, f, indent=2, ensure_ascii=False)

        print(f"\nResults saved: {filepath}")
        print(f"  Accuracy: {summary['accuracy']:.2%} ({summary['correct']}/{summary['samples']})")
        print(f"  Avg time: {summary.get('avg_time_s', 'N/A')}s")

        return summary

    # ── Run All ───────────────────────────────────────────────────

    def run_all(
        self,
        dataset: list,
        task_name: str = "gsm8k",
        methods: Optional[list] = None,
        n_chains_sc: int = 5,
        n_branches_tot: int = 3,
        max_depth_tot: int = 2,
    ) -> dict:
        """Run selected baseline methods and return a comparison summary."""
        if methods is None:
            methods = ["IO", "CoT", "CoT-SC", "ToT"]

        all_results = {}

        if "IO" in methods:
            all_results["IO"] = self.evaluate_io(dataset, task_name)

        if "CoT" in methods:
            all_results["CoT"] = self.evaluate_cot(dataset, task_name)

        if "CoT-SC" in methods:
            all_results["CoT-SC"] = self.evaluate_cot_sc(dataset, task_name, n_chains=n_chains_sc)

        if "ToT" in methods:
            all_results["ToT"] = self.evaluate_tot(dataset, task_name, n_branches=n_branches_tot, max_depth=max_depth_tot)

        # Print comparison table
        print(f"\n{'='*60}")
        print(f"Comparison — {self.model_name} on {task_name}")
        print(f"{'='*60}")
        print(f"{'Method':<20} {'Accuracy':<12} {'Time(s)':<10}")
        print(f"{'-'*42}")
        for method, res in all_results.items():
            acc = f"{res['accuracy']:.2%}"
            t = str(res.get("avg_time_s", "N/A"))
            print(f"{method:<20} {acc:<12} {t:<10}")

        return all_results


# ── CLI Entry Point ──────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Run LLM baseline evaluation")
    parser.add_argument("--model", default="qwen2.5:32b", help="Ollama model name")
    parser.add_argument("--task", choices=["gsm8k", "math", "aime", "c_gsm8k", "cmath", "c_aime"], default="gsm8k")
    parser.add_argument("--samples", type=int, default=5, help="Number of test samples")
    parser.add_argument("--methods", nargs="+", default=["IO", "CoT"],
                        help="Methods to run: IO, CoT, CoT-SC, ToT")
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()

    sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
    from src.samples import get_dataset
    dataset = get_dataset(args.task, args.samples)

    runner = BaselineRunner(model=args.model, temperature=args.temperature)
    runner.run_all(dataset, task_name=args.task, methods=args.methods)
