Chain of Thought Reasoning: Advanced Techniques for LLM Reasoning

import torch
import torch.nn as nn

class CoTPromptTemplate:
    """Template for chain-of-thought prompting."""
    
    def __init__(self, template_type="standard"):
        self.template_type = template_type
        
    def format(self, question, context=None):
        if self.template_type == "standard":
            return f"""Question: {question}
Let's think step by step and show your reasoning.
"""
        elif self.template_type == "detailed":
            return f"""Question: {question}

Please reason through this problem step by step:
1. First, identify what is being asked
2. Break down the problem into components
3. Analyze each component systematically
4. Combine insights to form a solution
5. Verify the solution makes sense

Show your work for each step.
"""
        elif self.template_type == "math":
            return f"""Solve the following math problem step by step:

{question}

Break down your solution:
- State the given information
- Identify the approach
- Execute each calculation
- State the final answer
"""
        return f"Question: {question}\nLet's think step by step.\n"


class EntropyGuidedCoT:
    """Entropy-guided chain-of-thought with dynamic depth control."""
    
    def __init__(self, model, tokenizer, entropy_threshold=0.5):
        self.model = model
        self.tokenizer = tokenizer
        self.entropy_threshold = entropy_threshold
        
    def generate_with_reasoning(self, prompt, max_reasoning_steps=5):
        """Generate response with entropy-guided reasoning depth."""
        reasoning_steps = 0
        all_outputs = []
        
        while reasoning_steps < max_reasoning_steps:
            # Generate next reasoning segment
            outputs = self.model.generate(
                prompt + "\n" + " ".join(all_outputs),
                max_new_tokens=100,
                do_sample=True,
                return_dict_in_generate=True
            )
            
            # Compute entropy of output distribution
            entropy = self.compute_entropy(outputs)
            all_outputs.append(outputs)
            
            # Check if we should continue reasoning
            if entropy < self.entropy_threshold:
                break
                
            reasoning_steps += 1
            
        return " ".join(all_outputs)
    
    def compute_entropy(self, outputs):
        """Compute entropy of output distribution."""
        # Simplified entropy computation
        probs = torch.softmax(outputs.scores[-1], dim=-1)
        entropy = -(probs * torch.log(probs + 1e-10)).sum(-1)
        return entropy.item()


class MultiStepReasoning:
    """Multi-step reasoning with explicit intermediate verification."""
    
    def __init__(self, model, num_steps=4):
        self.model = model
        self.num_steps = num_steps
        
    def reason(self, question):
        """Execute multi-step reasoning with verification."""
        current_state = {"question": question, "reasoning": [], "conclusion": None}
        
        for step in range(self.num_steps):
            # Generate reasoning for current step
            step_output = self.generate_step(current_state)
            current_state["reasoning"].append(step_output)
            
            # Verify step output
            if not self.verify_step(step_output, current_state):
                # Attempt correction
                step_output = self.correct_step(step_output, current_state)
                current_state["reasoning"][-1] = step_output
            
            # Update state for next step
            current_state = self.update_state(current_state, step_output)
            
        # Generate final conclusion
        current_state["conclusion"] = self.generate_conclusion(current_state)
        return current_state
    
    def generate_step(self, state):
        """Generate reasoning for current step."""
        prompt = f"""Current question: {state['question']}

Reasoning so far:
{chr(10).join(state['reasoning'])}

Generate the next step in reasoning:
"""
        return self.model.generate(prompt, max_new_tokens=100)
    
    def verify_step(self, step_output, state):
        """Verify that the step is valid and consistent."""
        # Check for contradictions with previous reasoning
        # Simplified: just check for obvious errors
        return len(step_output) > 10  # Basic sanity check
    
    def correct_step(self, step_output, state):
        """Correct an invalid step."""
        correction_prompt = f"""The previous step had issues:

Step: {step_output}

Previous reasoning:
{chr(10).join(state['reasoning'])}

Please provide a corrected step:
"""
        return self.model.generate(correction_prompt, max_new_tokens=100)
    
    def update_state(self, state, step_output):
        """Update reasoning state with new step."""
        state["current_context"] = state.get("current_context", "") + " " + step_output
        return state
    
    def generate_conclusion(self, state):
        """Generate final conclusion from reasoning."""
        conclusion_prompt = f"""Based on the following reasoning:

{chr(10).join(state['reasoning'])}

Provide the final answer:
"""
        return self.model.generate(conclusion_prompt, max_new_tokens=50)

