"""
Self-Play Reward Enhancement (SPRE) Data Generation

Advanced data generation system using self-play and reward modeling to create
high-quality training datasets through iterative improvement and validation.

Author: Nik Jois <nikjois@llamasearch.ai>
"""

import asyncio
import json
import logging
import random
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, cast

import numpy as np

logger = logging.getLogger(__name__)


class RewardType(str, Enum):
    """Types of rewards in the SPRE system"""
    
    QUALITY = "quality"
    DIVERSITY = "diversity"
    ACCURACY = "accuracy"
    ENGAGEMENT = "engagement"
    COHERENCE = "coherence"
    CREATIVITY = "creativity"
    SAFETY = "safety"


class PlayerRole(str, Enum):
    """Roles in the self-play system"""
    
    GENERATOR = "generator"
    EVALUATOR = "evaluator"
    REFINER = "refiner"
    VALIDATOR = "validator"


class DataType(str, Enum):
    """Data types for SPRE generation"""
    
    TEXT = "text"
    CONVERSATION = "conversation"
    REASONING = "reasoning"
    CREATIVE = "creative"
    TECHNICAL = "technical"
    EDUCATIONAL = "educational"


class ValidationStatus(str, Enum):
    """Validation status for generated items"""
    
    VALID = "valid"
    INVALID = "invalid"
    PENDING = "pending"
    NEEDS_REVIEW = "needs_review"


@dataclass
class SPREItem:
    """Individual item in SPRE dataset"""
    
    id: str
    data_type: DataType
    content: Dict[str, Any]
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    validation_status: ValidationStatus = ValidationStatus.PENDING
    tags: List[str] = field(default_factory=list)
    quality_score: float = 0.0
    reward_scores: Dict[str, float] = field(default_factory=dict)


@dataclass
class SPREDataset:
    """Dataset generated by SPRE system"""
    
    name: str
    description: str
    items: List[SPREItem] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    
    def add_item(self, item: SPREItem) -> None:
        """Add item to dataset"""
        self.items.append(item)
    
    def get_items_by_type(self, data_type: DataType) -> List[SPREItem]:
        """Get items by data type"""
        return [item for item in self.items if item.data_type == data_type]
    
    def get_valid_items(self) -> List[SPREItem]:
        """Get only valid items"""
        return [item for item in self.items if item.validation_status == ValidationStatus.VALID]


@dataclass
class PlaySession:
    """Individual self-play session"""
    
    session_id: str
    round_number: int
    players: Dict[str, PlayerRole]
    task_description: str
    generated_content: Dict[str, Any]
    reward_scores: Dict[str, float]
    feedback: List[str] = field(default_factory=list)
    session_duration: float = 0.0
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "session_id": self.session_id,
            "round_number": self.round_number,
            "players": {k: v.value for k, v in self.players.items()},
            "task_description": self.task_description,
            "generated_content": self.generated_content,
            "reward_scores": self.reward_scores,
            "feedback": self.feedback,
            "session_duration": self.session_duration,
            "created_at": self.created_at.isoformat()
        }


@dataclass
class RewardModel:
    """Reward model for evaluating generated content"""
    
    model_id: str
    reward_type: RewardType
    weights: Dict[str, float]
    baseline_score: float = 0.5
    learning_rate: float = 0.01
    performance_history: List[float] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "model_id": self.model_id,
            "reward_type": self.reward_type.value,
            "weights": self.weights,
            "baseline_score": self.baseline_score,
            "learning_rate": self.learning_rate,
            "performance_history": self.performance_history
        }


@dataclass
class SpreConfig:
    """Configuration for SPRE system"""
    
    max_rounds: int = 10
    players_per_session: int = 4
    reward_threshold: float = 0.8
    diversity_factor: float = 0.3
    learning_rate: float = 0.01
    enable_human_feedback: bool = False
    output_path: str = "./spre_output"
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "max_rounds": self.max_rounds,
            "players_per_session": self.players_per_session,
            "reward_threshold": self.reward_threshold,
            "diversity_factor": self.diversity_factor,
            "learning_rate": self.learning_rate,
            "enable_human_feedback": self.enable_human_feedback,
            "output_path": self.output_path
        }


