agents.py

"""
ContextFlow Multi-Agent Integration

Brings together the core ContextFlow agents for the OpenEnv environment:
- DoubtPredictorAgent: RL-based confusion prediction
- BehavioralAgent: Behavior signal analysis  
- HandGestureAgent: Gesture-based learning signals
"""

import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
import json


class ConfusionLevel(str, Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


class InterventionType(str, Enum):
    HINT = "hint"
    SIMPLIFY = "simplify"
    BREAKDOWN = "breakdown"
    EXAMPLE = "example"
    SCAFFOLD = "scaffold"
    PEER_CONNECT = "peer_connect"
    BREAK = "break"
    ENCOURAGE = "encourage"


@dataclass
class LearningState:
    topic: str
    subtopic: str
    progress_percentage: float
    time_spent_seconds: int
    confusion_signals: float
    eye_tracking_confidence: float
    scroll_reversals: int
    selection_count: int
    previous_doubts_count: int
    mastery_level: float
    difficulty_rating: float
    time_of_day: int
    streak_days: int


@dataclass
class BehavioralSignal:
    signal_type: str
    value: float
    timestamp: datetime
    source: str
    metadata: Dict = field(default_factory=dict)


@dataclass
class GestureTemplate:
    gesture_id: str
    name: str
    description: str
    samples: List[List[float]] = field(default_factory=list)
    centroid: Optional[List[float]] = None
    threshold: float = 0.3
    trained: bool = False


@dataclass
class AgentPrediction:
    confusion_probability: float
    confusion_level: ConfusionLevel
    confidence: float
    recommended_intervention: InterventionType
    intervention_intensity: float
    reasoning: str
    supporting_signals: Dict[str, float]


class MultiModalFusion:
    """Fuses signals from multiple modalities"""
    
    def __init__(self):
        self.weights = {
            "behavioral": 0.25,
            "gesture": 0.25,
            "biometric": 0.20,
            "temporal": 0.15,
            "content": 0.15,
        }
    
    def fuse(
        self,
        behavioral: float,
        gesture: float,
        biometric: float,
        temporal: float,
        content: float,
    ) -> float:
        return (
            self.weights["behavioral"] * behavioral +
            self.weights["gesture"] * gesture +
            self.weights["biometric"] * biometric +
            self.weights["temporal"] * temporal +
            self.weights["content"] * content
        )


class ContextFlowAgent:
    """
    Integrated ContextFlow agent combining:
    - RL-based doubt prediction
    - Behavioral signal analysis
    - Gesture recognition
    - Multi-modal fusion
    """
    
    def __init__(self, config: Optional[Dict] = None):
        self.config = config or {}
        
        self.doubt_predictor = RLBasedPredictor()
        self.behavioral_analyzer = BehavioralAnalyzer()
        self.gesture_recognizer = GestureRecognizer()
        self.fusion = MultiModalFusion()
        
        self.history: List[Dict] = []
        self.episode_rewards: List[float] = []
        
        self.epsilon = 1.0
        self.epsilon_decay = 0.995
        self.epsilon_min = 0.01
    
    def predict(
        self,
        observation: Dict[str, Any],
        use_exploration: bool = True,
    ) -> AgentPrediction:
        behavioral_signal = self.behavioral_analyzer.analyze(observation)
        gesture_signal = self.gesture_recognizer.recognize(observation)
        biometric_signal = self._extract_biometric_signal(observation)
        temporal_signal = self._extract_temporal_signal(observation)
        content_signal = self._extract_content_signal(observation)
        
        fused_signal = self.fusion.fuse(
            behavioral=behavioral_signal,
            gesture=gesture_signal,
            biometric=biometric_signal,
            temporal=temporal_signal,
            content=content_signal,
        )
        
        if use_exploration and np.random.random() < self.epsilon:
            confusion_prob = np.random.uniform(0.3, 0.8)
        else:
            confusion_prob = self.doubt_predictor.predict(fused_signal)
        
        confusion_level = self._get_confusion_level(confusion_prob)
        intervention, intensity = self._get_recommendation(confusion_prob, confusion_level)
        
        return AgentPrediction(
            confusion_probability=confusion_prob,
            confusion_level=confusion_level,
            confidence=0.85,
            recommended_intervention=intervention,
            intervention_intensity=intensity,
            reasoning=self._generate_reasoning(behavioral_signal, gesture_signal, biometric_signal),
            supporting_signals={
                "behavioral": behavioral_signal,
                "gesture": gesture_signal,
                "biometric": biometric_signal,
                "temporal": temporal_signal,
                "content": content_signal,
                "fused": fused_signal,
            }
        )
    
    def update(self, reward: float, observation: Dict[str, Any]):
        self.episode_rewards.append(reward)
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
        
        if len(self.episode_rewards) > 100:
            recent_avg = np.mean(self.episode_rewards[-100:])
            self.doubt_predictor.update_q_value(recent_avg)
    
    def _extract_biometric_signal(self, obs: Dict) -> float:
        biometric = obs.get("biometric_features", [])
        if not biometric:
            return 0.5
        
        hr = biometric[0] if len(biometric) > 0 else 70.0
        gsr = biometric[1] if len(biometric) > 1 else 0.5
        
        hr_signal = min(1.0, max(0.0, (hr - 60) / 40))
        gsr_signal = min(1.0, max(0.0, gsr * 2))
        
        return (hr_signal + gsr_signal) / 2
    
    def _extract_temporal_signal(self, obs: Dict) -> float:
        time_spent = obs.get("learning_context", {}).get("time_spent", 0)
        
        if time_spent < 300:
            return 0.2
        elif time_spent < 900:
            return 0.4
        elif time_spent < 1800:
            return 0.6
        else:
            return 0.8 + min(0.2, (time_spent - 1800) / 3600)
    
    def _extract_content_signal(self, obs: Dict) -> float:
        difficulty = obs.get("learning_context", {}).get("difficulty", "medium")
        difficulty_map = {"easy": 0.2, "medium": 0.5, "hard": 0.8}
        return difficulty_map.get(difficulty, 0.5)
    
    def _get_confusion_level(self, prob: float) -> ConfusionLevel:
        if prob < 0.25:
            return ConfusionLevel.LOW
        elif prob < 0.5:
            return ConfusionLevel.MEDIUM
        elif prob < 0.75:
            return ConfusionLevel.HIGH
        else:
            return ConfusionLevel.CRITICAL
    
    def _get_recommendation(self, prob: float, level: ConfusionLevel) -> Tuple[InterventionType, float]:
        recommendations = {
            ConfusionLevel.LOW: (InterventionType.ENCOURAGE, 0.3),
            ConfusionLevel.MEDIUM: (InterventionType.HINT, 0.5),
            ConfusionLevel.HIGH: (InterventionType.SIMPLIFY, 0.7),
            ConfusionLevel.CRITICAL: (InterventionType.SCAFFOLD, 0.9),
        }
        return recommendations[level]
    
    def _generate_reasoning(self, behavioral: float, gesture: float, biometric: float) -> str:
        reasons = []
        
        if behavioral > 0.6:
            reasons.append("High scroll reversals and hesitation detected")
        if gesture > 0.6:
            reasons.append("Confusion-related gestures identified")
        if biometric > 0.6:
            reasons.append("Elevated physiological stress indicators")
        
        if not reasons:
            reasons.append("All signals within normal range")
        
        return "; ".join(reasons)


class RLBasedPredictor:
    """Q-learning based confusion predictor"""
    
    def __init__(self):
        self.q_values: Dict[float, float] = {}
        self.gamma = 0.95
        self.learning_rate = 0.1
    
    def predict(self, state: float) -> float:
        if state not in self.q_values:
            self.q_values[state] = 0.5
        
        base = self.q_values[state]
        noise = np.random.normal(0, 0.05)
        return np.clip(base + noise, 0.0, 1.0)
    
    def update_q_value(self, reward: float, state: float = 0.5):
        if state not in self.q_values:
            self.q_values[state] = 0.5
        
        self.q_values[state] += self.learning_rate * (
            reward - self.q_values[state]
        )


class BehavioralAnalyzer:
    """Analyzes behavioral signals for confusion indicators"""
    
    def __init__(self):
        self.baseline_scroll_speed = 1.0
        self.baseline_click_rate = 1.0
    
    def analyze(self, observation: Dict[str, Any]) -> float:
        behavioral = observation.get("behavioral_features", [])
        
        if not behavioral or len(behavioral) < 4:
            return 0.5
        
        scroll_reversal = behavioral[0]
        hesitation = behavioral[1]
        click_pattern = behavioral[2]
        time_on_task = behavioral[3]
        
        signals = [
            min(1.0, scroll_reversal * 2),
            min(1.0, hesitation * 2),
            min(1.0, click_pattern * 2),
            min(1.0, time_on_task / 1800),
        ]
        
        return np.mean(signals)


class GestureRecognizer:
    """Recognizes confusion-related gestures"""
    
    CONFUCSION_GESTURES = {
        "head_scratch": {"pattern": [0.7, 0.8, 0.9], "confidence": 0.85},
        "brow_furrow": {"pattern": [0.6, 0.7, 0.8], "confidence": 0.75},
        "hand_wave": {"pattern": [0.5, 0.6, 0.7], "confidence": 0.70},
        "thinking": {"pattern": [0.4, 0.5, 0.6], "confidence": 0.65},
    }
    
    def __init__(self):
        self.last_gesture = None
        self.gesture_duration = 0
    
    def recognize(self, observation: Dict[str, Any]) -> float:
        gesture_features = observation.get("gesture_features", [])
        
        if not gesture_features or len(gesture_features) < 21:
            return 0.3
        
        hand_variance = np.var(gesture_features[:21])
        movement_intensity = np.mean(np.abs(np.diff(gesture_features[:21])))
        
        confusion_score = min(1.0, (hand_variance * 5 + movement_intensity * 3))
        
        return confusion_score


class KnowledgeGraphAgent:
    """Tracks concept relationships and prerequisite chains"""
    
    def __init__(self):
        self.concepts: Dict[str, Dict] = {}
        self.prerequisites: Dict[str, List[str]] = {}
    
    def add_concept(self, concept: str, mastery: float, prerequisites: List[str]):
        self.concepts[concept] = {
            "mastery": mastery,
            "last_accessed": datetime.now(),
        }
        self.prerequisites[concept] = prerequisites
    
    def get_prerequisite_mastery(self, concept: str) -> float:
        prereqs = self.prerequisites.get(concept, [])
        if not prereqs:
            return 1.0
        
        masteries = [self.concepts.get(p, {}).get("mastery", 0.0) for p in prereqs]
        return min(masteries) if masteries else 1.0
    
    def predict_confusion_risk(self, concept: str) -> float:
        mastery = self.concepts.get(concept, {}).get("mastery", 0.0)
        prereq_mastery = self.get_prerequisite_mastery(concept)
        
        risk = (1 - mastery) * 0.6 + (1 - prereq_mastery) * 0.4
        return risk


class PeerLearningAgent:
    """Connects learners with similar struggles"""
    
    def __init__(self):
        self.learners: Dict[str, Dict] = {}
        self.doubt_patterns: Dict[str, List[str]] = {}
    
    def register_doubt(self, user_id: str, doubt: str):
        if user_id not in self.doubt_patterns:
            self.doubt_patterns[user_id] = []
        self.doubt_patterns[user_id].append(doubt)
    
    def find_similar_learners(self, doubt: str, top_k: int = 3) -> List[Dict]:
        matches = []
        for user_id, doubts in self.doubt_patterns.items():
            overlap = len(set(doubts) & {doubt})
            if overlap > 0:
                matches.append({
                    "user_id": user_id,
                    "overlap": overlap,
                    "solutions_shared": len(doubts) - overlap,
                })
        
        matches.sort(key=lambda x: x["overlap"], reverse=True)
        return matches[:top_k]


class RecallAgent:
    """Spaced repetition for concept reinforcement"""
    
    def __init__(self):
        self.cards: Dict[str, Dict] = {}
    
    def add_card(self, concept: str, quality: int = 0):
        self.cards[concept] = {
            "interval": 1,
            "ease_factor": 2.5,
            "repetitions": 0,
            "next_review": datetime.now(),
            "quality": quality,
        }
    
    def process_review(self, concept: str, quality: int) -> Dict:
        if concept not in self.cards:
            self.add_card(concept, quality)
            return {"interval": 1, "message": "New card added"}
        
        card = self.cards[concept]
        
        if quality < 3:
            card["repetitions"] = 0
            card["interval"] = 1
        else:
            if card["repetitions"] == 0:
                card["interval"] = 1
            elif card["repetitions"] == 1:
                card["interval"] = 6
            else:
                card["interval"] = int(card["interval"] * card["ease_factor"])
            
            card["repetitions"] += 1
        
        card["ease_factor"] = max(1.3, card["ease_factor"] + (0.1 - (5 - quality) * (0.08 + (5 - quality) * 0.02)))
        card["next_review"] = datetime.now()
        
        return {"interval": card["interval"], "ease_factor": card["ease_factor"]}


__all__ = [
    "ContextFlowAgent",
    "RLBasedPredictor",
    "BehavioralAnalyzer",
    "GestureRecognizer",
    "KnowledgeGraphAgent",
    "PeerLearningAgent",
    "RecallAgent",
    "MultiModalFusion",
    "ConfusionLevel",
    "InterventionType",
    "AgentPrediction",
]