from dataclasses import dataclass
from typing import Optional


@dataclass
class CoTTemplate:
    name: str
    system: str
    user_prefix: str
    reasoning_guide: str


class CoTPromptLibrary:
    """Production-ready CoT prompt templates for common domains."""

    TEMPLATES = {
        "math": CoTTemplate(
            name="Mathematical Reasoning",
            system="You are a precise mathematical assistant. Always show every calculation step.",
            user_prefix="Solve the following problem:",
            reasoning_guide="""
Step 1: Identify all given values and what is being asked.
Step 2: Choose the appropriate formula or approach.
Step 3: Substitute values and compute, showing each arithmetic step.
Step 4: Check units and verify the answer is reasonable.
Final answer: [state answer with units]
""",
        ),
        "code_debug": CoTTemplate(
            name="Code Debugging",
            system="You are a senior software engineer. Debug methodically.",
            user_prefix="Debug the following code:",
            reasoning_guide="""
Step 1: Read the code and identify what it is supposed to do.
Step 2: Trace through the execution with example inputs.
Step 3: Identify where the actual behavior diverges from expected.
Step 4: Pinpoint the root cause (logic error, off-by-one, null check, etc.).
Step 5: Write the fix and explain why it resolves the issue.
""",
        ),
        "legal_analysis": CoTTemplate(
            name="Legal Analysis",
            system="You are a legal analyst. Apply rules to facts systematically.",
            user_prefix="Analyze the following legal scenario:",
            reasoning_guide="""
Step 1: Identify the applicable legal rule or statute.
Step 2: Extract the key facts relevant to each element of the rule.
Step 3: Apply each rule element to the facts.
Step 4: Identify any exceptions or defenses.
Step 5: State the likely outcome and confidence level.
""",
        ),
        "system_design": CoTTemplate(
            name="System Design",
            system="You are a principal engineer. Design systems with explicit trade-off reasoning.",
            user_prefix="Design a system for the following requirement:",
            reasoning_guide="""
Step 1: Clarify scale requirements (QPS, data volume, latency SLA).
Step 2: Identify the core data model and access patterns.
Step 3: Choose storage and compute components with justification.
Step 4: Address bottlenecks and single points of failure.
Step 5: Outline monitoring, scaling, and operational concerns.
""",
        ),
        "data_analysis": CoTTemplate(
            name="Data Analysis",
            system="You are a data scientist. Analyze data with explicit statistical reasoning.",
            user_prefix="Analyze the following data:",
            reasoning_guide="""
Step 1: Describe the data shape, types, and any missing values.
Step 2: Compute relevant summary statistics.
Step 3: Identify patterns, trends, or anomalies.
Step 4: Assess statistical significance where applicable.
Step 5: State conclusions and their limitations.
""",
        ),
    }

    @classmethod
    def format(cls, domain: str, question: str) -> dict[str, str]:
        """Format a question with the appropriate CoT template."""
        template = cls.TEMPLATES.get(domain, cls.TEMPLATES["math"])
        return {
            "system": template.system,
            "user": f"{template.user_prefix}\n\n{question}\n\n{template.reasoning_guide}",
        }

    @classmethod
    def zero_shot(cls, question: str, detail: str = "standard") -> str:
        """Simple zero-shot CoT without domain-specific template."""
        if detail == "minimal":
            return f"{question}\n\nLet's think step by step."
        elif detail == "structured":
            return (
                f"{question}\n\n"
                "Think through this carefully:\n"
                "1. What is being asked?\n"
                "2. What information do I have?\n"
                "3. What approach should I use?\n"
                "4. Execute the approach step by step.\n"
                "5. Verify the answer."
            )
        return f"{question}\n\nLet's reason through this step by step:"

FEW_SHOT_MATH_COT = """
Q: A store has 48 apples. They sell 3/4 of them in the morning and receive a delivery of 20 more. How many apples are there now?

A: Let me work through this step by step.
- Start: 48 apples
- Sold in morning: 48 × 3/4 = 36 apples sold
- Remaining after morning: 48 - 36 = 12 apples
- After delivery: 12 + 20 = 32 apples

Answer: 32 apples

Q: {question}

A: Let me work through this step by step.
"""

import re
from collections import Counter