class SPREGenerator:
    """Main SPRE Generator class for creating datasets"""
    
    def __init__(self, config: Optional[SpreConfig] = None):
        """Initialize SPRE generator"""
        self.config: SpreConfig = config or SpreConfig()
        self.engine: Optional['SpreEngine'] = None  # Will be initialized in _get_engine()
        self.generated_datasets: List[SPREDataset] = []
        
    def _get_engine(self) -> Optional['SpreEngine']:
        """Get or create SpreEngine instance"""
        if self.engine is None:
            # Initialize SpreEngine when needed - defined later in this file
            self.engine = globals()['SpreEngine'](self.config)
        return self.engine
        
    async def generate_dataset_async(
        self,
        name: str,
        count: int,
        description: str = "",
        **kwargs: Any,
    ) -> SPREDataset:
        """Public async interface for dataset generation."""
        return await self._generate_dataset_async(name, count, description, **kwargs)

    def generate_dataset(
        self,
        name: str,
        count: int,
        description: str = "",
        **kwargs: Any,
    ) -> SPREDataset:
        """Synchronous helper that is safe to call inside or outside an existing event loop."""
        try:
            loop = asyncio.get_running_loop()
            return loop.run_until_complete(
                self._generate_dataset_async(name, count, description, **kwargs)
            )
        except RuntimeError:
            # No running loop
            return asyncio.run(
                self._generate_dataset_async(name, count, description, **kwargs)
            )
    
    async def _generate_dataset_async(self, name: str, count: int, description: str = "", **kwargs: Any) -> SPREDataset:
        """Generate a SPRE dataset asynchronously"""
        # Ensure engine is initialized before we begin generation
        self._get_engine()

        dataset = SPREDataset(name=name, description=description)
        dataset.metadata.update(kwargs)
        
        # Generate items
        for i in range(count):
            item = await self._generate_item(i, **kwargs)
            dataset.add_item(item)
        
        # Validate dataset
        await self._validate_dataset(dataset)
        
        self.generated_datasets.append(dataset)
        return dataset
    
    async def _generate_item(self, index: int, **kwargs: Any) -> SPREItem:
        """Generate a single SPRE item"""
        # Determine data type
        data_type: DataType = kwargs.get("data_type", DataType.TEXT)
        # DataType is always an Enum, so no isinstance check needed
        # Generate content based on type
        content: Dict[str, Any]
        if data_type == DataType.CONVERSATION:
            content = await self._generate_conversation_item(**kwargs)
        elif data_type == DataType.REASONING:
            content = await self._generate_reasoning_item(**kwargs)
        elif data_type == DataType.CREATIVE:
            content = await self._generate_creative_item(**kwargs)
        elif data_type == DataType.TECHNICAL:
            content = await self._generate_technical_item(**kwargs)
        elif data_type == DataType.EDUCATIONAL:
            content = await self._generate_educational_item(**kwargs)
        else:
            content = await self._generate_text_item(**kwargs)
        
        # Create SPRE item
        item = SPREItem(
            id=f"spre_item_{index}_{int(datetime.now().timestamp())}",
            data_type=data_type,
            content=content,
            metadata={"generation_index": index, **kwargs},
            tags=cast(List[str], kwargs.get("tags", [])),
        )
        
        # Calculate quality score
        item.quality_score = await self._calculate_quality_score(item)
        
        return item
    
    async def _generate_conversation_item(self, **kwargs: Any) -> Dict[str, Any]:
        """Generate conversation content"""
        topic: str = kwargs.get("topic", "general discussion")
        num_turns: int = kwargs.get("num_turns", 4)
        # Use SPRE engine for generation
        task_description = f"Generate a conversation about {topic}"
        requirements = {"num_turns": num_turns, "topic": topic}
        engine = self._get_engine()
        if engine is None:
            raise RuntimeError("SPRE engine is not initialized")
        session = await engine.run_self_play_session(task_description, requirements)
        return {
            "type": "conversation",
            "topic": topic,
            "turns": session.generated_content.get("turns", []),
            "metadata": {"session_id": session.session_id, "quality_scores": session.reward_scores}
        }
    
    async def _generate_reasoning_item(self, **kwargs: Any) -> Dict[str, Any]:
        """Generate reasoning content"""
        difficulty: str = kwargs.get("difficulty", "medium")
        domain: str = kwargs.get("domain", "logic")
        task_description = f"Generate a reasoning problem in {domain} with {difficulty} difficulty"
        requirements = {"difficulty": difficulty, "domain": domain}
        engine = self._get_engine()
        if engine is None:
            raise RuntimeError("SPRE engine is not initialized")
        session = await engine.run_self_play_session(task_description, requirements)
        return {
            "type": "reasoning",
            "difficulty": difficulty,
            "domain": domain,
            "problem": session.generated_content.get("problem", "Generated reasoning problem"),
            "solution": session.generated_content.get("solution_steps", []),
            "metadata": {"session_id": session.session_id, "quality_scores": session.reward_scores}
        }
    
    async def _generate_creative_item(self, **kwargs: Any) -> Dict[str, Any]:
        """Generate creative content"""
        style: str = kwargs.get("style", "narrative")
        theme: str = kwargs.get("theme", "adventure")
        task_description = f"Generate creative content in {style} style with {theme} theme"
        requirements = {"style": style, "theme": theme}
        engine = self._get_engine()
        if engine is None:
            raise RuntimeError("SPRE engine is not initialized")
        session = await engine.run_self_play_session(task_description, requirements)
        return {
            "type": "creative",
            "style": style,
            "theme": theme,
            "content": session.generated_content.get("content", "Generated creative content"),
            "metadata": {"session_id": session.session_id, "quality_scores": session.reward_scores}
        }
    
    async def _generate_technical_item(self, **kwargs: Any) -> Dict[str, Any]:
        """Generate technical content"""
        domain: str = kwargs.get("domain", "software")
        complexity: str = kwargs.get("complexity", "intermediate")
        task_description = f"Generate technical content in {domain} with {complexity} complexity"
        requirements = {"domain": domain, "complexity": complexity}
        engine = self._get_engine()
        if engine is None:
            raise RuntimeError("SPRE engine is not initialized")
        session = await engine.run_self_play_session(task_description, requirements)
        return {
            "type": "technical",
            "domain": domain,
            "complexity": complexity,
            "content": session.generated_content.get("content", "Generated technical content"),
            "metadata": {"session_id": session.session_id, "quality_scores": session.reward_scores}
        }
    
    async def _generate_educational_item(self, **kwargs: Any) -> Dict[str, Any]:
        """Generate educational content"""
        subject: str = kwargs.get("subject", "general")
        level: str = kwargs.get("level", "beginner")
        task_description = f"Generate educational content in {subject} for {level} level"
        requirements = {"subject": subject, "level": level}
        engine = self._get_engine()
        if engine is None:
            raise RuntimeError("SPRE engine is not initialized")
        session = await engine.run_self_play_session(task_description, requirements)
        return {
            "type": "educational",
            "subject": subject,
            "level": level,
            "content": session.generated_content.get("content", "Generated educational content"),
            "metadata": {"session_id": session.session_id, "quality_scores": session.reward_scores}
        }
    
    async def _generate_text_item(self, **kwargs: Any) -> Dict[str, Any]:
        """Generate generic text content"""
        topic: str = kwargs.get("topic", "general")
        task_description = f"Generate text content about {topic}"
        requirements = {"topic": topic}
        engine = self._get_engine()
        if engine is None:
            raise RuntimeError("SPRE engine is not initialized")
        session = await engine.run_self_play_session(task_description, requirements)
        return {
            "type": "text",
            "topic": topic,
            "content": session.generated_content.get("content", "Generated text content"),
            "metadata": {"session_id": session.session_id, "quality_scores": session.reward_scores}
        }
    
    async def _calculate_quality_score(self, item: SPREItem) -> float:
        """Calculate quality score for an item"""
        # Use reward models to calculate quality
        if item.reward_scores:
            return float(np.mean(list(item.reward_scores.values())))
        
        # Fallback quality calculation
        content = item.content
        score = 0.5
        
        if content.get("type"):
            score += 0.1
        
        if len(str(content)) > 100:
            score += 0.2
        
        if content.get("metadata"):
            score += 0.1
        
        return min(1.0, score)
    
    async def _validate_dataset(self, dataset: SPREDataset) -> None:
        """Validate dataset items"""
        for item in dataset.items:
            # Basic validation
            if not item.content:
                item.validation_status = ValidationStatus.INVALID
                continue
            
            # Content validation
            if len(str(item.content)) < 10:
                item.validation_status = ValidationStatus.INVALID
                continue
            
            # Quality threshold
            if item.quality_score < 0.3:
                item.validation_status = ValidationStatus.NEEDS_REVIEW
                continue
            
            item.validation_status = ValidationStatus.VALID
    
    async def generate_from_prompts(self, prompts: List[str], output_file: Optional[str] = None) -> List[Dict[str, Any]]:
        """Generate dataset from list of prompts"""
        generated_items: List[Dict[str, Any]] = []
        engine = self._get_engine()
        if engine is None:
            raise RuntimeError("SPRE engine is not initialized")
        for i, prompt in enumerate(prompts):
            # Create task from prompt
            task_description = f"Process prompt: {prompt}"
            requirements = {"prompt": prompt, "index": i}
            try:
                session = await engine.run_self_play_session(task_description, requirements)
                item = {
                    "id": f"prompt_item_{i}",
                    "prompt": prompt,
                    "generated_content": session.generated_content,
                    "quality_scores": session.reward_scores,
                    "session_id": session.session_id,
                    "created_at": datetime.now(timezone.utc).isoformat()
                }
                generated_items.append(item)
            except Exception as e:
                logger.error(f"Failed to generate from prompt {i}: {e}")
                continue
        # Save to file if specified
        if output_file:
            with open(output_file, 'w') as f:
                json.dump(generated_items, f, indent=2)
        return generated_items
    
    def get_dataset_stats(self) -> Dict[str, Any]:
        """Get statistics about generated datasets"""
        total_items = sum(len(dataset.items) for dataset in self.generated_datasets)
        valid_items = sum(len(dataset.get_valid_items()) for dataset in self.generated_datasets)
        
        return {
            "total_datasets": len(self.generated_datasets),
            "total_items": total_items,
            "valid_items": valid_items,
            "validation_rate": valid_items / total_items if total_items > 0 else 0.0,
            "datasets": [
                {
                    "name": dataset.name,
                    "description": dataset.description,
                    "items": len(dataset.items),
                    "valid_items": len(dataset.get_valid_items()),
                    "created_at": dataset.created_at
                }
                for dataset in self.generated_datasets
            ]
        }


