biomind_recursive_dmn.py

"""
BIOMIND Recursive DMN Reflection Loop Implementation
==================================================

Implements a metacognitive reflection cycle within the Decoupled HIPECA v2.5 framework
using a Qwen2-0.5B-Instruct (Q4_K_M) as the DMN Reflector for System-2 thinking.

Enhanced with Memory Integration:
- Similarity-driven memory recall from Wiki-trained knowledge
- Confidence-gated memory access bandwidth
- Counterfactual recall for hypothesis stress-testing
- Self-adversarial memory checking
- Multi-specialist + memory consensus

Key Components:
1. DMNReflector - Qwen2-0.5B metacognitive SLM
2. MemoryRecallSystem - Similarity-driven Wiki knowledge access
3. Recursive GLW Update Rule with reflection injection
4. HIPECA Executive Gating (ACC & Cerebellum)
5. Memory-Integrated Specialist Dispatch with confidence gating

Author: Principal Neuro-AI Engineer
Date: January 7, 2026
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, field
from enum import Enum
import time
import json
import re
from pathlib import Path

from real_model_inference import RealSpecialistSystem

@dataclass
class MemoryTrace:
    """Represents a recalled memory trace from Wiki-trained knowledge"""
    content: str
    embedding: torch.Tensor
    relevance_score: float
    activation_strength: float
    source_type: str  # 'wiki_episodic', 'wiki_semantic', 'reasoning_episodic'
    timestamp: float
    confidence_weight: float

@dataclass
class CounterfactualMemory:
    """Represents contrasting memories for hypothesis testing"""
    original_concept: str
    contrasting_concept: str
    distinction_explanation: str
    memory_trace: MemoryTrace
    stress_test_score: float

@dataclass
class CognitiveState:
    """Represents the current cognitive state that must change each iteration"""
    memory_recall_radius: float = 1.0  # How broad to search memory
    abstraction_level: float = 0.5     # How abstract/concrete the thinking
    attention_distribution: torch.Tensor = None  # Where attention is focused
    specialist_role: str = "general"    # Current specialist role
    question_representation: str = ""   # How question is internally framed
    confidence_policy: float = 0.5      # "I know how to think about this"
    confidence_answer: float = 0.5      # "I trust the result I produced"
    last_executed_action: Optional[str] = None
    execution_confidence_pre: float = 0.0  # Confidence before execution
    execution_confidence_post: float = 0.0 # Confidence after execution
    reasoning_history: List[str] = field(default_factory=list)     # Raw specialist responses for iterative reasoning

class MDNAction(Enum):
    """MDN planning actions for cognitive control"""
    RECALL_MORE_MEMORY = "recall_more_memory"
    REFRAME_QUESTION = "reframe_question"
    COUNTERFACTUAL_SIMULATE = "counterfactual_simulate"
    SWITCH_SPECIALIST = "switch_specialist"
    DEEPEN_REASONING = "deepen_reasoning"
    REEVALUATE_ANSWER = "reevaluate_answer"
    EXPAND_ATTENTION = "expand_attention"
    NARROW_FOCUS = "narrow_focus"
    INCREASE_ABSTRACTION = "increase_abstraction"
    DECREASE_ABSTRACTION = "decrease_abstraction"
    EXECUTE_AND_EVALUATE = "execute_and_evaluate"

@dataclass
class MDNActionPlan:
    """MDN action planning result"""
    selected_action: MDNAction
    expected_confidence_gain: float
    state_change_required: Dict[str, Any]
    epistemic_value: float  # Expected information gain

class MDNPlanner(nn.Module):
    """
    MDN Action Planner - selects cognitive actions to maximize epistemic value

    Predicts expected confidence gain for each possible cognitive action
    and chooses the action that maximizes information gain, not speed.
    """

    def __init__(self, glw_dim: int = 512, num_actions: int = len(MDNAction)):
        super().__init__()
        self.glw_dim = glw_dim
        self.num_actions = num_actions

        # State encoder
        self.state_encoder = nn.Sequential(
            nn.Linear(glw_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128)
        )

        # Action value predictor (expected confidence gain)
        self.action_value_head = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, num_actions)
        )

        # Epistemic value predictor (information gain)
        self.epistemic_value_head = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, num_actions)
        )

        # State change predictor
        self.state_change_head = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 32)
        )

    def forward(self, glw_state: torch.Tensor, cognitive_state: CognitiveState,
                memory_context: Dict[str, Any], task_input: str = "", execution_history: List[Dict] = None) -> MDNActionPlan:
        """
        Plan the next cognitive action based on current state and memory context
        """
        execution_history = execution_history or []
        
        # Encode current GLW state
        state_encoded = self.state_encoder(glw_state.squeeze(0))

        # Predict action values (expected confidence gain)
        action_values = self.action_value_head(state_encoded)

        # Predict epistemic values (information gain)
        epistemic_values = self.epistemic_value_head(state_encoded)

        # Combine confidence gain and information gain
        # Weight epistemic value higher for exploration
        combined_values = action_values + 1.5 * epistemic_values

        # Boost EXECUTE action if confidences are high enough
        execute_idx = list(MDNAction).index(MDNAction.EXECUTE_AND_EVALUATE)
        if cognitive_state.confidence_policy > 0.6 and cognitive_state.confidence_answer > 0.5:
            combined_values[execute_idx] += 2.0  # Significant boost for EXECUTE when ready
        elif cognitive_state.confidence_policy > 0.4 or cognitive_state.confidence_answer > 0.3:
            combined_values[execute_idx] += 1.0  # Moderate boost for EXECUTE with reasonable confidence

        # Select best action
        best_action_idx = torch.argmax(combined_values).item()
        best_action = list(MDNAction)[best_action_idx]

        # Predict state changes required for this action
        state_changes = self._predict_state_changes(best_action, cognitive_state, memory_context, task_input)

        return MDNActionPlan(
            selected_action=best_action,
            expected_confidence_gain=action_values[best_action_idx].item(),
            state_change_required=state_changes,
            epistemic_value=epistemic_values[best_action_idx].item()
        )

    def _predict_state_changes(self, action: MDNAction, cognitive_state: CognitiveState,
                              memory_context: Dict[str, Any], task_input: str = "") -> Dict[str, Any]:
        """Predict what state changes are required for the selected action"""

        changes = {}

        if action == MDNAction.RECALL_MORE_MEMORY:
            changes['memory_recall_radius'] = min(2.0, cognitive_state.memory_recall_radius * 1.5)
            changes['attention_distribution'] = 'memory_focused'

        elif action == MDNAction.REFRAME_QUESTION:
            changes['question_representation'] = self._generate_reframed_question(
                cognitive_state.question_representation, memory_context
            )
            changes['abstraction_level'] = cognitive_state.abstraction_level * 0.8  # More concrete

        elif action == MDNAction.COUNTERFACTUAL_SIMULATE:
            changes['attention_distribution'] = 'counterfactual_focused'
            changes['memory_recall_radius'] = cognitive_state.memory_recall_radius * 1.2

        elif action == MDNAction.SWITCH_SPECIALIST:
            current_role = cognitive_state.specialist_role
            changes['specialist_role'] = self._select_alternative_specialist(current_role, task_input)

        elif action == MDNAction.DEEPEN_REASONING:
            changes['abstraction_level'] = min(1.0, cognitive_state.abstraction_level + 0.2)
            changes['memory_recall_radius'] = cognitive_state.memory_recall_radius * 0.9

        elif action == MDNAction.REEVALUATE_ANSWER:
            changes['attention_distribution'] = 'answer_focused'
            changes['confidence_answer'] = cognitive_state.confidence_answer * 0.7  # Reset

        elif action == MDNAction.EXPAND_ATTENTION:
            changes['attention_distribution'] = 'broad_focus'
            changes['memory_recall_radius'] = cognitive_state.memory_recall_radius * 1.3

        elif action == MDNAction.NARROW_FOCUS:
            changes['attention_distribution'] = 'narrow_focus'
            changes['memory_recall_radius'] = cognitive_state.memory_recall_radius * 0.7

        elif action == MDNAction.INCREASE_ABSTRACTION:
            changes['abstraction_level'] = min(1.0, cognitive_state.abstraction_level + 0.3)

        elif action == MDNAction.DECREASE_ABSTRACTION:
            changes['abstraction_level'] = max(0.0, cognitive_state.abstraction_level - 0.3)

        elif action == MDNAction.EXECUTE_AND_EVALUATE:
            changes['attention_distribution'] = 'execution_focused'
            # Include explicit answer format instructions for specialists
            changes['specialist_instructions'] = self._generate_specialist_instructions(task_input, cognitive_state)

        return changes

    def _generate_specialist_instructions(self, task_input: str, cognitive_state: CognitiveState) -> str:
        """Generate explicit instructions for specialists based on task type and cognitive context"""

        # Determine task type and generate appropriate instructions
        task_lower = task_input.lower()

        # Math questions - require single letter answers
        if ("math" in task_lower or "calculate" in task_lower or "algebra" in task_lower or
            "field" in task_lower or "extension" in task_lower or "equation" in task_lower or
            "degree" in task_lower or "solve" in task_lower):

            return ("You are a mathematical expert. Provide step-by-step reasoning, then give ONLY the single letter "
                   "of your final answer (A, B, C, or D) at the very end. Format: 'Final answer: B'")

        # General reasoning questions
        elif ("reasoning" in task_lower or "logic" in task_lower or "analyze" in task_lower):

            return ("You are a reasoning expert. Provide clear analysis, then give ONLY the single letter "
                   "of your final answer (A, B, C, or D) at the very end. Format: 'Final answer: B'")

        # General knowledge questions
        else:
            return ("You are a knowledge expert. Provide a clear explanation, then give ONLY the single letter "
                   "of your final answer (A, B, C, or D) at the very end. Format: 'Final answer: B'")

    def _generate_reframed_question(self, current_question: str, memory_context: Dict[str, Any]) -> str:
        """Generate a reframed version of the question based on memory insights"""
        # Simple reframing based on available memories
        memories = memory_context.get('recalled_memories', [])
        if memories:
            # Use memory content to reframe
            key_concept = memories[0].content.split(':')[0] if ':' in memories[0].content else memories[0].content[:20]
            return f"Considering {key_concept}: {current_question}"
        return current_question

    def _select_alternative_specialist(self, current_role: str, task_input: str = "") -> str:
        """Select an alternative specialist role using trained router"""
        if self.trained_router and task_input:
            try:
                with torch.no_grad():
                    router_result = self.trained_router(task_input)
                    selected = router_result['predicted_specialist']
                    
                    # Override obviously wrong predictions
                    task_lower = task_input.lower()
                    is_math_question = ("math" in task_lower or "calculate" in task_lower or "algebra" in task_lower or 
                                      "field" in task_lower or "extension" in task_lower or "group" in task_lower or 
                                      "ring" in task_lower or "polynomial" in task_lower or "equation" in task_lower or
                                      "sqrt" in task_lower or "degree" in task_lower or "finite" in task_lower)
                    
                    if is_math_question and selected != "qwen_math_expert" and selected != current_role:
                        print(f"   [CYCLE] Overriding alternative router ({selected}) -> qwen_math_expert for math question")
                        return "qwen_math_expert"
                    
                    if selected != current_role:
                        return selected
            except Exception as e:
                print(f"[WARN]  Alternative specialist selection failed: {e}")
        
        # Fallback to random selection
        roles = ['qwen_math_expert', 'qwen_general_reasoner', 'tiny_llama_planner', 'tiny_llama_critic']
        if current_role in roles:
            roles.remove(current_role)
        return np.random.choice(roles)

    def _evaluate_specialist_performance(self, cognitive_state: CognitiveState, 
                                       execution_history: List[Dict], task_input: str) -> float:
        """
        Evaluate how well the current specialist is performing based on execution history
        
        Returns performance score between 0-1, where 1 is excellent performance
        """
        if not execution_history:
            return 0.5  # Neutral score with no history
            
        # Analyze recent executions (last 3 cycles)
        recent_executions = execution_history[-3:]
        
        # Performance metrics
        confidence_stability = []
        answer_consistency = []
        basal_ganglia_rejections = 0
        
        for execution in recent_executions:
            # Confidence stability: how much confidence changes between pre/post execution
            conf_pre = execution.get('confidence_pre', 0.5)
            conf_post = execution.get('confidence_post', 0.5)
            stability = 1.0 - abs(conf_pre - conf_post)  # Higher when confidence is stable
            confidence_stability.append(stability)
            
            # Basal ganglia rejections: how often cognition is reopened
            bg_eval = execution.get('basal_ganglia_evaluation', {})
            if bg_eval.get('should_reopen_cognition', False):
                basal_ganglia_rejections += 1
        
        # Check if this is a math question that should be handled by math expert
        task_lower = task_input.lower()
        is_math_question = ("math" in task_lower or "calculate" in task_lower or "algebra" in task_lower or 
                          "field" in task_lower or "extension" in task_lower or "group" in task_lower or 
                          "ring" in task_lower or "polynomial" in task_lower or "equation" in task_lower or
                          "sqrt" in task_lower or "degree" in task_lower or "finite" in task_lower or
                          "s_" in task_lower or "symmetric" in task_lower or "permutation" in task_lower or
                          "index" in task_lower or "order" in task_lower or "subgroup" in task_lower or
                          "<" in task_lower or ">" in task_lower or "z_" in task_lower or "c_" in task_lower)
        
        # Penalize non-math specialists on math questions
        specialist_penalty = 0.0
        if is_math_question and cognitive_state.specialist_role != "qwen_math_expert":
            specialist_penalty = 0.3  # Significant penalty for wrong specialist on math
            
        # Calculate overall performance score
        avg_stability = np.mean(confidence_stability) if confidence_stability else 0.5
        rejection_rate = basal_ganglia_rejections / len(recent_executions) if recent_executions else 0.0
        
        # Performance = stability - rejection penalty - specialist penalty
        performance = avg_stability - (rejection_rate * 0.5) - specialist_penalty
        
        # Clamp to [0, 1]
        performance = max(0.0, min(1.0, performance))
        
        return performance

class BasalGangliaEvaluator(nn.Module):
    """
    Basal Ganglia-style post-execution evaluation gate

    Measures confidence drop, internal contradiction, memory disagreement
    after answer generation. If confidence decreases, suppresses commitment
    and reopens DMN loop.
    """

    def __init__(self, glw_dim: int = 512):
        super().__init__()
        self.glw_dim = glw_dim

        # Confidence drop detector
        self.confidence_drop_detector = nn.Sequential(
            nn.Linear(glw_dim + 2, 128),  # GLW + pre/post confidence
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()  # Probability of significant drop
        )

        # Internal contradiction detector
        self.contradiction_detector = nn.Sequential(
            nn.Linear(glw_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

        # Memory disagreement scorer
        self.memory_disagreement_scorer = nn.Sequential(
            nn.Linear(glw_dim + 10, 64),  # GLW + memory features
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )

    def evaluate_post_execution(self,
                               glw_state: torch.Tensor,
                               confidence_pre: float,
                               confidence_post: float,
                               generated_answer: str,
                               recalled_memories: List[MemoryTrace],
                               task_input: str) -> Dict[str, Any]:
        """
        Evaluate the generated answer and decide whether to commit or reopen cognition
        """

        # Detect confidence drop
        confidence_input = torch.cat([
            glw_state.squeeze(0),
            torch.tensor([confidence_pre, confidence_post], dtype=torch.float32)
        ])
        confidence_drop_prob = self.confidence_drop_detector(confidence_input.unsqueeze(0)).item()

        # Detect internal contradictions
        contradiction_prob = self.contradiction_detector(glw_state).item()

        # Score memory disagreement
        memory_features = self._extract_memory_features(recalled_memories, generated_answer)
        memory_input = torch.cat([glw_state.squeeze(0), torch.tensor(memory_features, dtype=torch.float32)])
        memory_disagreement = self.memory_disagreement_scorer(memory_input.unsqueeze(0)).item()

        # Overall evaluation
        confidence_drop_threshold = 0.3  # 30% drop is significant
        actual_drop = confidence_pre - confidence_post

        evaluation = {
            'confidence_drop': actual_drop,
            'confidence_drop_significant': actual_drop > confidence_drop_threshold,
            'internal_contradictions': contradiction_prob > 0.5,
            'memory_disagreement': memory_disagreement > 0.6,
            'should_reopen_cognition': False,
            'reopen_reason': None
        }

        # Decision logic: reopen if confidence drops significantly OR high contradiction/disagreement
        if evaluation['confidence_drop_significant']:
            evaluation['should_reopen_cognition'] = True
            evaluation['reopen_reason'] = f"Confidence dropped {actual_drop:.2f} (>{confidence_drop_threshold})"
        elif evaluation['internal_contradictions'] and evaluation['memory_disagreement']:
            evaluation['should_reopen_cognition'] = True
            evaluation['reopen_reason'] = "High internal contradiction and memory disagreement"

        return evaluation

    def _extract_memory_features(self, memories: List[MemoryTrace], answer: str) -> List[float]:
        """Extract features about memory-answer agreement"""
        if not memories:
            return [0.0] * 10

        features = []

        # Average relevance score
        features.append(np.mean([m.relevance_score for m in memories]))

        # Max relevance score
        features.append(max([m.relevance_score for m in memories]))

        # Number of memories
        features.append(len(memories))

        # Semantic overlap with answer (simplified)
        answer_words = set(answer.lower().split())
        memory_words = set()
        for m in memories:
            memory_words.update(m.content.lower().split())

        overlap = len(answer_words.intersection(memory_words))
        features.append(overlap / max(len(answer_words), 1))

        # Fill remaining features
        while len(features) < 10:
            features.append(0.0)

        return features[:10]

class CounterfactualReDecision(nn.Module):
    """
    Counterfactual re-decision mechanism

    When confidence is unstable or low, generates alternative answers
    by simulating different cognitive states and memory configurations.
    """

    def __init__(self, glw_dim: int = 512, num_alternatives: int = 3):
        super().__init__()
        self.glw_dim = glw_dim
        self.num_alternatives = num_alternatives

        # Alternative answer generator
        self.alternative_generator = nn.Sequential(
            nn.Linear(glw_dim + 100, 256),  # GLW + counterfactual features
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1)  # Confidence score for alternative
        )

        # Counterfactual state simulator
        self.state_simulator = nn.Sequential(
            nn.Linear(glw_dim + 20, 128),  # GLW + state perturbation
            nn.ReLU(),
            nn.Linear(128, glw_dim)
        )

    def generate_alternatives(self,
                            glw_state: torch.Tensor,
                            original_answer: str,
                            confidence: float,
                            cognitive_state: CognitiveState,
                            recalled_memories: List[MemoryTrace]) -> List[Dict[str, Any]]:
        """
        Generate alternative answers when confidence is unstable
        """

        alternatives = []

        # Generate counterfactual scenarios
        for i in range(self.num_alternatives):
            # Perturb cognitive state
            state_perturbation = self._generate_state_perturbation(cognitive_state)
            perturbed_glw = self.state_simulator(
                torch.cat([glw_state.squeeze(0), torch.tensor(state_perturbation)])
            )

            # Generate alternative based on perturbed state
            counterfactual_features = self._extract_counterfactual_features(
                original_answer, cognitive_state, recalled_memories, i
            )

            generator_input = torch.cat([perturbed_glw, torch.tensor(counterfactual_features)])
            alt_confidence = self.alternative_generator(generator_input.unsqueeze(0)).item()

            # Create alternative answer (simplified - in practice would use full generation)
            alternative = self._create_alternative_answer(
                original_answer, cognitive_state, recalled_memories, i
            )

            alternatives.append({
                'answer': alternative,
                'confidence': alt_confidence,
                'state_perturbation': state_perturbation,
                'counterfactual_scenario': f"scenario_{i}"
            })

        # Sort by confidence
        alternatives.sort(key=lambda x: x['confidence'], reverse=True)

        return alternatives

    def _generate_state_perturbation(self, cognitive_state: CognitiveState) -> List[float]:
        """Generate random perturbations to cognitive state"""
        perturbation = []

        # Perturb memory radius
        perturbation.append(cognitive_state.memory_recall_radius * np.random.uniform(0.5, 1.5))

        # Perturb abstraction level
        perturbation.append(np.clip(
            cognitive_state.abstraction_level + np.random.normal(0, 0.2), 0, 1
        ))

        # Perturb attention distribution (simplified)
        attention_perturbation = np.random.normal(0, 0.1, 5).tolist()
        perturbation.extend(attention_perturbation)

        # Perturb specialist role (random switch)
        perturbation.append(np.random.randint(0, 5))

        return perturbation

    def _extract_counterfactual_features(self,
                                       original_answer: str,
                                       cognitive_state: CognitiveState,
                                       memories: List[MemoryTrace],
                                       scenario_idx: int) -> List[float]:
        """Extract features for counterfactual generation"""
        features = []

        # Original answer length
        features.append(len(original_answer))

        # Memory count
        features.append(len(memories))

        # Cognitive state features
        features.append(cognitive_state.memory_recall_radius)
        features.append(cognitive_state.abstraction_level)
        features.extend(cognitive_state.attention_distribution)

        # Scenario index (one-hot style)
        scenario_onehot = [0.0] * self.num_alternatives
        scenario_onehot[scenario_idx] = 1.0
        features.extend(scenario_onehot)

        # Pad to 100
        while len(features) < 100:
            features.append(0.0)

        return features[:100]

    def _create_alternative_answer(self,
                                 original_answer: str,
                                 cognitive_state: CognitiveState,
                                 memories: List[MemoryTrace],
                                 scenario_idx: int) -> str:
        """
        Create alternative answer based on counterfactual scenario
        (Simplified implementation - would use full generation in practice)
        """

        # Different strategies based on scenario
        if scenario_idx == 0:
            # More conservative approach
            return f"Alternative (conservative): {original_answer} (with additional verification needed)"
        elif scenario_idx == 1:
            # More aggressive approach
            return f"Alternative (aggressive): {original_answer} (high confidence despite uncertainty)"
        else:
            # Balanced approach
            return f"Alternative (balanced): {original_answer} (considering {len(memories)} memory traces)"

class TerminationCriteria:
    """
    Convergence-based termination criteria

    Terminates when information gain diminishes and confidence converges,
    not when cycle/time limits are reached.
    """

    def __init__(self, min_information_gain: float = 0.01, confidence_stability_threshold: float = 0.05):
        self.min_information_gain = min_information_gain
        self.confidence_stability_threshold = confidence_stability_threshold

        # Track convergence history
        self.confidence_history = []
        self.information_gain_history = []
        self.answer_history = []  # Track recent answers for stability
        self.max_history_length = 5

    def should_terminate(self,
                        current_confidence_policy: float,
                        current_confidence_answer: float,
                        information_gain: float,
                        cycle_count: int,
                        max_cycles: int = 50,
                        current_answer: str = None,
                        answer_confidence: float = None) -> Tuple[bool, str]:
        """
        Determine if cognitive processing should terminate
        """

        # Always terminate after max cycles (safety net)
        if cycle_count >= max_cycles:
            return True, "Maximum cycles reached"

        # Update history
        self.confidence_history.append((current_confidence_policy, current_confidence_answer))
        self.information_gain_history.append(information_gain)
        
        # Track answer history if provided
        if current_answer is not None and answer_confidence is not None:
            self.answer_history.append((current_answer, answer_confidence))

        # Keep only recent history
        if len(self.confidence_history) > self.max_history_length:
            self.confidence_history.pop(0)
            self.information_gain_history.pop(0)
            if len(self.answer_history) > self.max_history_length:
                self.answer_history.pop(0)

        # Check answer stability (early termination for consistent high-confidence answers)
        if len(self.answer_history) >= 2:
            recent_answers = self.answer_history[-3:]  # Check last 3 answers
            
            # Check if last 2 answers are the same and have high confidence
            if len(recent_answers) >= 2:
                last_answer, last_confidence = recent_answers[-1]
                prev_answer, prev_confidence = recent_answers[-2]
                
                # Same answer with high confidence twice in a row
                if (last_answer == prev_answer and 
                    last_confidence > 0.8 and prev_confidence > 0.8):
                    return True, f"Answer stabilized: '{last_answer}' (confidence: {last_confidence:.2f})"

        # Check information gain convergence
        if len(self.information_gain_history) >= 3:
            recent_gains = self.information_gain_history[-3:]
            avg_recent_gain = np.mean(recent_gains)

            if avg_recent_gain < self.min_information_gain:
                return True, f"Information gain converged ({avg_recent_gain:.4f} < {self.min_information_gain})"

        # Check confidence stability
        if len(self.confidence_history) >= 3:
            recent_confidences = self.confidence_history[-3:]
            policy_stable = np.std([c[0] for c in recent_confidences]) < self.confidence_stability_threshold
            answer_stable = np.std([c[1] for c in recent_confidences]) < self.confidence_stability_threshold

            if policy_stable and answer_stable:
                avg_policy = np.mean([c[0] for c in recent_confidences])
                avg_answer = np.mean([c[1] for c in recent_confidences])

                # Only terminate if both confidences are reasonably high
                if avg_policy > 0.8 and avg_answer > 0.7:
                    return True, f"Confidence converged (policy: {avg_policy:.2f}, answer: {avg_answer:.2f})"

        return False, "Continue processing"

    def reset(self):
        """Reset convergence tracking for new problem"""
        self.confidence_history.clear()
        self.information_gain_history.clear()
        self.answer_history.clear()

class MemoryRecallSystem(nn.Module):
    """
    Similarity-driven memory recall system for Wiki-trained knowledge

    Features:
    - Confidence-gated recall bandwidth
    - Counterfactual memory activation
    - Self-adversarial memory checking
    - Episodic and semantic memory integration
    """

    def __init__(self, embedding_dim: int = 512, memory_capacity: int = 10000):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.memory_capacity = memory_capacity

        # Memory storage (simulated - in production would load from Wiki training)
        self.episodic_memories: List[MemoryTrace] = []
        self.semantic_memories: List[MemoryTrace] = []

        # Recall mechanisms
        self.similarity_encoder = nn.Sequential(
            nn.Linear(embedding_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128)
        )

        self.confidence_gate = nn.Sequential(
            nn.Linear(1, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )

        self.counterfactual_detector = nn.Sequential(
            nn.Linear(embedding_dim * 2, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

        # Initialize with simulated Wiki-trained knowledge
        self._initialize_simulated_wiki_memories()

    def _initialize_simulated_wiki_memories(self):
        """Initialize with simulated Wiki-trained knowledge (episodic and semantic)"""
        # Simulate episodic memories from Wiki training
        episodic_concepts = [
            ("group theory", "mathematical structures with binary operations"),
            ("abelian group", "commutative group where order doesn't matter"),
            ("cyclic group", "group generated by a single element"),
            ("normal subgroup", "subgroup invariant under conjugation"),
            ("kernel homomorphism", "set where homomorphism maps to identity"),
            ("ring structure", "set with addition and multiplication"),
            ("field definition", "ring where nonzero elements have multiplicative inverses"),
            ("vector space", "structure with scalars and vector addition"),
            ("linear transformation", "function preserving vector space structure"),
            ("eigenvalue problem", "finding scalars where matrix action scales vectors")
        ]

        for concept, description in episodic_concepts:
            embedding = torch.randn(self.embedding_dim) * 0.1
            memory = MemoryTrace(
                content=f"Concept: {concept} - {description}",
                embedding=embedding,
                relevance_score=0.5,
                activation_strength=0.3,
                source_type="wiki_episodic",
                timestamp=time.time(),
                confidence_weight=0.7
            )
            self.episodic_memories.append(memory)

        # Simulate semantic associations
        semantic_associations = [
            ("abelian", "commutative", "groups where ab = ba for all elements"),
            ("cyclic", "generated", "groups with single generator element"),
            ("normal", "invariant", "subgroups unchanged by conjugation"),
            ("kernel", "identity", "elements mapping to group identity"),
            ("field", "division", "rings with multiplicative inverses"),
            ("vector", "linear", "structures with scalar multiplication"),
            ("eigen", "characteristic", "values where matrix scales vectors")
        ]

        for term1, term2, association in semantic_associations:
            embedding = torch.randn(self.embedding_dim) * 0.1
            memory = MemoryTrace(
                content=f"Association: {term1} ↔ {term2} - {association}",
                embedding=embedding,
                relevance_score=0.4,
                activation_strength=0.2,
                source_type="wiki_semantic",
                timestamp=time.time(),
                confidence_weight=0.6
            )
            self.semantic_memories.append(memory)

    def recall_memories(self, query_embedding: torch.Tensor, current_confidence: float,
                       attention_focus: str = "balanced") -> Tuple[List[MemoryTrace], List[CounterfactualMemory]]:
        """
        Recall relevant memories based on similarity and confidence gating

        Args:
            query_embedding: Current reasoning state embedding
            current_confidence: Current inference confidence (0-1)
            attention_focus: 'balanced', 'exploratory', 'focused'

        Returns:
            Tuple of (relevant_memories, counterfactual_memories)
        """
        # Confidence-gated recall bandwidth
        confidence_tensor = torch.tensor([[current_confidence]], dtype=torch.float32)
        recall_bandwidth = self.confidence_gate(confidence_tensor).item()

        # Adjust bandwidth based on confidence and attention
        if current_confidence < 0.4:  # Low confidence - increase recall
            recall_bandwidth = min(1.0, recall_bandwidth * 1.5)
        elif current_confidence > 0.8:  # High confidence - focused recall
            recall_bandwidth = max(0.3, recall_bandwidth * 0.7)

        if attention_focus == "exploratory":
            recall_bandwidth = min(1.0, recall_bandwidth * 1.3)
        elif attention_focus == "focused":
            recall_bandwidth = max(0.2, recall_bandwidth * 0.6)

        # Calculate similarity scores for all memories
        relevant_memories = []
        all_memories = self.episodic_memories + self.semantic_memories

        for memory in all_memories:
            # Cosine similarity between query and memory
            similarity = F.cosine_similarity(
                query_embedding.unsqueeze(0),
                memory.embedding.unsqueeze(0)
            ).item()

            # Relevance weighting based on similarity and confidence
            relevance_score = similarity * recall_bandwidth * memory.confidence_weight

            if relevance_score > 0.01:  # Lowered activation threshold for testing
                activated_memory = MemoryTrace(
                    content=memory.content,
                    embedding=memory.embedding,
                    relevance_score=relevance_score,
                    activation_strength=min(1.0, relevance_score * 2.0),
                    source_type=memory.source_type,
                    timestamp=time.time(),
                    confidence_weight=memory.confidence_weight
                )
                relevant_memories.append(activated_memory)

        # Sort by relevance and limit results
        relevant_memories.sort(key=lambda x: x.relevance_score, reverse=True)
        relevant_memories = relevant_memories[:int(len(relevant_memories) * recall_bandwidth)]

        # Generate counterfactual memories for hypothesis testing
        counterfactual_memories = self._generate_counterfactual_memories(
            query_embedding, relevant_memories, current_confidence
        )

        return relevant_memories, counterfactual_memories

    def _generate_counterfactual_memories(self, query_embedding: torch.Tensor,
                                        relevant_memories: List[MemoryTrace],
                                        current_confidence: float) -> List[CounterfactualMemory]:
        """Generate contrasting memories for hypothesis stress-testing"""
        counterfactuals = []

        # Common mathematical distinctions to test
        distinction_pairs = [
            ("abelian", "cyclic", "Abelian groups are commutative, cyclic groups are generated by one element"),
            ("normal", "subgroup", "Normal subgroups are invariant under conjugation"),
            ("field", "ring", "Fields have multiplicative inverses, rings may not"),
            ("kernel", "image", "Kernel maps to identity, image is homomorphism range"),
            ("eigenvalue", "eigenvector", "Eigenvalues are scalars, eigenvectors are vectors"),
            ("linear", "affine", "Linear maps preserve origin, affine may translate"),
            ("group", "monoid", "Groups have inverses, monoids may not"),
            ("ring", "semiring", "Rings have additive inverses, semirings may not")
        ]

        for memory in relevant_memories:
            if memory.relevance_score > 0.3:  # Only for highly relevant memories
                memory_text = memory.content.lower()

                # Check if memory contains concepts that have common confusions
                for concept1, concept2, distinction in distinction_pairs:
                    if concept1 in memory_text and concept2 not in memory_text:
                        # Create counterfactual by finding contrasting memory
                        contrasting_memory = self._find_contrasting_memory(concept2, memory)

                        if contrasting_memory:
                            counterfactual = CounterfactualMemory(
                                original_concept=concept1,
                                contrasting_concept=concept2,
                                distinction_explanation=distinction,
                                memory_trace=contrasting_memory,
                                stress_test_score=self._calculate_stress_test_score(
                                    query_embedding, memory, contrasting_memory
                                )
                            )
                            counterfactuals.append(counterfactual)

        return counterfactuals

    def _find_contrasting_memory(self, concept: str, original_memory: MemoryTrace) -> Optional[MemoryTrace]:
        """Find a memory that contrasts with the original concept"""
        all_memories = self.episodic_memories + self.semantic_memories

        for memory in all_memories:
            if concept.lower() in memory.content.lower():
                # Ensure it's different from original
                if memory.content != original_memory.content:
                    return memory

        return None

    def _calculate_stress_test_score(self, query_embedding: torch.Tensor,
                                   original: MemoryTrace, contrasting: MemoryTrace) -> float:
        """Calculate how much this counterfactual challenges the current hypothesis"""
        # Use embedding similarity to measure potential conflict
        similarity = F.cosine_similarity(
            original.embedding.unsqueeze(0),
            contrasting.embedding.unsqueeze(0)
        ).item()

        # Lower similarity = higher stress test potential
        stress_score = 1.0 - similarity

        # Weight by relevance to query
        query_similarity = F.cosine_similarity(
            query_embedding.unsqueeze(0),
            contrasting.embedding.unsqueeze(0)
        ).item()

        return stress_score * query_similarity

    def self_adversarial_check(self, hypothesis: str, recalled_memories: List[MemoryTrace]) -> Dict[str, Any]:
        """
        Self-adversarial memory checking for contradictions and overgeneralizations

        Returns critique of the hypothesis based on recalled memories
        """
        contradictions = []
        overgeneralizations = []
        supporting_evidence = []

        hypothesis_lower = hypothesis.lower()

        for memory in recalled_memories:
            memory_lower = memory.content.lower()

            # Check for direct contradictions
            if any(phrase in memory_lower for phrase in ['not always', 'does not imply', 'counterexample']):
                if any(word in hypothesis_lower for word in memory_lower.split()):
                    contradictions.append({
                        'memory': memory.content,
                        'conflict_type': 'direct_contradiction',
                        'strength': memory.relevance_score
                    })

            # Check for overgeneralization warnings
            if any(phrase in memory_lower for phrase in ['in general', 'typically', 'usually', 'not necessarily']):
                overgeneralizations.append({
                    'memory': memory.content,
                    'warning_type': 'overgeneralization_risk',
                    'strength': memory.relevance_score
                })

            # Check for supporting evidence
            if any(word in hypothesis_lower for word in memory_lower.split() if len(word) > 3):
                supporting_evidence.append({
                    'memory': memory.content,
                    'support_type': 'semantic_overlap',
                    'strength': memory.relevance_score
                })

        return {
            'contradictions': contradictions,
            'overgeneralizations': overgeneralizations,
            'supporting_evidence': supporting_evidence,
            'overall_critique': self._synthesize_critique(contradictions, overgeneralizations, supporting_evidence)
        }

    def _synthesize_critique(self, contradictions: List[Dict], overgeneralizations: List[Dict],
                           supporting: List[Dict]) -> str:
        """Synthesize overall critique from memory analysis"""
        if not contradictions and not overgeneralizations and not supporting:
            return "Memory recall shows no strong conflicts or support for this hypothesis."

        critique_parts = []

        if contradictions:
            strongest = max(contradictions, key=lambda x: x['strength'])
            critique_parts.append(f"Memory evidence suggests potential contradiction: {strongest['memory'][:100]}...")

        if overgeneralizations:
            strongest = max(overgeneralizations, key=lambda x: x['strength'])
            critique_parts.append(f"Risk of overgeneralization detected: {strongest['memory'][:100]}...")

        if supporting:
            strongest = max(supporting, key=lambda x: x['strength'])
            critique_parts.append(f"Memory provides supporting context: {strongest['memory'][:100]}...")

        return " ".join(critique_parts)

class ReflectionAction(Enum):
    """DMN Reflector action decisions"""
    CONTINUE = "CONTINUE"
    REFINE = "REFINE" 
    REPLAN = "REPLAN"
    EXECUTE = "EXECUTE"

@dataclass
class ReflectionOutput:
    """Structured output from DMN Reflector"""
    critique: str
    alternative: str
    confidence: float
    action: ReflectionAction
    reflection_vector: torch.Tensor

class DMNReflector(nn.Module):
    """
    DMN Reflector using Qwen2-0.5B-Instruct for metacognitive analysis
    
    This module analyzes GLW states and generates structured meta-commentary
    to enable System-2 deliberate reasoning.
    """
    
    def __init__(self, glw_dim: int = 512, reflection_dim: int = 128):
        super().__init__()
        self.glw_dim = glw_dim
        self.reflection_dim = reflection_dim
        
        # GLW state encoder for reflection context
        self.glw_encoder = nn.Sequential(
            nn.Linear(glw_dim, 256),
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(256, reflection_dim)
        )
        
        # Text embedding for reflection output
        self.text_embedding = nn.Linear(768, reflection_dim)  # Assume 768-dim text embeddings
        
        # Confidence predictor
        self.confidence_head = nn.Sequential(
            nn.Linear(reflection_dim * 2, 64),  # GLW + text embedding
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )
        
        # System prompt for Qwen2-0.5B
        self.system_prompt = """You are a metacognitive critic analyzing reasoning processes. 