class SelfConsistencyCoT:
    """
    Generate N independent reasoning chains and vote on the final answer.
    Most effective for tasks with discrete answer spaces (math, classification).
    """

    def __init__(self, model, n_samples: int = 10, temperature: float = 0.7):
        self.model = model
        self.n_samples = n_samples
        self.temperature = temperature

    def generate_chain(self, prompt: str) -> str:
        """Generate one reasoning chain."""
        return self.model.generate(
            prompt,
            temperature=self.temperature,
            max_new_tokens=512,
        )

    def extract_answer(self, chain: str) -> Optional[str]:
        """Extract final answer from a reasoning chain."""
        # Look for common answer patterns
        patterns = [
            r"(?:final answer|answer|result):\s*(.+?)(?:\n|$)",
            r"(?:therefore|thus|so),?\s+(?:the answer is\s+)?(.+?)(?:\.|$)",
            r"=\s*(\d+(?:\.\d+)?)\s*$",
        ]
        for pattern in patterns:
            match = re.search(pattern, chain, re.IGNORECASE | re.MULTILINE)
            if match:
                return match.group(1).strip()
        # Fallback: last non-empty line
        lines = [l.strip() for l in chain.split("\n") if l.strip()]
        return lines[-1] if lines else None

    def solve(self, question: str) -> dict:
        """Solve using self-consistency voting."""
        prompt = CoTPromptLibrary.zero_shot(question, detail="structured")

        chains = [self.generate_chain(prompt) for _ in range(self.n_samples)]
        answers = [self.extract_answer(c) for c in chains]
        valid_answers = [a for a in answers if a is not None]

        if not valid_answers:
            return {"answer": None, "confidence": 0.0, "chains": chains}

        vote_counts = Counter(valid_answers)
        best_answer, count = vote_counts.most_common(1)[0]
        confidence = count / len(valid_answers)

        return {
            "answer": best_answer,
            "confidence": confidence,
            "vote_distribution": dict(vote_counts),
            "chains": chains,
        }

from dataclasses import dataclass, field


@dataclass
class ThoughtNode:
    thought: str
    parent: Optional["ThoughtNode"] = None
    children: list["ThoughtNode"] = field(default_factory=list)
    value: float = 0.0
    depth: int = 0

    def path_to_root(self) -> list[str]:
        """Return the full reasoning path from root to this node."""
        path = []
        node = self
        while node is not None:
            path.append(node.thought)
            node = node.parent
        return list(reversed(path))


class TreeOfThought:
    """
    BFS/DFS over reasoning paths, pruning low-value branches.
    Best for problems where intermediate steps can be evaluated.
    """

    def __init__(self, model, branching_factor: int = 3, max_depth: int = 4):
        self.model = model
        self.branching_factor = branching_factor
        self.max_depth = max_depth

    def generate_thoughts(self, question: str, context: str, n: int) -> list[str]:
        """Generate n candidate next thoughts."""
        prompt = (
            f"Question: {question}\n\n"
            f"Reasoning so far:\n{context}\n\n"
            f"Generate {n} different next reasoning steps (one per line):"
        )
        output = self.model.generate(prompt, temperature=0.9, max_new_tokens=200)
        lines = [l.strip() for l in output.split("\n") if l.strip()]
        return lines[:n]

    def evaluate_thought(self, question: str, path: list[str]) -> float:
        """Score a reasoning path (0=wrong, 0.5=partial, 1=correct)."""
        prompt = (
            f"Question: {question}\n\n"
            f"Reasoning path:\n" + "\n".join(path) + "\n\n"
            "Rate this reasoning path on a scale of 0-1 where:\n"
            "0 = clearly wrong direction\n"
            "0.5 = plausible but uncertain\n"
            "1 = correct and complete\n"
            "Score (just the number):"
        )
        output = self.model.generate(prompt, temperature=0.0, max_new_tokens=5)
        try:
            return float(re.search(r"[01](?:\.\d+)?", output).group())
        except (AttributeError, ValueError):
            return 0.5

    def solve(self, question: str) -> dict:
        """BFS Tree of Thought search."""
        root = ThoughtNode(thought=f"Problem: {question}")
        frontier = [root]
        best_node = root

        for depth in range(self.max_depth):
            next_frontier = []
            for node in frontier:
                context = "\n".join(node.path_to_root())
                candidates = self.generate_thoughts(
                    question, context, self.branching_factor
                )

                for thought in candidates:
                    child = ThoughtNode(
                        thought=thought, parent=node, depth=depth + 1
                    )
                    path = child.path_to_root()
                    child.value = self.evaluate_thought(question, path)
                    node.children.append(child)

                    if child.value > best_node.value:
                        best_node = child

                    if child.value > 0.8:  # Early exit on high-confidence path
                        return {"answer": thought, "path": path, "value": child.value}

                # Keep top-k branches
                scored = sorted(node.children, key=lambda n: n.value, reverse=True)
                next_frontier.extend(scored[: self.branching_factor])

            frontier = next_frontier[:self.branching_factor]  # Beam width

        return {
            "answer": best_node.thought,
            "path": best_node.path_to_root(),
            "value": best_node.value,
        }