class SprePlayer:
    """Base class for self-play players"""
    
    def __init__(self, player_id: str, role: PlayerRole):
        """Initialize player"""
        self.player_id = player_id
        self.role = role
        self.performance_history: List[float] = []
        self.last_session_score: float = 0.0
        self.improvement_rate: float = 0.0
        
    async def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
        """Take action based on role"""
        if self.role == PlayerRole.GENERATOR:
            return await self._generate_content(context)
        elif self.role == PlayerRole.EVALUATOR:
            return await self._evaluate_content(context)
        elif self.role == PlayerRole.REFINER:
            return await self._refine_content(context)
        elif self.role == PlayerRole.VALIDATOR:
            return await self._validate_content(context)
        else:
            raise ValueError(f"Unknown role: {self.role}")
    
    async def _generate_content(self, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate content based on context"""
        task_type = context.get("task_type", "general")
        requirements = context.get("requirements", {})
        
        # Generate content based on task type
        if task_type == "conversation":
            return await self._generate_conversation(requirements)
        elif task_type == "reasoning":
            return await self._generate_reasoning_problem(requirements)
        elif task_type == "creative":
            return await self._generate_creative_content(requirements)
        else:
            return await self._generate_generic_content(requirements)
    
    async def _generate_conversation(self, requirements: Dict[str, Any]) -> Dict[str, Any]:
        """Generate conversation data"""
        topic: str = requirements.get("topic", "general discussion")
        num_turns: int = requirements.get("num_turns", 4)
        conversation: Dict[str, Any] = {
            "type": "conversation",
            "topic": topic,
            "turns": []  # type: List[Dict[str, Any]]
        }
        for i in range(num_turns):
            turn: Dict[str, Any] = {
                "turn_id": i,
                "speaker": "user" if i % 2 == 0 else "assistant",
                "text": f"Turn {i}: Discussion about {topic}",
                "intent": "information_seeking" if i % 2 == 0 else "informative_response"
            }
            conversation["turns"].append(turn)
        return conversation
    
    async def _generate_reasoning_problem(self, requirements: Dict[str, Any]) -> Dict[str, Any]:
        """Generate reasoning problem"""
        difficulty = requirements.get("difficulty", "medium")
        domain = requirements.get("domain", "logic")
        
        return {
            "type": "reasoning",
            "difficulty": difficulty,
            "domain": domain,
            "problem": f"Given the constraints in {domain}, solve for the optimal solution",
            "solution_steps": [
                "1. Analyze the problem constraints",
                "2. Identify key variables and relationships",
                "3. Apply logical reasoning principles",
                "4. Verify the solution"
            ],
            "expected_answer": f"Solution for {difficulty} {domain} problem"
        }
    
    async def _generate_creative_content(self, requirements: Dict[str, Any]) -> Dict[str, Any]:
        """Generate creative content"""
        style = requirements.get("style", "narrative")
        theme = requirements.get("theme", "adventure")
        
        return {
            "type": "creative",
            "style": style,
            "theme": theme,
            "content": f"A {style} piece exploring the theme of {theme}",
            "elements": ["character_development", "plot_progression", "thematic_exploration"]
        }
    
    async def _generate_generic_content(self, requirements: Dict[str, Any]) -> Dict[str, Any]:
        """Generate generic content"""
        return {
            "type": "generic",
            "content": "Generated content based on requirements",
            "metadata": {"generated_by": self.player_id, "requirements": requirements}
        }
    
    async def _evaluate_content(self, context: Dict[str, Any]) -> Dict[str, Any]:
        """Evaluate content quality"""
        content = context.get("content", {})
        # Evaluate different aspects
        quality_score = await self._assess_quality(content)
        coherence_score = await self._assess_coherence(content)
        creativity_score = await self._assess_creativity(content)
        feedback = await self._generate_feedback(content)
        overall: float = (
            float(quality_score) + float(coherence_score) + float(creativity_score)
        ) / 3.0
        evaluation = {
            "quality": quality_score,
            "coherence": coherence_score,
            "creativity": creativity_score,
            "overall": overall,
            "feedback": feedback
        }
        return evaluation
    
    async def _assess_quality(self, content: Dict[str, Any]) -> float:
        """Assess content quality"""
        # Basic quality metrics
        score = 0.5
        
        if content.get("type"):
            score += 0.1
        
        if len(str(content)) > 100:
            score += 0.2
        
        # Type-specific quality checks
        if content.get("type") == "conversation":
            if content.get("turns") and len(content["turns"]) > 0:
                score += 0.2
        elif content.get("type") == "reasoning":
            if content.get("solution_steps"):
                score += 0.2
        
        return min(1.0, score)
    
    async def _assess_coherence(self, content: Dict[str, Any]) -> float:
        """Assess content coherence"""
        # Simple coherence assessment
        score = 0.7  # Base coherence score
        
        if content.get("type") == "conversation":
            turns = content.get("turns", [])
            if len(turns) >= 2:
                score += 0.1
        
        return min(1.0, score)
    
    async def _assess_creativity(self, content: Dict[str, Any]) -> float:
        """Assess content creativity"""
        # Basic creativity assessment
        score = 0.6
        
        if content.get("type") == "creative":
            score += 0.2
        
        if content.get("elements"):
            score += 0.1
        
        return min(1.0, score)
    
    async def _generate_feedback(self, content: Dict[str, Any]) -> List[str]:
        """Generate feedback for content"""
        feedback: List[str] = []
        if not content.get("type"):
            feedback.append("Content should specify a clear type")
        if len(str(content)) < 50:
            feedback.append("Content could be more detailed")
        if content.get("type") == "conversation":
            if not content.get("turns"):
                feedback.append("Conversation should include dialogue turns")
        if not feedback:
            feedback.append("Good quality content")
        return feedback
    
    async def _refine_content(self, context: Dict[str, Any]) -> Dict[str, Any]:
        """Refine content based on feedback"""
        content = context.get("content", {})
        feedback = context.get("feedback", [])
        
        refined_content = content.copy()
        
        # Apply refinements based on feedback
        for feedback_item in feedback:
            if "more detailed" in feedback_item.lower():
                refined_content["enhanced"] = True
                refined_content["additional_details"] = "Enhanced with more details"
            
            if "dialogue turns" in feedback_item.lower():
                if refined_content.get("type") == "conversation":
                    if not refined_content.get("turns"):
                        refined_content["turns"] = [
                            {"turn_id": 0, "speaker": "user", "text": "Hello"},
                            {"turn_id": 1, "speaker": "assistant", "text": "Hi there!"}
                        ]
        
        refined_content["refined_by"] = self.player_id
        return refined_content
    
    async def _validate_content(self, context: Dict[str, Any]) -> Dict[str, Any]:
        """Validate content meets requirements"""
        content = context.get("content", {})
        requirements = context.get("requirements", {})
        issues_list: List[str] = []
        valid: bool = True
        score: float = 1.0
        # Check requirements
        if requirements.get("min_length"):
            content_length = len(str(content))
            if content_length < requirements["min_length"]:
                valid = False
                issues_list.append(f"Content too short: {content_length} < {requirements['min_length']}")
                score *= 0.8
        if requirements.get("required_type"):
            if content.get("type") != requirements["required_type"]:
                valid = False
                issues_list.append(f"Wrong type: {content.get('type')} != {requirements['required_type']}")
                score *= 0.7
        validation_result: Dict[str, Any] = {
            "valid": valid,
            "issues": issues_list,
            "score": score
        }
        return validation_result
    
    def update_performance(self, session_score: float) -> None:
        """Update player performance metrics"""
        self.performance_history.append(session_score)
        
        if len(self.performance_history) > 1:
            self.improvement_rate = session_score - self.performance_history[-2]
        
        self.last_session_score = session_score
    
    def get_performance_summary(self) -> Dict[str, Any]:
        """Get performance summary"""
        return {
            "player_id": self.player_id,
            "role": self.role.value,
            "sessions_played": len(self.performance_history),
            "average_score": np.mean(self.performance_history) if self.performance_history else 0.0,
            "best_score": max(self.performance_history) if self.performance_history else 0.0,
            "latest_score": self.last_session_score,
            "improvement_rate": self.improvement_rate
        }


class SpreRewardModel:
    """Reward model for SPRE system"""
    
    def __init__(self, model_id: str, reward_type: RewardType):
        """Initialize reward model"""
        self.model_id = model_id
        self.reward_type = reward_type
        self.weights = self._initialize_weights()
        self.baseline_score = 0.5
        self.learning_rate = 0.01
        self.training_history: List[Dict[str, Any]] = []
    
    def _initialize_weights(self) -> Dict[str, float]:
        """Initialize reward model weights"""
        if self.reward_type == RewardType.QUALITY:
            return {
                "completeness": 0.3,
                "accuracy": 0.3,
                "relevance": 0.2,
                "clarity": 0.2
            }
        elif self.reward_type == RewardType.DIVERSITY:
            return {
                "uniqueness": 0.4,
                "variety": 0.3,
                "novelty": 0.3
            }
        elif self.reward_type == RewardType.CREATIVITY:
            return {
                "originality": 0.4,
                "imagination": 0.3,
                "innovation": 0.3
            }
        else:
            return {
                "general_quality": 0.5,
                "structure": 0.3,
                "content": 0.2
            }
    
    async def calculate_reward(self, content: Dict[str, Any], context: Dict[str, Any]) -> float:
        """Calculate reward for generated content"""
        feature_scores = await self._extract_features(content, context)
        
        # Calculate weighted score
        total_score = 0.0
        total_weight = 0.0
        
        for feature, weight in self.weights.items():
            if feature in feature_scores:
                total_score += feature_scores[feature] * weight
                total_weight += weight
        
        # Normalize score
        reward = total_score / max(total_weight, 1.0) if total_weight > 0 else self.baseline_score
        
        # Apply learning adjustments
        adjusted_reward = self._apply_learning_adjustments(reward, content, context)
        
        return max(0.0, min(1.0, adjusted_reward))
    
    async def _extract_features(self, content: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, float]:
        """Extract features for reward calculation"""
        features: Dict[str, float] = {}
        # Basic features
        features["completeness"] = await self._assess_completeness(content)
        features["structure"] = await self._assess_structure(content)
        features["content"] = await self._assess_content_quality(content)
        # Type-specific features
        if self.reward_type == RewardType.QUALITY:
            features.update(await self._extract_quality_features(content, context))
        elif self.reward_type == RewardType.DIVERSITY:
            features.update(await self._extract_diversity_features(content, context))
        elif self.reward_type == RewardType.CREATIVITY:
            features.update(await self._extract_creativity_features(content, context))
        return features
    
    async def _assess_completeness(self, content: Dict[str, Any]) -> float:
        """Assess content completeness"""
        required_fields = ["type"]
        present_fields = sum(1 for field in required_fields if field in content)
        return present_fields / len(required_fields)
    
    async def _assess_structure(self, content: Dict[str, Any]) -> float:
        """Assess content structure"""
        score = 0.5
        score += 0.2  # Dict[str, Any] is always a dict
        if content.get("type"):
            score += 0.3
        return min(1.0, score)
    
    async def _assess_content_quality(self, content: Dict[str, Any]) -> float:
        """Assess general content quality"""
        score = 0.5
        
        content_str = str(content)
        if len(content_str) > 50:
            score += 0.2
        
        if len(content_str) > 200:
            score += 0.3
        
        return min(1.0, score)
    
    async def _extract_quality_features(self, content: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, float]:
        """Extract quality-specific features"""
        return {
            "accuracy": 0.8,  # Would use actual accuracy assessment
            "relevance": 0.9,  # Would use relevance scoring
            "clarity": 0.85   # Would use clarity assessment
        }
    
    async def _extract_diversity_features(self, content: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, float]:
        """Extract diversity-specific features"""
        return {
            "uniqueness": random.uniform(0.6, 1.0),  # Would use actual uniqueness scoring
            "variety": random.uniform(0.7, 1.0),     # Would use variety assessment
            "novelty": random.uniform(0.5, 0.9)      # Would use novelty detection
        }
    
    async def _extract_creativity_features(self, content: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, float]:
        """Extract creativity-specific features"""
        return {
            "originality": random.uniform(0.6, 1.0),   # Would use originality assessment
            "imagination": random.uniform(0.7, 0.95),  # Would use imagination scoring
            "innovation": random.uniform(0.5, 0.9)     # Would use innovation detection
        }
    
    def _apply_learning_adjustments(self, base_reward: float, content: Dict[str, Any], context: Dict[str, Any]) -> float:
        """Apply learning-based adjustments to reward"""
        # Simple learning adjustment based on history
        if self.training_history:
            recent_scores: List[float] = [float(entry["reward"]) for entry in self.training_history[-10:]]
            avg_recent: float = float(np.mean(recent_scores)) if recent_scores else 0.0
            # Adjust reward based on recent performance
            if base_reward > avg_recent:
                adjustment = self.learning_rate * (base_reward - avg_recent)
            else:
                adjustment = -self.learning_rate * (avg_recent - base_reward)
            return float(base_reward + adjustment)
        return float(base_reward)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "model_id": self.model_id,
            "reward_type": self.reward_type.value,
            "weights": self.weights,
            "baseline_score": self.baseline_score,
            "learning_rate": self.learning_rate,
            "training_history": self.training_history,
        }
    
    def update_model(self, training_data: List[Dict[str, Any]]) -> None:
        """Update reward model with new training data"""
        self.training_history.extend(training_data)
        
        # Simple weight update based on feedback
        if len(training_data) > 0:
            avg_reward = np.mean([entry.get("reward", 0.5) for entry in training_data])
            
            # Adjust weights slightly based on performance
            for feature in self.weights:
                if avg_reward > self.baseline_score:
                    self.weights[feature] *= (1 + self.learning_rate)
                else:
                    self.weights[feature] *= (1 - self.learning_rate)
        
        # Keep recent history
        if len(self.training_history) > 1000:
            self.training_history = self.training_history[-1000:]


class SpreEngine:
    """Main SPRE (Self-Play Reward Enhancement) engine"""
    
    def __init__(self, config: Optional[SpreConfig] = None):
        """Initialize SPRE engine"""
        self.config = config or SpreConfig()
        self.players: Dict[str, SprePlayer] = {}
        self.reward_models: Dict[str, SpreRewardModel] = {}
        self.session_history: List[PlaySession] = []
        self.current_round = 0
        
        # Initialize default components
        self._initialize_players()
        self._initialize_reward_models()
    
    def _initialize_players(self) -> None:
        """Initialize default players"""
        roles = [PlayerRole.GENERATOR, PlayerRole.EVALUATOR, PlayerRole.REFINER, PlayerRole.VALIDATOR]
        
        for i, role in enumerate(roles):
            player_id = f"{role.value}_{i}"
            player = SprePlayer(player_id, role)
            self.players[player_id] = player
    
    def _initialize_reward_models(self) -> None:
        """Initialize reward models"""
        reward_types = [RewardType.QUALITY, RewardType.DIVERSITY, RewardType.CREATIVITY]
        
        for reward_type in reward_types:
            model_id = f"reward_{reward_type.value}"
            model = SpreRewardModel(model_id, reward_type)
            self.reward_models[model_id] = model
    
    def add_player(self, player: SprePlayer) -> None:
        """Add a player to the system"""
        self.players[player.player_id] = player
        logger.info(f"Added player {player.player_id} with role {player.role.value}")
    
    def add_reward_model(self, model: SpreRewardModel) -> None:
        """Add a reward model to the system"""
        self.reward_models[model.model_id] = model
        logger.info(f"Added reward model {model.model_id} for {model.reward_type.value}")
    
    async def run_self_play_session(self, task_description: str, requirements: Dict[str, Any]) -> PlaySession:
        """Run a single self-play session"""
        session_id = f"session_{int(datetime.now().timestamp())}_{self.current_round}"
        start_time = asyncio.get_event_loop().time()
        logger.info(f"Starting self-play session: {session_id}")
        # Select players for this session
        selected_players = self._select_players_for_session()
        # Create session context
        context = {
            "task_description": task_description,
            "requirements": requirements,
            "session_id": session_id,
            "round_number": self.current_round
        }
        # Phase 1: Generation
        generator = next(
            (p for p in selected_players.values() if p.role == PlayerRole.GENERATOR),
            None,
        )
        if not generator:
            raise ValueError("No generator player available")
        generated_content = await generator.act(context)
        # Phase 2: Evaluation
        evaluator = next(
            (p for p in selected_players.values() if p.role == PlayerRole.EVALUATOR),
            None,
        )
        if evaluator:
            eval_context = {**context, "content": generated_content}
            evaluation = await evaluator.act(eval_context)
        else:
            evaluation = {"overall": 0.5, "feedback": ["No evaluator available"]}
        # Phase 3: Refinement (if needed)
        overall_score = evaluation.get("overall", 0.0)
        if isinstance(overall_score, list):
            overall_score = 0.0
        if isinstance(overall_score, float) and overall_score < self.config.reward_threshold:
            refiner = next(
                (p for p in selected_players.values() if p.role == PlayerRole.REFINER),
                None,
            )
            if refiner:
                refine_context = {
                    **context,
                    "content": generated_content,
                    "feedback": evaluation.get("feedback", [])
                }
                generated_content = await refiner.act(refine_context)
        # Phase 4: Validation
        validator = next(
            (p for p in selected_players.values() if p.role == PlayerRole.VALIDATOR),
            None,
        )
        if validator:
            validate_context = {**context, "content": generated_content}
            await validator.act(validate_context)
        # Calculate rewards
        reward_scores = await self._calculate_session_rewards(generated_content, context)
        # Create session record
        feedback_val = evaluation.get("feedback", [])
        if not isinstance(feedback_val, list):
            feedback_val = []
        else:
            feedback_val = cast(List[str], feedback_val)
        session = PlaySession(
            session_id=session_id,
            round_number=self.current_round,
            players={p.player_id: p.role for p in selected_players.values()},
            task_description=task_description,
            generated_content=generated_content,
            reward_scores=reward_scores,
            feedback=feedback_val,
            session_duration=asyncio.get_event_loop().time() - start_time
        )
        # Update player performance
        overall_score = float(np.mean(list(reward_scores.values()))) if reward_scores else 0.5
        for player in selected_players.values():
            player.update_performance(overall_score)
        # Store session
        self.session_history.append(session)
        logger.info(f"Completed session {session_id} with overall score: {overall_score:.3f}")
        return session
    
    def _select_players_for_session(self) -> Dict[str, SprePlayer]:
        """Select players for a session"""
        selected: Dict[str, SprePlayer] = {}
        # Select one player per role
        for role in PlayerRole:
            role_players = [p for p in self.players.values() if p.role == role]
            if role_players:
                # Select based on performance (with some randomness)
                if len(role_players) == 1:
                    selected[role.value] = role_players[0]
                else:
                    # Weighted selection based on performance
                    weights = [p.last_session_score + 0.1 for p in role_players]  # Add small base weight
                    selected_player = random.choices(role_players, weights=weights, k=1)[0]
                    selected[role.value] = selected_player
        return selected
    
    async def _calculate_session_rewards(self, content: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, float]:
        """Calculate rewards for session content"""
        rewards: Dict[str, float] = {}
        for model_id, model in self.reward_models.items():
            try:
                reward = await model.calculate_reward(content, context)
                rewards[model_id] = reward
            except Exception as e:
                logger.warning(f"Failed to calculate reward for {model_id}: {e}")
                rewards[model_id] = 0.5  # Default neutral reward
        return rewards
    
    async def run_training_rounds(self, num_rounds: int, tasks: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Run multiple training rounds"""
        logger.info(f"Starting SPRE training for {num_rounds} rounds")
        training_results: Dict[str, Any] = {
            "rounds_completed": 0,
            "sessions": [],  # type: List[Dict[str, Any]]
            "average_scores": [],  # type: List[float]
            "improvement_trend": []  # type: List[float]
        }
        for round_num in range(num_rounds):
            self.current_round = round_num
            # Run sessions for this round
            round_sessions: List[PlaySession] = []
            round_scores: List[float] = []
            for task in tasks:
                try:
                    session = await self.run_self_play_session(
                        task["description"],
                        task.get("requirements", {})
                    )
                    round_sessions.append(session)
                    # Calculate session score
                    session_score: float = (
                        float(np.mean(list(session.reward_scores.values())))
                        if session.reward_scores
                        else 0.5
                    )
                    round_scores.append(session_score)
                except Exception as e:
                    logger.error(f"Failed to run session in round {round_num}: {e}")
                    continue
            # Update training results
            training_results["rounds_completed"] = round_num + 1
            training_results["sessions"].extend([s.to_dict() for s in round_sessions])
            if round_scores:
                avg_score: float = float(np.mean(round_scores))
                training_results["average_scores"].append(avg_score)
                # Calculate improvement trend
                prev_avg_scores: List[float] = training_results["average_scores"]
                if len(prev_avg_scores) > 1:
                    improvement: float = avg_score - prev_avg_scores[-2]
                    training_results["improvement_trend"].append(improvement)
                logger.info(f"Round {round_num} completed. Average score: {avg_score:.3f}")
            # Update reward models with session data
            await self._update_reward_models(round_sessions)
        # Calculate final statistics
        avg_scores: List[float] = training_results["average_scores"]
        improvement_trend: List[float] = training_results["improvement_trend"]
        training_results["final_average"] = float(np.mean(avg_scores)) if avg_scores else 0.0
        training_results["total_improvement"] = float(sum(improvement_trend)) if improvement_trend else 0.0
        logger.info(f"SPRE training completed. Final average: {training_results['final_average']:.3f}")
        return training_results
    
    async def _update_reward_models(self, sessions: List[PlaySession]) -> None:
        """Update reward models based on session results"""
        for model_id, model in self.reward_models.items():
            training_data: List[Dict[str, Any]] = []
            for session in sessions:
                if model_id in session.reward_scores:
                    training_data.append(
                        {
                            "session_id": session.session_id,
                            "content": session.generated_content,
                            "reward": session.reward_scores[model_id],
                            "feedback": session.feedback,
                        }
                    )
            if training_data:
                model.update_model(training_data)
    
    def get_system_status(self) -> Dict[str, Any]:
        """Get current system status"""
        return {
            "config": self.config.to_dict(),
            "players": {
                player_id: player.get_performance_summary() 
                for player_id, player in self.players.items()
            },
            "reward_models": {
                model_id: {
                    "type": model.reward_type.value,
                    "baseline_score": model.baseline_score,
                    "training_sessions": len(model.training_history)
                }
                for model_id, model in self.reward_models.items()
            },
            "sessions_completed": len(self.session_history),
            "current_round": self.current_round,
            "average_session_score": (
                np.mean(
                    [
                        np.mean(list(s.reward_scores.values()))
                        for s in self.session_history
                        if s.reward_scores
                    ]
                )
                if self.session_history
                else 0.0
            )
        }
    
    async def export_training_data(self, output_path: Optional[str] = None) -> str:
        """Export training data and results"""
        export_path = output_path or f"{self.config.output_path}/spre_export_{int(datetime.now().timestamp())}.json"
        
        export_data = {
            "config": self.config.to_dict(),
            "sessions": [session.to_dict() for session in self.session_history],
            "player_performance": {
                player_id: player.get_performance_summary() 
                for player_id, player in self.players.items()
            },
            "reward_models": {
                model_id: model.to_dict() 
                for model_id, model in self.reward_models.items()
            },
            "export_timestamp": datetime.now(timezone.utc).isoformat()
        }
        
        # Ensure output directory exists
        Path(export_path).parent.mkdir(parents=True, exist_ok=True)
        
        with open(export_path, "w") as f:
            json.dump(export_data, f, indent=2)
        
        logger.info(f"Exported SPRE training data to {export_path}")
        return export_path


async def create_sample_spre_system() -> SpreEngine:
    """Create a sample SPRE system for demonstration"""
    config = SpreConfig(
        max_rounds=5,
        players_per_session=4,
        reward_threshold=0.7,
        diversity_factor=0.3
    )
    
    engine = SpreEngine(config)
    
    # Sample tasks
    tasks = [
        {
            "description": "Generate a creative dialogue about future technology",
            "requirements": {"num_turns": 4, "topic": "future_tech", "style": "creative"}
        },
        {
            "description": "Create a reasoning problem about resource allocation",
            "requirements": {"difficulty": "medium", "domain": "optimization"}
        },
        {
            "description": "Generate educational content about machine learning",
            "requirements": {"target_audience": "beginners", "format": "explanation"}
        }
    ]
    
    # Run a few training rounds
    results = await engine.run_training_rounds(3, tasks)
    logger.info(f"Sample SPRE system created with {results['rounds_completed']} rounds completed")
    
    return engine