Your task is to examine the current thought state and provide structured feedback.

Output format (exactly as shown):
CRITIQUE: [Identify potential errors or weak reasoning steps]
ALTERNATIVE: [Propose a different approach or perspective] 
CONFIDENCE: [0.0-1.0 numerical score]
ACTION: [Choose: CONTINUE, REFINE, REPLAN, or EXECUTE]

Be concise but thorough. Focus on logical consistency and reasoning quality."""
        
        print("[BRAIN] DMN Reflector initialized")
        print(f"   GLW dimension: {glw_dim}")
        print(f"   Reflection dimension: {reflection_dim}")
    
    def _simulate_qwen_reflection(self, glw_state: torch.Tensor, task_context: str) -> str:
        """
        Simulate Qwen2-0.5B reflection output with better dynamics
        In production, this would interface with llama.cpp Q4_K_M model
        """
        # Extract GLW statistics for context
        glw_mean = glw_state.mean().item()
        glw_std = glw_state.std().item()
        glw_max = glw_state.max().item()
        glw_min = glw_state.min().item()
        
        # Analyze task context for better reflection
        task_lower = task_context.lower()
        is_math = any(word in task_lower for word in ['calculate', 'compound', 'interest', 'rate', 'math'])
        is_complex = any(word in task_lower for word in ['strategy', 'comprehensive', 'design', 'plan'])
        is_simple = any(word in task_lower for word in ['capital', 'what is', 'simple'])
        
        # Dynamic reflection based on GLW evolution and task type
        if glw_std > 0.15:  # High variance = good exploration
            if is_math:
                critique = "Mathematical reasoning pathway shows good activation diversity. Formula application appears on track."
                alternative = "Consider double-checking calculation steps for precision."
                confidence = 0.85
                action = "EXECUTE"
            elif is_complex:
                critique = "Complex reasoning showing good multi-faceted analysis. Multiple solution paths active."
                alternative = "Could benefit from prioritizing implementation feasibility."
                confidence = 0.75
                action = "REFINE"
            else:
                critique = "Reasoning shows healthy activation patterns with good information integration."
                alternative = "Current approach is solid, minor verification recommended."
                confidence = 0.80
                action = "EXECUTE"
        elif glw_std > 0.05:  # Medium variance = developing understanding
            if glw_mean > 0.3:
                critique = "Moderate confidence in reasoning path. Some uncertainty remains in approach."
                alternative = "Gather more specific information or break down into smaller steps."
                confidence = 0.65
                action = "REFINE"
            else:
                critique = "Reasoning developing but needs more focused activation on key concepts."
                alternative = "Re-examine core problem requirements and strengthen relevant knowledge activation."
                confidence = 0.45
                action = "REFINE"
        else:  # Low variance = need more exploration
            if glw_mean < 0.1:  # Very low activation
                critique = "Minimal reasoning activation detected. Problem analysis insufficient."
                alternative = "Start with fundamental concept activation and build systematic understanding."
                confidence = 0.25
                action = "REPLAN"
            else:  # Some activation but low variance
                critique = "Reasoning too narrow or focused. Missing broader context integration."
                alternative = "Explore alternative approaches and activate related knowledge domains."
                confidence = 0.40
                action = "REFINE"
        
        # Boost confidence for appropriate tasks
        if is_simple and confidence < 0.9:
            confidence = min(0.95, confidence + 0.4)  # Bigger boost for simple tasks
            action = "EXECUTE"
            critique = "Simple factual query with clear answer pathway activated."
        elif is_math and confidence > 0.7:
            confidence = min(0.9, confidence + 0.15)  # Boost math confidence
            action = "EXECUTE"
        
        # Format as structured output
        reflection_text = f"""CRITIQUE: {critique}