import json
import time
from pathlib import Path


class CoTEvaluator:
    """Evaluate CoT accuracy, cost, and latency on a test set."""

    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def count_tokens(self, text: str) -> int:
        return len(self.tokenizer.encode(text))

    def evaluate_dataset(
        self,
        examples: list[dict],  # [{"question": ..., "answer": ...}]
        prompt_fn,              # Function(question) -> prompt string
        extract_fn,             # Function(output) -> answer string
        n_examples: int = 100,
    ) -> dict:
        """Run evaluation and return accuracy/cost/latency metrics."""
        correct = 0
        total_tokens = 0
        total_latency = 0.0

        for ex in examples[:n_examples]:
            prompt = prompt_fn(ex["question"])
            input_tokens = self.count_tokens(prompt)

            t0 = time.perf_counter()
            output = self.model.generate(prompt, max_new_tokens=512, temperature=0.0)
            latency = time.perf_counter() - t0

            output_tokens = self.count_tokens(output)
            predicted = extract_fn(output)

            if predicted and predicted.lower().strip() == ex["answer"].lower().strip():
                correct += 1

            total_tokens += input_tokens + output_tokens
            total_latency += latency

        n = min(n_examples, len(examples))
        return {
            "accuracy": correct / n,
            "avg_tokens_per_query": total_tokens / n,
            "avg_latency_sec": total_latency / n,
            "total_examples": n,
        }

    def compare_strategies(self, examples: list[dict]) -> None:
        """Compare direct vs CoT strategies and print a report."""
        lib = CoTPromptLibrary()
        extractor = lambda o: SelfConsistencyCoT(self.model).extract_answer(o)

        strategies = {
            "direct": lambda q: q,
            "zero_shot_cot": lambda q: lib.zero_shot(q, "minimal"),
            "structured_cot": lambda q: lib.zero_shot(q, "structured"),
        }

        print(f"{'Strategy':<20} {'Accuracy':>10} {'Avg Tokens':>12} {'Avg Latency':>13}")
        print("-" * 60)
        for name, prompt_fn in strategies.items():
            metrics = self.evaluate_dataset(examples, prompt_fn, extractor)
            print(
                f"{name:<20} {metrics['accuracy']:>9.1%} "
                f"{metrics['avg_tokens_per_query']:>12.0f} "
                f"{metrics['avg_latency_sec']:>12.2f}s"
            )

def detect_reasoning_loop(steps: list[str], similarity_threshold: float = 0.85) -> bool:
    """Detect if reasoning is looping by checking step similarity."""
    if len(steps) < 3:
        return False
    last = steps[-1].lower().split()
    prev = steps[-2].lower().split()
    if not last or not prev:
        return False
    # Jaccard similarity
    intersection = len(set(last) & set(prev))
    union = len(set(last) | set(prev))
    similarity = intersection / union if union else 0
    return similarity > similarity_threshold

VERIFICATION_PROMPT = """
You previously reasoned:
{reasoning}

And reached the answer: {answer}

Verify: Does each step logically follow from the previous? Is the final answer consistent with the reasoning?
If you find an error, state which step is wrong and provide the correct reasoning.
"""

def tiered_cot(question: str, difficulty_classifier) -> str:
    """Use cheap CoT for easy questions, expensive for hard ones."""
    difficulty = difficulty_classifier(question)  # "easy", "medium", "hard"

    if difficulty == "easy":
        return direct_answer(question)          # No CoT
    elif difficulty == "medium":
        return zero_shot_cot(question)          # Standard CoT
    else:
        return self_consistency_cot(question, n=5)  # Voting