ALTERNATIVE: {alternative}
CONFIDENCE: {confidence:.2f}
ACTION: {action}"""
        
        return reflection_text
    
    def _parse_reflection_output(self, reflection_text: str) -> Dict[str, Any]:
        """Parse structured reflection output from Qwen2-0.5B"""
        try:
            # Extract components using regex
            critique_match = re.search(r'CRITIQUE:\s*(.*?)(?=\nALTERNATIVE:|$)', reflection_text, re.DOTALL)
            alternative_match = re.search(r'ALTERNATIVE:\s*(.*?)(?=\nCONFIDENCE:|$)', reflection_text, re.DOTALL)
            confidence_match = re.search(r'CONFIDENCE:\s*([0-9.]+)', reflection_text)
            action_match = re.search(r'ACTION:\s*(CONTINUE|REFINE|REPLAN|EXECUTE)', reflection_text)
            
            critique = critique_match.group(1).strip() if critique_match else "No critique provided"
            alternative = alternative_match.group(1).strip() if alternative_match else "No alternative provided"
            confidence = float(confidence_match.group(1)) if confidence_match else 0.5
            action_str = action_match.group(1) if action_match else "CONTINUE"
            
            # Clamp confidence to [0, 1]
            confidence = max(0.0, min(1.0, confidence))
            
            return {
                'critique': critique,
                'alternative': alternative,
                'confidence': confidence,
                'action': ReflectionAction(action_str)
            }
        except Exception as e:
            print(f"[WARN] Error parsing reflection output: {e}")
            return {
                'critique': "Parse error in reflection",
                'alternative': "Fallback to default processing",
                'confidence': 0.5,
                'action': ReflectionAction.CONTINUE
            }
    
    def forward(self, glw_state: torch.Tensor, task_context: str = "") -> ReflectionOutput:
        """
        Perform DMN reflection on current GLW state
        
        Args:
            glw_state: Current Global Latent Workspace state [1, glw_dim]
            task_context: Optional task description for context
            
        Returns:
            ReflectionOutput with structured feedback
        """
        batch_size = glw_state.shape[0]
        
        # Encode GLW state for reflection
        glw_encoded = self.glw_encoder(glw_state)  # [1, reflection_dim]
        
        # Generate reflection using simulated Qwen2-0.5B
        # In production: interface with llama.cpp here
        reflection_text = self._simulate_qwen_reflection(glw_state[0], task_context)
        
        # Parse structured output
        parsed = self._parse_reflection_output(reflection_text)
        
        # Create text embedding (simulated)
        # In production: use actual text encoder
        text_hash = hash(reflection_text) % 1000000
        text_embedding = torch.randn(1, self.reflection_dim) * 0.1
        text_embedding += torch.tensor([[text_hash / 1000000.0] * self.reflection_dim])
        
        # Combine GLW and text embeddings
        combined = torch.cat([glw_encoded, text_embedding], dim=1)
        
        # Predict confidence (neural network validation of LLM confidence)
        predicted_confidence = self.confidence_head(combined).item()
        
        # Use maximum of LLM confidence and neural predictor (trust the higher one)
        final_confidence = max(parsed['confidence'], predicted_confidence * 0.5 + 0.5)  # Boost neural predictor
        
        # Create reflection vector for GLW injection
        reflection_vector = glw_encoded[0] * final_confidence
        
        return ReflectionOutput(
            critique=parsed['critique'],
            alternative=parsed['alternative'],
            confidence=final_confidence,
            action=parsed['action'],
            reflection_vector=reflection_vector
        )

class HIPECAExecutiveGating(nn.Module):
    """
    HIPECA Executive Gating with ACC Error Monitor and Cerebellar Timing
    
    Implements state-guard transactions and confidence-based execution control
    """
    
    def __init__(self, glw_dim: int = 512):
        super().__init__()
        self.glw_dim = glw_dim
        
        # ACC Error Monitor
        self.acc_monitor = nn.Sequential(
            nn.Linear(glw_dim + 1, 128),  # GLW + confidence
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )
        
        # Cerebellar Timing Predictor
        self.cerebellar_timing = nn.Sequential(
            nn.Linear(glw_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(), 
            nn.Linear(32, 1)  # Predicted reflection cycles needed
        )
        
        # State snapshots for rollback
        self.state_snapshots = []
        self.max_snapshots = 5
        
        # Thresholds
        self.confidence_threshold = 0.6
        self.error_threshold = 0.7
        
        print("[TARGET] HIPECA Executive Gating initialized")
    
    def create_snapshot(self, glw_state: torch.Tensor, metadata: Dict[str, Any]):
        """Create state snapshot for potential rollback"""
        snapshot = {
            'glw_state': glw_state.clone(),
            'metadata': metadata.copy(),
            'timestamp': time.time()
        }
        
        self.state_snapshots.append(snapshot)
        
        # Limit snapshot buffer
        if len(self.state_snapshots) > self.max_snapshots:
            self.state_snapshots.pop(0)
    
    def rollback_to_snapshot(self, steps_back: int = 1) -> Optional[Dict[str, Any]]:
        """Rollback to previous state snapshot"""
        if len(self.state_snapshots) >= steps_back:
            snapshot = self.state_snapshots[-(steps_back)]
            print(f"[CYCLE] Rolling back {steps_back} steps")
            return snapshot
        return None
    
    def predict_reflection_cycles(self, glw_state: torch.Tensor, task_type: str = "general") -> int:
        """Predict optimal number of reflection cycles for task"""
        predicted_cycles = self.cerebellar_timing(glw_state).item()
        
        # Task-specific adjustments
        base_cycles = max(1, int(predicted_cycles))
        
        if "math" in task_type.lower():
            return min(base_cycles, 2)  # Math needs fewer cycles
        elif "novel" in task_type.lower() or "complex" in task_type.lower():
            return min(base_cycles + 2, 5)  # Complex problems need more
        else:
            return min(base_cycles, 3)  # General cap at 3
    
    def should_force_refinement(self, glw_state: torch.Tensor, confidence: float) -> Tuple[bool, str]:
        """Check if ACC should force system into refinement mode"""
        
        # Confidence check
        if confidence < self.confidence_threshold:
            return True, f"Low confidence: {confidence:.2f} < {self.confidence_threshold}"
        
        # Error monitor check
        confidence_tensor = torch.tensor([[confidence]])
        combined_input = torch.cat([glw_state, confidence_tensor], dim=1)
        error_prediction = self.acc_monitor(combined_input).item()
        
        if error_prediction > self.error_threshold:
            return True, f"High error prediction: {error_prediction:.2f} > {self.error_threshold}"
        
        return False, "Executive clearance granted"
    
    def forward(self, glw_state: torch.Tensor, confidence: float, 
                task_type: str = "general") -> Dict[str, Any]:
        """Execute HIPECA gating logic"""
        
        # Predict reflection cycles needed
        predicted_cycles = self.predict_reflection_cycles(glw_state, task_type)
        
        # Check if refinement is needed
        force_refine, reason = self.should_force_refinement(glw_state, confidence)
        
        return {
            'predicted_reflection_cycles': predicted_cycles,
            'force_refinement': force_refine,
            'refinement_reason': reason,
            'executive_clearance': not force_refine
        }

class RecursiveGLWManager(nn.Module):
    """
    Manages the Recursive GLW Update Rule with reflection injection
    
    Implements: GLW_t+1 = alpha·GLW_t + beta·Sensory + gamma·DMN_reflection + delta·Action_feedback
    """
    
    def __init__(self, glw_dim: int = 512, reflection_dim: int = 128):
        super().__init__()
        self.glw_dim = glw_dim
        self.reflection_dim = reflection_dim
        
        # Learnable mixing coefficients
        self.alpha = nn.Parameter(torch.tensor(0.7))  # GLW persistence
        self.beta = nn.Parameter(torch.tensor(0.2))   # Sensory weight
        self.gamma = nn.Parameter(torch.tensor(0.1))  # Reflection weight
        self.delta = nn.Parameter(torch.tensor(0.05)) # Action feedback weight
        
        # Projection layers
        self.reflection_projector = nn.Linear(reflection_dim, glw_dim)
        self.action_projector = nn.Linear(64, glw_dim)  # Assume 64-dim action feedback
        
        # Normalization
        self.layer_norm = nn.LayerNorm(glw_dim)
        
        print(f"[CYCLE] Recursive GLW Manager initialized")
        print(f"   GLW dimension: {glw_dim}")
        print(f"   Initial mixing: alpha={self.alpha.item():.2f}, beta={self.beta.item():.2f}, "
              f"gamma={self.gamma.item():.2f}, delta={self.delta.item():.2f}")
    
    def update_glw(self, 
                   current_glw: torch.Tensor,
                   sensory_input: torch.Tensor,
                   reflection_vector: torch.Tensor,
                   action_feedback: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Execute recursive GLW update with reflection injection
        """
        batch_size = current_glw.shape[0]
        
        # Project reflection to GLW space
        reflection_projected = self.reflection_projector(reflection_vector.unsqueeze(0))
        
        # Handle action feedback
        if action_feedback is None:
            action_feedback = torch.zeros(batch_size, 64)
        action_projected = self.action_projector(action_feedback)
        
        # Ensure sensory input matches GLW dimension
        if sensory_input.shape[1] != self.glw_dim:
            # Simple projection if dimensions don't match
            sensory_projected = F.linear(sensory_input, 
                                       torch.randn(self.glw_dim, sensory_input.shape[1]) * 0.1)
        else:
            sensory_projected = sensory_input
        
        # Normalize coefficients to sum <= 1.0
        total_weight = torch.abs(self.alpha) + torch.abs(self.beta) + torch.abs(self.gamma) + torch.abs(self.delta)
        norm_alpha = self.alpha / total_weight
        norm_beta = self.beta / total_weight
        norm_gamma = self.gamma / total_weight
        norm_delta = self.delta / total_weight
        
        # Recursive update rule
        new_glw = (norm_alpha * current_glw + 
                   norm_beta * sensory_projected +
                   norm_gamma * reflection_projected +
                   norm_delta * action_projected)
        
        # Add small positive bias to prevent zero-mean collapse
        new_glw = new_glw + 0.01
        
        # Layer normalization for stability (but preserve some mean)
        normalized_glw = self.layer_norm(new_glw)
        
        # Blend to preserve some original magnitude
        final_glw = 0.8 * normalized_glw + 0.2 * new_glw
        
        return final_glw

class BIOMINDRecursiveDMN(nn.Module):
    """
    Complete BIOMIND Cognitive Control Architecture

    Implements sophisticated cognitive control with:
    - Dual confidence signals (policy vs answer)
    - State-changing DMN iterations
    - MDN planning over cognitive actions
    - Basal ganglia-style post-execution evaluation
    - Counterfactual re-decision when confidence unstable
    - Convergence-based termination criteria
    """

    def __init__(self,
                 specialists: List[str],
                 glw_dim: int = 512,
                 reflection_dim: int = 128,
                 max_reflection_cycles: int = 50,
                 use_real_models: bool = True,
                 use_trained_router: bool = True):
        super().__init__()

        self.specialists = specialists
        self.glw_dim = glw_dim
        self.reflection_dim = reflection_dim
        self.max_reflection_cycles = max_reflection_cycles
        self.use_real_models = use_real_models
        self.use_trained_router = use_trained_router

        # Initialize REAL model inference system
        self.real_specialist_system = RealSpecialistSystem(use_real_models=use_real_models)

        # Load trained thalamic router if available
        self.trained_router = None
        if use_trained_router:
            try:
                from train_simple_thalamic import SimpleThalamicRouter
                checkpoint_path = Path(__file__).parent / "trained_router_checkpoint.pth"
                if checkpoint_path.exists():
                    print("[CYCLE] Loading trained thalamic router checkpoint...")
                    self.trained_router = SimpleThalamicRouter(specialists)
                    checkpoint = torch.load(checkpoint_path, map_location='cpu')
                    # Handle nested checkpoint structure
                    if 'model_state_dict' in checkpoint:
                        state_dict = checkpoint['model_state_dict']
                    else:
                        state_dict = checkpoint
                    self.trained_router.load_state_dict(state_dict)
                    self.trained_router.eval()
                    print("[OK] Trained thalamic router loaded successfully!")
                else:
                    print("[WARN]  Trained router checkpoint not found, using keyword-based routing")
            except Exception as e:
                print(f"[WARN]  Failed to load trained router: {e}, using keyword-based routing")

        # Core cognitive control components
        self.dmn_reflector = DMNReflector(glw_dim, reflection_dim)
        self.executive_gating = HIPECAExecutiveGating(glw_dim)
        self.glw_manager = RecursiveGLWManager(glw_dim, reflection_dim)

        # New cognitive control components
        self.mdn_planner = MDNPlanner(glw_dim)
        self.basal_ganglia = BasalGangliaEvaluator(glw_dim)
        self.counterfactual_redecision = CounterfactualReDecision(glw_dim)
        self.termination_criteria = TerminationCriteria()

        # Memory integration components
        self.memory_system = MemoryRecallSystem(embedding_dim=glw_dim)

        # Specialist dispatcher (fallback if trained router not available)
        if not self.trained_router:
            self.specialist_dispatcher = nn.Sequential(
                nn.Linear(glw_dim, 256),
                nn.ReLU(),
                nn.Dropout(0.1),
                nn.Linear(256, len(specialists)),
                nn.Softmax(dim=-1)
            )

        # Initialize GLW state
        self.register_buffer('current_glw', torch.randn(1, glw_dim) * 0.1)

        print("[BRAIN] BIOMIND Cognitive Control Architecture initialized")
        print(f"   Specialists: {specialists}")
        print(f"   GLW dimension: {glw_dim}")
        print(f"   Real models: {'YES' if use_real_models else 'NO'}")
        print(f"   Trained router: {'YES' if self.trained_router else 'NO'}")
        print("   New components: MDN Planner, Basal Ganglia Evaluator, Counterfactual Re-Decision")

    def recursive_reflection_cycle(self,
                                   task_input: str,
                                   task_type: str = "general",
                                   max_cycles: Optional[int] = None) -> Dict[str, Any]:
        """
        Execute sophisticated cognitive control loop with state-changing iterations

        Implements true MDN planning, dual confidence signals, post-execution evaluation,
        and convergence-based termination.
        """

        if max_cycles is None:
            max_cycles = self.max_reflection_cycles

        print("\n[CYCLE] Starting BIOMIND Cognitive Control Loop")
        print(f"   Task: {task_input[:100]}...")
        print(f"   Type: {task_type}")

        # Initialize cognitive state
        # Use trained router for initial specialist selection
        initial_specialist = "general"
        if self.trained_router:
            try:
                with torch.no_grad():
                    router_result = self.trained_router(task_input)
                    initial_specialist = router_result['predicted_specialist']
                    
                    # Override obviously wrong predictions with keyword-based routing
                    task_lower = task_input.lower()
                    is_math_question = ("math" in task_lower or "calculate" in task_lower or "algebra" in task_lower or 
                                      "field" in task_lower or "extension" in task_lower or "group" in task_lower or 
                                      "ring" in task_lower or "polynomial" in task_lower or "equation" in task_lower or
                                      "sqrt" in task_lower or "degree" in task_lower or "finite" in task_lower or
                                      "s_" in task_lower or "symmetric" in task_lower or "permutation" in task_lower or
                                      "index" in task_lower or "order" in task_lower or "subgroup" in task_lower or
                                      "<" in task_lower or ">" in task_lower or "z_" in task_lower or "c_" in task_lower)
                    
                    if is_math_question and initial_specialist != "qwen_math_expert":
                        print(f"   [CYCLE] Overriding trained router -> qwen_math_expert for math question")
                        initial_specialist = "qwen_math_expert"
                    
                    print(f"   Initial routing: {initial_specialist} (trained router)")
            except Exception as e:
                print(f"   Initial routing failed: {e}, using general")
        
        cognitive_state = CognitiveState(
            question_representation=task_input,
            specialist_role=initial_specialist
        )

        # Initialize sensory input
        sensory_input = self._initialize_sensory_input(task_input)

        # Initialize GLW with meaningful content
        self.current_glw = sensory_input.clone()

        # Reset termination criteria
        self.termination_criteria.reset()

        # Control loop variables
        cycle = 0
        last_executed_answer = None
        last_execution_confidence_pre = 0.0
        last_execution_confidence_post = 0.0
        information_gain = 0.0

        reflection_history = []
        execution_history = []
        specialist_responses = []  # Track responses from different specialists

        # Enhanced task input that accumulates context
        current_task_input = task_input

        while cycle < max_cycles:
            cycle += 1
            print(f"\n[CYCLE] Cycle {cycle}: Cognitive Control Loop")

            # Step 1: Memory recall with current cognitive state
            recalled_memories, counterfactual_memories = self.memory_system.recall_memories(
                self.current_glw.squeeze(0),
                cognitive_state.confidence_policy,
                self._attention_focus_from_state(cognitive_state)
            )

            print(f"   Context: {len(recalled_memories)} memories, {len(cognitive_state.reasoning_history)} previous reasoning traces")
            if cognitive_state.reasoning_history:
                print(f"   Prior reasoning: {len(cognitive_state.reasoning_history)} traces available")

            memory_context = {
                'recalled_memories': recalled_memories,
                'counterfactual_memories': counterfactual_memories,
                'num_memories': len(recalled_memories),
                'num_counterfactuals': len(counterfactual_memories)
            }

            # Step 2: MDN Planning - select cognitive action
            print(f"   [MDN] MDN Planning: Selecting specialist and action...")
            action_plan = self.mdn_planner(
                self.current_glw,
                cognitive_state,
                memory_context,
                task_input
            )
            print(f"   [MDN] MDN Decision: {action_plan.selected_action.value} with {cognitive_state.specialist_role}")

            # Step 3: Execute selected cognitive action
            print(f"   [BRAIN] PFC Execution: Generating specialist response...")
            state_changed = self._execute_cognitive_action(
                action_plan.selected_action,
                cognitive_state,
                memory_context,
                current_task_input,  # Use enhanced task input
                specialist_responses  # Pass specialist responses for context
            )

            # Step 4: Update GLW with reflection
            reflection_output = self.dmn_reflector(self.current_glw, task_input)
            old_glw = self.current_glw.clone()

            self.current_glw = self.glw_manager.update_glw(
                self.current_glw,
                sensory_input,
                reflection_output.reflection_vector
            )

            # Step 5: Update confidence signals
            old_policy_confidence = cognitive_state.confidence_policy
            old_answer_confidence = cognitive_state.confidence_answer

            cognitive_state.confidence_policy = self._update_policy_confidence(
                cognitive_state, action_plan, reflection_output, memory_context
            )

            cognitive_state.confidence_answer = self._update_answer_confidence(
                cognitive_state, reflection_output, memory_context
            )

            # Calculate information gain
            confidence_change = (cognitive_state.confidence_policy - old_policy_confidence) + \
                              (cognitive_state.confidence_answer - old_answer_confidence)
            information_gain = abs(confidence_change) + action_plan.epistemic_value

            print(f"   [TARGET] Basal Ganglia: Evaluating response quality...")
            print(f"      Confidence - Policy: {cognitive_state.confidence_policy:.2f}, Answer: {cognitive_state.confidence_answer:.2f}")

            # Step 6: Check termination criteria
            current_answer = generated_answer if 'generated_answer' in locals() else None
            current_answer_confidence = execution_confidence_post if 'execution_confidence_post' in locals() else None
            
            should_terminate, termination_reason = self.termination_criteria.should_terminate(
                cognitive_state.confidence_policy,
                cognitive_state.confidence_answer,
                information_gain,
                cycle,
                max_cycles,
                current_answer,
                current_answer_confidence
            )

            # Step 7: Handle EXECUTE action
            if action_plan.selected_action == MDNAction.EXECUTE_AND_EVALUATE:
                # Extract question for logging
                task_lines = current_task_input.split('\n')
                question = ""
                for line in task_lines:
                    if line.startswith('Question:'):
                        question = line.replace('Question:', '').strip()
                        break
                if not question:
                    question = current_task_input[:100]

                print(f"   ? Specialist Execution: {cognitive_state.specialist_role}")
                print(f"      Input: '{question[:60]}...'")
                print(f"      Context: {len(recalled_memories)} recalled memories")

                # Generate answer (simplified - would use specialist dispatch)
                generated_answer = self._generate_answer(current_task_input, cognitive_state, recalled_memories, action_plan)
                execution_confidence_pre = cognitive_state.confidence_answer

                print(f"      Response: '{generated_answer}'")

                # Simulate post-execution confidence (would be from actual execution)
                execution_confidence_post = self._simulate_execution_confidence(generated_answer, task_input)

                print(f"   [BRAIN] PFC Reintegration: Specialist knowledge integrated")
                print(f"      Pre-execution confidence: {execution_confidence_pre:.2f}")
                print(f"      Post-execution confidence: {execution_confidence_post:.2f}")

                cognitive_state.execution_confidence_pre = execution_confidence_pre
                cognitive_state.execution_confidence_post = execution_confidence_post

                # Basal Ganglia Evaluation
                evaluation = self.basal_ganglia.evaluate_post_execution(
                    self.current_glw,
                    execution_confidence_pre,
                    execution_confidence_post,
                    generated_answer,
                    recalled_memories,
                    task_input
                )

                if evaluation['should_reopen_cognition']:
                    print(f"   [CYCLE] Counterfactual: Reopening cognition - {evaluation['reopen_reason']}")
                    # Reset confidences and continue
                    cognitive_state.confidence_answer *= 0.7
                    cognitive_state.confidence_policy *= 0.8
                    should_terminate = False
                    termination_reason = "Basal ganglia reopened cognition"
                else:
                    print(f"   [OK] Answer committed to working memory")
                    last_executed_answer = generated_answer
                    last_execution_confidence_pre = execution_confidence_pre
                    last_execution_confidence_post = execution_confidence_post

                execution_history.append({
                    'cycle': cycle,
                    'answer': generated_answer,
                    'confidence_pre': execution_confidence_pre,
                    'confidence_post': execution_confidence_post,
                    'basal_ganglia_evaluation': evaluation,
                    'specialist': cognitive_state.specialist_role
                })

                # Track specialist responses for context in future iterations
                specialist_responses.append({
                    'cycle': cycle,
                    'specialist': cognitive_state.specialist_role,
                    'answer': generated_answer,
                    'confidence': execution_confidence_post
                })

            # Step 8: Counterfactual re-decision if confidence unstable
            if (cognitive_state.confidence_answer < 0.5 and
                len(execution_history) > 0 and
                execution_history[-1]['basal_ganglia_evaluation']['should_reopen_cognition']):

                print(f"   [CYCLE] Counterfactual: Generating alternative hypotheses...")
                alternatives = self.counterfactual_redecision.generate_alternatives(
                    self.current_glw,
                    last_executed_answer,
                    cognitive_state.confidence_answer,
                    cognitive_state,
                    recalled_memories
                )

                if alternatives:
                    best_alternative = alternatives[0]
                    print(f"      Best Alternative: '{best_alternative['answer'][:50]}...' (conf: {best_alternative['confidence']:.2f})")
                    # Update answer with best alternative
                    last_executed_answer = best_alternative['answer']
                    cognitive_state.confidence_answer = best_alternative['confidence']

            # Record reflection history
            reflection_history.append({
                'cycle': cycle,
                'cognitive_state': {
                    'memory_radius': cognitive_state.memory_recall_radius,
                    'abstraction_level': cognitive_state.abstraction_level,
                    'specialist_role': cognitive_state.specialist_role,
                    'confidence_policy': cognitive_state.confidence_policy,
                    'confidence_answer': cognitive_state.confidence_answer
                },
                'mdn_action': action_plan.selected_action.value,
                'expected_gain': action_plan.expected_confidence_gain,
                'epistemic_value': action_plan.epistemic_value,
                'information_gain': information_gain,
                'memory_recall': memory_context,
                'state_changed': state_changed,
                'termination_check': termination_reason
            })

            # Check termination
            if should_terminate:
                print(f"   [OK] Termination: {termination_reason}")
                break

            print(f"   ? Continuing to cycle {cycle + 1}...")

        # Final specialist - use the one that was actually used during execution
        selected_specialist = cognitive_state.specialist_role
        specialist_confidence = cognitive_state.confidence_answer  # Use answer confidence as specialist confidence

        print("\n   [TARGET] Final Results")
        print(f"   Specialist: {selected_specialist}")
        print(f"   Final Answer: {last_executed_answer[:100] if last_executed_answer else 'No answer generated'}...")
        print(f"   Policy Confidence: {cognitive_state.confidence_policy:.2f}")
        print(f"   Answer Confidence: {cognitive_state.confidence_answer:.2f}")
        print(f"   Total Cycles: {cycle}")

        return {
            'selected_specialist': selected_specialist,
            'specialist_confidence': specialist_confidence,
            'final_answer': last_executed_answer,
            'cognitive_state': cognitive_state,
            'reflection_cycles': cycle,
            'reflection_history': reflection_history,
            'execution_history': execution_history,
            'final_glw_state': self.current_glw.clone(),
            'termination_reason': termination_reason,
            'executive_summary': {
                'total_cycles': cycle,
                'final_policy_confidence': cognitive_state.confidence_policy,
                'final_answer_confidence': cognitive_state.confidence_answer,
                'glw_variance': self.current_glw.var().item(),
                'total_executions': len(execution_history),
                'successful_commits': sum(1 for e in execution_history
                                        if not e['basal_ganglia_evaluation']['should_reopen_cognition'])
            }
        }

    def _initialize_sensory_input(self, task_input: str) -> torch.Tensor:
        """Initialize sensory input based on task content"""
        task_words = task_input.lower().split()

        # Create content-based activation patterns
        sensory_input = torch.zeros(1, self.glw_dim)

        # Math-related activation
        math_words = ['calculate', 'compound', 'interest', 'rate', 'formula', 'equation', 'solve']
        math_activation = sum(1 for word in task_words if word in math_words) / len(task_words)

        # Complex reasoning activation
        complex_words = ['strategy', 'comprehensive', 'design', 'plan', 'analyze', 'evaluate']
        complex_activation = sum(1 for word in task_words if word in complex_words) / len(task_words)

        # Simple factual activation
        simple_words = ['what', 'is', 'capital', 'who', 'when', 'where']
        simple_activation = sum(1 for word in task_words if word in simple_words) / len(task_words)

        # Build meaningful GLW initialization
        for i in range(self.glw_dim):
            if i < self.glw_dim // 4:  # Math region
                sensory_input[0, i] = math_activation + np.random.normal(0, 0.1)
            elif i < self.glw_dim // 2:  # Complex reasoning region
                sensory_input[0, i] = complex_activation + np.random.normal(0, 0.1)
            elif i < 3 * self.glw_dim // 4:  # Simple factual region
                sensory_input[0, i] = simple_activation + np.random.normal(0, 0.1)
            else:  # General activation region
                sensory_input[0, i] = 0.3 + np.random.normal(0, 0.15)

        return sensory_input

    def _attention_focus_from_state(self, cognitive_state: CognitiveState) -> str:
        """Determine attention focus from cognitive state"""
        if cognitive_state.attention_distribution is not None:
            # Simplified - would analyze attention distribution
            return "balanced"
        return "balanced"

    def _execute_cognitive_action(self, action: MDNAction, cognitive_state: CognitiveState,
                                 memory_context: Dict[str, Any], task_input: str,
                                 specialist_responses: List[Dict[str, Any]] = None) -> bool:
        """Execute the selected cognitive action and modify state"""
        if specialist_responses is None:
            specialist_responses = []

        state_changed = False

        if action == MDNAction.RECALL_MORE_MEMORY:
            old_radius = cognitive_state.memory_recall_radius
            cognitive_state.memory_recall_radius = min(2.0, old_radius * 1.5)
            state_changed = cognitive_state.memory_recall_radius != old_radius

        elif action == MDNAction.REFRAME_QUESTION:
            old_question = cognitive_state.question_representation
            cognitive_state.question_representation = self._reframe_question(cognitive_state, memory_context)
            state_changed = cognitive_state.question_representation != old_question

        elif action == MDNAction.COUNTERFACTUAL_SIMULATE:
            old_radius = cognitive_state.memory_recall_radius
            cognitive_state.memory_recall_radius = min(2.0, old_radius * 1.2)
            state_changed = cognitive_state.memory_recall_radius != old_radius

        elif action == MDNAction.SWITCH_SPECIALIST:
            old_role = cognitive_state.specialist_role
            cognitive_state.specialist_role = self._select_optimal_specialist(cognitive_state, task_input)

            # Update task input with previous specialist responses for context
            if specialist_responses:
                context_lines = []
                for i, response in enumerate(specialist_responses[-3:]):  # Include last 3 responses for context
                    context_lines.append(f"Previous_Answer_{i+1} ({response['specialist']}): {response['answer']}")

                # Reconstruct task input with context
                current_task_input = task_input
                if context_lines:
                    current_task_input += "\n" + "\n".join(context_lines)
                    print(f"      ? Enhanced task input with {len(context_lines)} previous specialist responses")

            state_changed = cognitive_state.specialist_role != old_role

        elif action == MDNAction.DEEPEN_REASONING:
            old_level = cognitive_state.abstraction_level
            cognitive_state.abstraction_level = min(1.0, old_level + 0.2)
            state_changed = cognitive_state.abstraction_level != old_level

        elif action == MDNAction.REEVALUATE_ANSWER:
            old_confidence = cognitive_state.confidence_answer
            cognitive_state.confidence_answer = max(0.0, old_confidence * 0.7)
            state_changed = True  # Always consider this a state change

        elif action == MDNAction.EXPAND_ATTENTION:
            # Simplified attention expansion
            state_changed = True

        elif action == MDNAction.NARROW_FOCUS:
            # Simplified focus narrowing
            state_changed = True

        elif action == MDNAction.INCREASE_ABSTRACTION:
            old_level = cognitive_state.abstraction_level
            cognitive_state.abstraction_level = min(1.0, old_level + 0.3)
            state_changed = cognitive_state.abstraction_level != old_level

        elif action == MDNAction.DECREASE_ABSTRACTION:
            old_level = cognitive_state.abstraction_level
            cognitive_state.abstraction_level = max(0.0, old_level - 0.3)
            state_changed = cognitive_state.abstraction_level != old_level

        return state_changed

    def _reframe_question(self, cognitive_state: CognitiveState, memory_context: Dict[str, Any]) -> str:
        """Reframe the question based on memory insights"""
        memories = memory_context.get('recalled_memories', [])
        if memories:
            key_concept = memories[0].content.split(':')[0] if ':' in memories[0].content else memories[0].content[:20]
            return f"Considering {key_concept}: {cognitive_state.question_representation}"
        return cognitive_state.question_representation

    def _select_optimal_specialist(self, cognitive_state: CognitiveState, task_input: str) -> str:
        """Select the most appropriate specialist based on current state"""
        # Use trained router if available
        if self.trained_router:
            try:
                with torch.no_grad():
                    result = self.trained_router(task_input)
                    predicted = result['predicted_specialist']
                    
                    # Override obviously wrong predictions with keyword-based routing
                    task_lower = task_input.lower()
                    is_math_question = ("math" in task_lower or "calculate" in task_lower or "algebra" in task_lower or 
                                      "field" in task_lower or "extension" in task_lower or "group" in task_lower or 
                                      "ring" in task_lower or "polynomial" in task_lower or "equation" in task_lower or
                                      "sqrt" in task_lower or "degree" in task_lower or "finite" in task_lower)
                    
                    if is_math_question and predicted != "qwen_math_expert":
                        print(f"   [CYCLE] Overriding trained router ({predicted}) -> qwen_math_expert for math question")
                        return "qwen_math_expert"
                    
                    return predicted
            except Exception as e:
                print(f"[WARN]  Trained router failed: {e}, falling back to keyword matching")

        # Fallback to simple keyword matching
        task_lower = task_input.lower()
        if ("math" in task_lower or "calculate" in task_lower or "algebra" in task_lower or 
            "field" in task_lower or "extension" in task_lower or "group" in task_lower or 
            "ring" in task_lower or "polynomial" in task_lower or "equation" in task_lower or
            "sqrt" in task_lower or "degree" in task_lower or "finite" in task_lower):
            return "qwen_math_expert"
        elif "strategy" in task_lower or "plan" in task_lower:
            return "tiny_llama_planner"
        else:
            return "qwen_general_reasoner"

    def build_contextual_prompt(self, base_prompt: str, cognitive_state: CognitiveState, 
                               execution_history: List[Dict[str, Any]]) -> str:
        """Build a contextual prompt by injecting the last 3 execution_history entries"""
        if not cognitive_state.reasoning_history:
            return base_prompt
            
        # Get last 3 raw reasoning responses
        recent_reasoning = cognitive_state.reasoning_history[-3:]
        context_parts = []
        
        for i, raw_response in enumerate(recent_reasoning):
            context_parts.append(f"Previous Reasoning {i+1}:\n{raw_response}\n")
        
        context_text = "\n".join(context_parts)
        
        # Inject context before the main prompt
        contextual_prompt = f"""Based on previous reasoning attempts, refine your approach:

{context_text}
Now, building on the above reasoning, provide an improved answer:

{base_prompt}"""
        
        return contextual_prompt

    def _update_policy_confidence(self, cognitive_state: CognitiveState, action_plan: MDNActionPlan,
                                 reflection_output: ReflectionOutput, memory_context: Dict[str, Any]) -> float:
        """Update policy confidence (how well we know how to think about this)"""
        base_confidence = cognitive_state.confidence_policy

        # Increase confidence based on successful actions and memory support
        memory_support = len(memory_context.get('recalled_memories', [])) > 0
        action_success = action_plan.expected_confidence_gain > 0.1

        confidence_boost = 0.0
        if memory_support:
            confidence_boost += 0.1
        if action_success:
            confidence_boost += action_plan.expected_confidence_gain * 0.5

        new_confidence = min(1.0, base_confidence + confidence_boost)

        # Slight decay if no progress
        if action_plan.expected_confidence_gain < 0.05:
            new_confidence *= 0.98

        return new_confidence

    def _update_answer_confidence(self, cognitive_state: CognitiveState, reflection_output: ReflectionOutput,
                                 memory_context: Dict[str, Any]) -> float:
        """Update answer confidence (how trustworthy the result is)"""
        base_confidence = cognitive_state.confidence_answer

        # Answer confidence is more conservative - requires strong evidence
        memory_support = len(memory_context.get('recalled_memories', [])) >= 2
        reflection_support = reflection_output.confidence > 0.7

        if memory_support and reflection_support:
            new_confidence = min(1.0, base_confidence + 0.15)
        elif memory_support or reflection_support:
            new_confidence = min(1.0, base_confidence + 0.05)
        else:
            new_confidence = max(0.0, base_confidence - 0.02)  # Slow decay

        return new_confidence

    def _generate_answer(self, task_input: str, cognitive_state: CognitiveState,
                        recalled_memories: List[MemoryTrace], action_plan: MDNActionPlan = None) -> str:
        """Generate an answer using REAL model inference with context integration and detailed logging"""

        # Extract question and choices from task_input
        lines = task_input.split('\n')
        question = ""
        choices = []
        context_history = []

        for line in lines:
            if line.startswith('Question:'):
                question = line.replace('Question:', '').strip()
            elif line.startswith('Choices:'):
                choices_text = line.replace('Choices:', '').strip()
                choices = [c.strip() for c in choices_text.split(' | ')]
            elif line.startswith('Previous_Answer_'):
                # Extract previous specialist responses for context
                context_history.append(line)

        if not question:
            question = task_input  # Fallback

        # Build base prompt with choices and instructions (MDN planner formats the prompt)
        choice_letters = ['A', 'B', 'C', 'D', 'E', 'F'][:len(choices)]
        
        base_prompt = f"{question}\n\n"
        for i, choice in enumerate(choices):
            base_prompt += f"{choice_letters[i]}. {choice}\n"
        
        # Use specialist instructions from MDN planner if available
        if action_plan and 'specialist_instructions' in action_plan.state_change_required:
            specialist_instructions = action_plan.state_change_required['specialist_instructions']
            base_prompt += f"\n{specialist_instructions}\n\nAnswer:"
        else:
            # Fallback to default instructions
            base_prompt += "\nProvide step-by-step mathematical reasoning. Show your work clearly, then at the very end, respond with ONLY the single letter of your final answer (A, B, C, or D).\n\nExample format:\nStep 1: Analyze the problem and identify key concepts.\nStep 2: Apply relevant mathematical principles.\nStep 3: Perform calculations if needed.\nFinal answer: B\n\nAnswer:"
        
        # Build contextual prompt using reasoning history
        prompt = self.build_contextual_prompt(base_prompt, cognitive_state, [])  # execution_history not needed here
        
        # Add context from previous iterations if available
        if context_history:
            context_text = "\n".join(context_history)
            prompt = f"""Previous reasoning and answers from other specialists:
{context_text}

Based on the previous responses above, please provide a refined answer.

{prompt}"""

        # Use real model inference
        if self.use_real_models and hasattr(self, 'real_specialist_system'):
            # Map specialist role to model name
            specialist_map = {
                'qwen_math_expert': 'qwen_math',
                'qwen_general_reasoner': 'qwen_general',
                'tiny_llama_planner': 'llama_reasoning',
                'tiny_llama_critic': 'llama_reasoning'
            }

            model_name = specialist_map.get(cognitive_state.specialist_role, 'qwen_general')

            print(f"         Model {model_name} input prepared")
            print(f"         Generating response...")

            try:
                # Load model if needed
                if not self.real_specialist_system.inference_engine.load_model(model_name):
                    print(f"         [FAIL] Model loading failed for {model_name}")
                    return f"Based on {len(recalled_memories)} memory traces and {len(context_history)} previous responses, the answer requires careful consideration."

                # Generate real answer with complete formatted prompt
                predicted_answer, confidence, raw_response = self.real_specialist_system.inference_engine.generate_answer(
                    model_name, prompt, []  # Empty choices since prompt is already complete
                )
                print(f"         Response: '{predicted_answer}' (confidence: {confidence:.2f})")
                
                # Store raw response in reasoning history for iterative refinement
                cognitive_state.reasoning_history.append(raw_response)
                
                return predicted_answer

            except Exception as e:
                print(f"         [FAIL] Real inference failed: {e}")
                # Fallback to simulation
                return f"Based on {len(recalled_memories)} memory traces and {len(context_history)} previous responses, the answer requires careful consideration."
        else:
            # Fallback when real models disabled or unavailable
            print(f"         Real models disabled or unavailable")
            return f"Based on {len(recalled_memories)} memory traces and {len(context_history)} previous responses, the answer requires careful consideration."

    def _simulate_execution_confidence(self, answer: str, task_input: str) -> float:
        """Simulate post-execution confidence (would be from actual execution)"""
        # Simple simulation - in practice would be from specialist execution
        if "Paris" in answer:
            return 0.95  # High confidence for simple facts
        elif "formula" in answer:
            return 0.85  # Good confidence for mathematical answers
        else:
            return 0.6   # Moderate confidence for complex answers

def main():
    """Demonstration of BIOMIND Cognitive Control Architecture"""

    print("[BRAIN] BIOMIND Cognitive Control Architecture")
    print("=" * 60)

    # Initialize system
    specialists = ['qwen_math_expert', 'qwen_general_reasoner', 'tiny_llama_planner', 'tiny_llama_critic']
    biomind_system = BIOMINDRecursiveDMN(specialists, max_reflection_cycles=10)  # Increased for convergence

    # Test cases
    test_cases = [
        {
            'task': "Calculate the compound interest on $1000 at 5% annual rate for 3 years",
            'type': "math"
        },
        {
            'task': "Design a comprehensive strategy for reducing carbon emissions in urban areas",
            'type': "complex"
        },
        {
            'task': "What is the capital of France?",
            'type': "simple"
        }
    ]

    print(f"\n[TEST] Testing Cognitive Control Architecture")
    print("=" * 50)

    for i, test_case in enumerate(test_cases):
        print(f"\n[LAB] Test Case {i+1}: {test_case['type'].upper()}")

        start_time = time.time()
        result = biomind_system.recursive_reflection_cycle(
            test_case['task'],
            test_case['type']
        )
        elapsed_time = (time.time() - start_time) * 1000

        print(f"\n[CHART] Results Summary:")
        print(f"   Specialist: {result['selected_specialist']}")
        print(f"   Final Answer: {result.get('final_answer', 'No answer')[:80] if result.get('final_answer') else 'No answer'}...")
        print(f"   Policy Confidence: {result['executive_summary']['final_policy_confidence']:.3f}")
        print(f"   Answer Confidence: {result['executive_summary']['final_answer_confidence']:.3f}")
        print(f"   Reflection Cycles: {result['reflection_cycles']}")
        print(f"   Processing Time: {elapsed_time:.1f}ms")
        print(f"   Successful Commits: {result['executive_summary']['successful_commits']}/{result['executive_summary']['total_executions']}")

        # Print key cognitive control events
        print(f"\n   [BRAIN] Cognitive Control Summary:")
        mdn_actions = [r['mdn_action'] for r in result['reflection_history']]
        print(f"      MDN Actions: {', '.join(mdn_actions[:5])}{'...' if len(mdn_actions) > 5 else ''}")
        print(f"      Termination: {result.get('termination_reason', 'Unknown')}")

if __name__ == "__main__":
    main()