"""
Enhanced Toxo Training Engine with Interactive Capabilities.

This module implements comprehensive training functionality including
interactive training flows, professional AI patterns, and dynamic
response generation capabilities.
"""

import asyncio
import json
import os
from pathlib import Path
from typing import Dict, Any, List, Optional, Union, Tuple
from dataclasses import dataclass, asdict
import numpy as np
from datetime import datetime
import yaml
import uuid

from .config import ToxoConfig
from .toxo_layer import ToxoLayer, ToxoLayerMetadata
from .rich_prompt_engine import RichPromptEngine, ProfessionalAIPatterns, DynamicDomainAnalyzer, EnhancedMetaPromptGenerator
from ..integrations.gemini_client import GeminiClient
from ..integrations.llm_client import MultiProviderLLMClient, create_llm_client
from ..training.soft_prompt import SoftPromptEngine
from ..processing.document_processor import DocumentProcessor
from ..utils.logger import get_logger, PerformanceTimer
from ..utils.exceptions import ToxoError, TrainingError, ValidationError


@dataclass
class TrainingExample:
    """A single training example."""
    input_text: str
    expected_output: str
    domain: str = "general"
    task_type: str = "general"
    context: Optional[Dict[str, Any]] = None
    source: str = "synthetic"  # "synthetic", "user", "document"
    quality_score: float = 1.0
    metadata: Optional[Dict[str, Any]] = None
    
    def __post_init__(self):
        """Initialize metadata if not provided."""
        if self.metadata is None:
            self.metadata = {
                "created_at": datetime.now().isoformat(),
                "domain": self.domain,
                "task_type": self.task_type
            }


@dataclass
class TrainingSession:
    """Represents a complete training session."""
    id: str
    name: str
    description: str
    domain: str
    examples: List[TrainingExample]
    performance_metrics: Dict[str, float]
    created_at: str
    completed_at: Optional[str] = None
    final_prompt: Optional[np.ndarray] = None


@dataclass
class InteractiveTrainingSession:
    """Interactive training session with real-time feedback"""
    id: str
    name: str
    description: str
    domain: str
    layer_id: Optional[str] = None
    contexts: List[str] = None
    training_metrics: Dict[str, float] = None
    created_at: str = None
    completed_at: Optional[str] = None
    
    def __post_init__(self):
        if self.contexts is None:
            self.contexts = []
        if self.training_metrics is None:
            self.training_metrics = {}
        if self.created_at is None:
            self.created_at = datetime.now().isoformat()


class InteractiveToxoTrainer:
    """Enhanced Interactive Toxo Training System with Professional AI Integration"""
    
    def __init__(
        self,
        config: Optional[ToxoConfig] = None,
        layer_name: str = "interactive_layer",
        domain: str = "general",
        description: str = "",
        llm_provider: str = "gemini",
        llm_model: str = None,
        llm_api_key: str = None
    ):
        """
        Initialize Interactive Toxo Trainer.
        
        Args:
            config: Toxo configuration
            layer_name: Name for the layer being trained
            domain: Domain specialization
            description: Task description in natural language
            llm_provider: LLM provider ("gemini", "openai", "claude")
            llm_model: Specific model name (optional, uses defaults)
            llm_api_key: API key for the LLM provider
        """
        self.config = config or ToxoConfig()
        self.layer_name = layer_name
        self.domain = domain
        self.description = description
        
        # LLM Provider Configuration
        self.llm_provider = llm_provider
        self.llm_model = llm_model
        self.llm_api_key = llm_api_key
        
        self.logger = get_logger(f"toxo.interactive_trainer.{layer_name}")
        
        # Initialize LLM client
        if self.llm_api_key and self.llm_model:
            self.gemini_client = create_llm_client(self.llm_provider, self.llm_api_key, self.llm_model)
            self.logger.info(f"InteractiveToxoTrainer initialized with {self.llm_provider} client")
        elif self.llm_api_key and not self.llm_model:
            raise ValueError("Model name is required when API key is provided")
        else:
            # Fallback to Gemini client for backward compatibility
            self.gemini_client = GeminiClient(self.config.gemini)
            self.logger.info("InteractiveToxoTrainer initialized with default Gemini client")
        
        # Never log API keys (even partially).
        
        self.document_processor = DocumentProcessor(self.config)
        self.prompt_engine = RichPromptEngine(gemini_client=self.gemini_client)
        
        # Enhanced components from interactive script
        self.professional_patterns = ProfessionalAIPatterns()
        self.domain_analyzer = DynamicDomainAnalyzer()
        self.meta_prompt_generator = EnhancedMetaPromptGenerator()
        
        # Training state
        self.current_session = None
        self.current_layer = None
        self.training_active = False
        
        self.logger.info(f"Interactive ToxoTrainer initialized for '{layer_name}' with professional AI patterns")
    
    def update_config(self, new_config: ToxoConfig):
        """Update the trainer's configuration and reinitialize components with new config."""
        self.config = new_config
        
        # Reinitialize GeminiClient with new config
        self.gemini_client = GeminiClient(self.config.gemini)
        
        # 🔑 DEBUG: Log the updated API key
        api_key_suffix = self.config.gemini.api_key[-8:] if self.config.gemini.api_key else "None"
        self.logger.info(f"🔑 InteractiveToxoTrainer config updated, GeminiClient now using API key: ...{api_key_suffix}")
        
        # Also update the layer's config if it exists
        if self.current_layer:
            self.current_layer.config = new_config
            # Force reinitialize the layer's GeminiClient
            self.current_layer.gemini_client = GeminiClient(new_config.gemini)
            layer_api_key_suffix = new_config.gemini.api_key[-8:] if new_config.gemini.api_key else "None"
            self.logger.info(f"🔑 Updated layer's GeminiClient to use API key: ...{layer_api_key_suffix}")
    
    async def create_interactive_layer(self, description: str) -> str:
        """Create a new Toxo layer with interactive training capabilities"""
        self.logger.info(f"Creating interactive Toxo layer: {description}")
        
        # Analyze domain from description
        domain_analysis = self.domain_analyzer.analyze_query_domain(description, {})
        detected_domain = domain_analysis['domain'].lower().replace(' ', '_')
        
        # Create enhanced layer with professional patterns
        layer = ToxoLayer(
            domain=detected_domain,
            task_type="interactive",
            config=self.config,
            name=self.layer_name,
            description=description,
            llm_provider=self.llm_provider,
            llm_model=self.llm_model,
            llm_api_key=self.llm_api_key
        )
        
        # Generate professional base prompt
        pattern = self.professional_patterns.get_pattern(detected_domain)
        
        # Create rich master prompt with professional patterns
        enhanced_prompt = f"""{pattern['identity']}

{pattern['approach']}

## DOMAIN EXPERTISE
Your specialized knowledge areas include:
{chr(10).join([f"• {area}" for area in pattern['expertise_areas']])}

## TASK DESCRIPTION
{description}

## TRAINING CONTEXT
You are being trained through interactive examples and contexts. Each interaction
will help you become more specialized and effective at this specific task.

{pattern.get('response_structure', '')}

## QUALITY STANDARDS
- Provide specific, actionable advice with quantitative backing where applicable
- Reference domain-specific knowledge and best practices
- Include risk assessments and mitigation strategies when relevant
- Demonstrate sophisticated understanding of domain dynamics
- Use professional terminology appropriately
- Support recommendations with logical reasoning and evidence
"""
        
        # Update layer with enhanced prompt
        layer.master_prompt = enhanced_prompt
        
        # Initialize training session
        session = InteractiveTrainingSession(
            id=str(uuid.uuid4()),
            name=f"Interactive Training - {self.layer_name}",
            description=description,
            domain=detected_domain,
            layer_id=layer.name
        )
        
        self.current_session = session
        self.current_layer = layer
        self.training_active = True
        
        self.logger.info(f"Created interactive layer with domain: {detected_domain}")
        return layer.name
    
    async def add_training_context(self, context_text: str) -> Dict[str, Any]:
        """Add training context with real-time adaptation"""
        if not self.current_session or not self.current_layer:
            raise TrainingError("No active training session. Create a layer first.")
        
        self.logger.info(f"Adding training context: {context_text[:100]}...")
        
        # Add context to session
        self.current_session.contexts.append(context_text)
        
        # Simulate soft prompt adaptation (in real implementation, this would use gradients)
        if hasattr(self.current_layer, 'soft_prompt_engine'):
            # Get current parameters
            current_params = self.current_layer.soft_prompt_engine.get_current_prompt()
            
            # Context-driven parameter updates (simplified simulation)
            context_influence = np.random.normal(0, 0.01, current_params.shape)
            updated_params = current_params + context_influence
            
            # Update parameters
            self.current_layer.soft_prompt_engine.set_prompt(updated_params)
            
            # Calculate adaptation metrics
            param_change = np.mean(np.abs(context_influence))
            
            # Update training metrics
            self.current_session.training_metrics.update({
                'context_count': len(self.current_session.contexts),
                'parameter_adaptation': param_change,
                'last_update': datetime.now().isoformat()
            })
            
            return {
                'context_added': True,
                'parameter_change': param_change,
                'total_contexts': len(self.current_session.contexts),
                'adaptation_strength': param_change * 1000  # Scale for display
            }
        
        return {
            'context_added': True,
            'total_contexts': len(self.current_session.contexts)
        }
    
    async def complete_interactive_training(self) -> Dict[str, Any]:
        """Complete interactive training and finalize the layer"""
        if not self.current_session or not self.current_layer:
            raise TrainingError("No active training session to complete.")
        
        # Finalize training session
        self.current_session.completed_at = datetime.now().isoformat()
        
        # Update layer metadata
        self.current_layer.metadata.training_examples = len(self.current_session.contexts)
        self.current_layer.metadata.performance_metrics = self.current_session.training_metrics
        self.current_layer.metadata.updated_at = datetime.now().isoformat()
        
        # Mark as trained
        self.current_layer.is_trained = True
        self.training_active = False
        
        training_summary = {
            'session_id': self.current_session.id,
            'layer_name': self.current_layer.name,
            'domain': self.current_session.domain,
            'total_contexts': len(self.current_session.contexts),
            'training_metrics': self.current_session.training_metrics,
            'completion_time': self.current_session.completed_at
        }
        
        self.logger.info(f"Completed interactive training for layer: {self.current_layer.name}")
        return training_summary
    
    async def process_interactive_query(
        self, 
        query: str, 
        use_professional_patterns: bool = True
    ) -> str:
        """Process query through trained layer with professional AI patterns"""
        if not self.current_layer:
            raise TrainingError("No trained layer available. Complete training first.")
        
        self.logger.info(f"Processing interactive query: {query[:100]}...")
        
        # Prepare context data
        context_data = {
            'layer_description': self.current_layer.description,
            'training_contexts': self.current_session.contexts if self.current_session else [],
            'context_count': len(self.current_session.contexts) if self.current_session else 0,
            'training_metrics': self.current_session.training_metrics if self.current_session else {},
            'layer_influence': 'trained' if self.current_layer.is_trained else 'training'
        }
        
        if use_professional_patterns:
            # Generate professional meta-prompt
            meta_prompt = self.meta_prompt_generator.create_professional_meta_prompt(
                base_prompt=self.current_layer.master_prompt,
                query=query,
                context_data=context_data,
                domain=self.current_session.domain if self.current_session else None
            )
            
            # Generate response using Gemini with professional patterns
            try:
                response = await self.gemini_client.generate(meta_prompt)
                
                # Format professional response
                formatted_response = self._format_professional_response(
                    response, 
                    context_data, 
                    "LIVE_PROFESSIONAL_AI"
                )
                
                return formatted_response
                
            except Exception as e:
                self.logger.warning(f"Professional pattern generation failed: {e}, falling back to standard")
                # Fallback to standard layer query
                return await self.current_layer.query(query, context=context_data)
        else:
            # Use standard layer query
            return await self.current_layer.query(query, context=context_data)
    
    def _format_professional_response(
        self, 
        response_content: str, 
        context_data: Dict[str, Any], 
        generation_type: str
    ) -> str:
        """Format response with professional metadata"""
        
        generation_labels = {
            "LIVE_PROFESSIONAL_AI": "🚀 **TOXO PROFESSIONAL AI - LIVE ANALYSIS**",
            "ENHANCED_SIMULATION": "🏦 **TOXO PROFESSIONAL AI - ENHANCED ANALYSIS**"
        }
        
        return f"""{generation_labels.get(generation_type, "🧠 **TOXO PROFESSIONAL AI**")}

{response_content}

---
**TOXO LAYER METADATA:**
- **Training Contexts**: {context_data.get('context_count', 0)} specialized contexts
- **Layer Status**: {context_data.get('layer_influence', 'N/A')}
- **Domain Specialization**: {context_data.get('domain', 'General')}
- **Response Generation**: {generation_type.replace('_', ' ').title()}
- **Enhancement Pattern**: Professional AI patterns from leading tools
- **Analysis Framework**: Dynamic domain intelligence with context sensitivity

*Professional analysis generated using Toxo's advanced training with complete context integration.*
"""

class ToxoTrainer:
    """
    Advanced trainer for Toxo layers with full system integration.
    
    Supports soft prompt training, memory integration, vector store knowledge retrieval,
    and comprehensive performance optimization.
    """
    
    def __init__(
        self,
        config: ToxoConfig,
        layer_name: str = "toxo_layer",
        domain: str = "general",
        description: str = "",
        task_type: str = "general",
        llm_provider: str = "gemini",
        llm_model: str = None,
        llm_api_key: str = None
    ):
        """Initialize the trainer."""
        self.config = config
        self.layer_name = layer_name
        self.domain = domain
        self.description = description
        self.task_type = task_type
        
        # LLM Provider Configuration
        self.llm_provider = llm_provider
        self.llm_model = llm_model
        self.llm_api_key = llm_api_key
        
        self.logger = get_logger(f"toxo.trainer.{layer_name}")
        
        # Training state
        self.layer: Optional[ToxoLayer] = None
        self.training_sessions: Dict[str, Any] = {}
        self.is_training = False
        
        # CRITICAL FIX: Add vector store integration
        self.vector_store = None
        self.knowledge_enhanced = False
        
        # Initialize LLM client for training data generation
        if self.llm_api_key and self.llm_model:
            self.gemini_client = create_llm_client(self.llm_provider, self.llm_api_key, self.llm_model)
            self.logger.info(f"ToxoTrainer initialized with {self.llm_provider} client")
        elif self.llm_api_key and not self.llm_model:
            raise ValueError("Model name is required when API key is provided")
        else:
            # Fallback to Gemini client for backward compatibility
            self.gemini_client = GeminiClient(config.gemini)
            self.logger.info("ToxoTrainer initialized with default Gemini client")
        
        self.logger.info(f"ToxoTrainer initialized for domain: {domain}")
    
    def set_vector_store(self, vector_store):
        """Set vector store for knowledge-enhanced training."""
        self.vector_store = vector_store
        self.knowledge_enhanced = True
        self.logger.info("Vector store integrated for knowledge-enhanced training")
    
    def create_training_session(
        self,
        name: str,
        description: str,
        domain: str,
        examples: List[TrainingExample]
    ) -> TrainingSession:
        """
        Create a new training session.
        
        Args:
            name: Session name
            description: Session description
            domain: Domain for the session
            examples: Training examples
            
        Returns:
            TrainingSession instance
        """
        if not examples:
            raise ValidationError("Training session must have at least one example")
        
        if not domain or domain.strip() == "":
            raise ValidationError("Domain cannot be empty")
        
        session = TrainingSession(
            id=str(uuid.uuid4()),
            name=name,
            description=description,
            domain=domain,
            examples=examples,
            performance_metrics={},
            created_at=datetime.now().isoformat()
        )
        
        self.training_sessions[session.id] = session
        self.logger.info(f"Created training session: {name} with {len(examples)} examples")
        
        return session
    
    async def define_task_from_description(
        self,
        task_description: str,
        session_name: str,
        num_examples: int = 5
    ) -> TrainingSession:
        """
        Define a task from natural language description.
        
        Args:
            task_description: Natural language task description
            session_name: Name for the training session
            num_examples: Number of synthetic examples to generate
            
        Returns:
            TrainingSession instance
        """
        try:
            # Analyze task to determine domain
            domain = self._infer_domain_from_description(task_description)
            
            # Generate synthetic examples
            examples = await self.generate_synthetic_examples(
                task_description=task_description,
                domain=domain,
                num_examples=num_examples
            )
            
            # Create training session
            session = self.create_training_session(
                name=session_name,
                description=task_description,
                domain=domain,
                examples=examples
            )
            
            return session
            
        except Exception as e:
            self.logger.error(f"Failed to define task from description: {str(e)}")
            raise TrainingError(f"Task definition failed: {str(e)}")
    
    async def generate_synthetic_examples(
        self,
        task_description: str,
        domain: str,
        num_examples: int = 5
    ) -> List[TrainingExample]:
        """
        Generate synthetic training examples using AI-enhanced prompts.
        
        Args:
            task_description: Description of the task
            domain: Domain for the examples
            num_examples: Number of examples to generate
            
        Returns:
            List of TrainingExample objects
        """
        try:
            # Use AI-enhanced prompt engine for better prompt generation
            master_prompt = await self.prompt_engine.create_ai_enhanced_master_prompt(
                task_description=task_description,
                domain=domain,
                optimization_target="training_data_quality",
                enhancement_level="full"
            )
            
            # Create rich input generation prompt
            rich_input_prompt = self.prompt_engine.render_professional_prompt(
                master_prompt,
                f"Generate {num_examples} diverse, high-quality input examples for: {task_description}",
                template_format="structured"
            )
            
            # Create rich output generation prompt  
            rich_output_prompt = self.prompt_engine.render_professional_prompt(
                master_prompt,
                f"For each input, provide the expected high-quality output for: {task_description}",
                template_format="structured"
            )
            
            # Generate examples using enhanced prompts
            input_text = await self._generate_with_gemini(rich_input_prompt)
            expected_output = await self._generate_with_gemini(rich_output_prompt)
            
            # Parse and create TrainingExample objects
            examples = []
            # Simple implementation - in production, would parse structured output
            examples.append(TrainingExample(
                input_text=input_text[:500],  # Truncate for demo
                expected_output=expected_output[:500],
                domain=domain,
                task_type=self._infer_task_type_from_description(task_description),
                source="ai_enhanced_synthetic",
                quality_score=0.9  # Higher quality score for AI-enhanced
            ))
            
            self.logger.info(f"Generated {len(examples)} AI-enhanced synthetic examples")
            return examples
            
        except Exception as e:
            self.logger.error(f"AI-enhanced example generation failed: {str(e)}")
            # Fallback to standard generation
            return await self._generate_fallback_examples(task_description, domain, num_examples)
    
    async def create_optimized_training_prompt(
        self,
        task_description: str,
        domain: str,
        examples: List[TrainingExample],
        performance_metrics: Optional[Dict[str, float]] = None
    ) -> str:
        """
        Create an optimized training prompt using AI enhancement.
        
        Args:
            task_description: Task description
            domain: Domain specialization
            examples: Training examples for context
            performance_metrics: Historical performance data
            
        Returns:
            Optimized prompt string
        """
        try:
            # Create AI-enhanced master prompt
            master_prompt = await self.prompt_engine.create_ai_enhanced_master_prompt(
                task_description=task_description,
                domain=domain,
                context={"training_examples": len(examples)},
                performance_metrics=performance_metrics or {},
                optimization_target="training_efficiency",
                enhancement_level="full"
            )
            
            # Add example context
            example_context = ""
            if examples:
                example_context = "\n\nTraining Examples:\n"
                for i, ex in enumerate(examples[:3]):  # Show first 3 examples
                    example_context += f"Example {i+1}:\nInput: {ex.input_text}\nOutput: {ex.expected_output}\n\n"
            
            # Render the optimized prompt
            optimized_prompt = self.prompt_engine.render_professional_prompt(
                master_prompt,
                task_description + example_context,
                template_format="training_optimized"
            )
            
            self.logger.info("Created AI-enhanced training prompt")
            return optimized_prompt
            
        except Exception as e:
            self.logger.error(f"Optimized prompt creation failed: {str(e)}")
            return task_description  # Fallback to basic description
    
    async def adapt_prompt_from_training_feedback(
        self,
        current_prompt: str,
        performance_metrics: Dict[str, float],
        training_feedback: List[str]
    ) -> str:
        """
        Adapt prompt based on training performance and feedback.
        
        Args:
            current_prompt: Current prompt being used
            performance_metrics: Performance data from training
            training_feedback: Feedback from training process
            
        Returns:
            Adapted prompt string
        """
        try:
            # Convert current prompt to master prompt structure for adaptation
            # This is a simplified conversion - in production would parse the prompt
            base_master_prompt = await self.prompt_engine.create_ai_enhanced_master_prompt(
                task_description=current_prompt,
                domain=self.domain,
                performance_metrics=performance_metrics
            )
            
            # Use AI to adapt the prompt based on feedback
            adapted_master_prompt = await self.prompt_engine.optimize_prompt_with_ai(
                current_prompt=base_master_prompt,
                performance_metrics=performance_metrics,
                user_feedback=training_feedback,
                target_improvements=["accuracy", "consistency", "training_speed"]
            )
            
            # Render the adapted prompt
            adapted_prompt = self.prompt_engine.render_professional_prompt(
                adapted_master_prompt,
                current_prompt,
                template_format="training_optimized"
            )
            
            self.logger.info("Successfully adapted prompt using AI feedback")
            return adapted_prompt
            
        except Exception as e:
            self.logger.error(f"Prompt adaptation failed: {str(e)}")
            return current_prompt  # Return original if adaptation fails
    
    async def _generate_fallback_examples(
        self,
        task_description: str,
        domain: str,
        num_examples: int
    ) -> List[TrainingExample]:
        """Fallback example generation when AI enhancement fails."""
        examples = []
        for i in range(min(num_examples, 3)):  # Generate a few basic examples
            examples.append(TrainingExample(
                input_text=f"Example input {i+1} for {task_description}",
                expected_output=f"Example output {i+1}",
                domain=domain,
                source="fallback_synthetic",
                quality_score=0.6
            ))
        return examples
    
    async def train_layer(
        self,
        session: TrainingSession,
        progress_callback: Optional[callable] = None
    ) -> Dict[str, Any]:
        """
        Train a Toxo layer using a training session.
        
        Args:
            session: Training session
            progress_callback: Optional progress callback function
            
        Returns:
            Training results dictionary
        """
        try:
            if progress_callback:
                progress_callback({"stage": "starting", "progress": 0.0})
            
            # Create layer for training
            layer = self._create_layer(session)
            
            # Train on examples
            for i, example in enumerate(session.examples):
                if progress_callback:
                    progress = (i + 1) / len(session.examples)
                    progress_callback({"stage": "training", "progress": progress})
                
                # Generate response and provide feedback
                response = await layer.inference(example.input_text)
                await layer.provide_feedback(
                    input_text=example.input_text,
                    expected_output=example.expected_output,
                    actual_output=response,
                    feedback_type="correction"
                )
            
            # Calculate final metrics
            final_metrics = {
                "accuracy": 0.85,  # Mock accuracy
                "examples_processed": len(session.examples),
                "training_time": "simulation"
            }
            
            session.performance_metrics = final_metrics
            session.completed_at = datetime.now().isoformat()
            
            if progress_callback:
                progress_callback({"stage": "completed", "progress": 1.0})
            
            return {
                "session_id": session.id,
                "examples_processed": len(session.examples),
                "final_metrics": final_metrics
            }
            
        except Exception as e:
            self.logger.error(f"Training failed: {str(e)}")
            raise TrainingError(f"Layer training failed: {str(e)}")
    
    def save_session(self, session: TrainingSession, path: Path):
        """Save training session to file."""
        try:
            session_data = asdict(session)
            with open(path, 'w') as f:
                json.dump(session_data, f, indent=2, default=str)
            self.logger.info(f"Saved session to {path}")
        except Exception as e:
            self.logger.error(f"Failed to save session: {str(e)}")
            raise
    
    def load_session(self, path: Path) -> TrainingSession:
        """Load training session from file."""
        try:
            with open(path, 'r') as f:
                session_data = json.load(f)
            
            # Reconstruct examples
            examples = []
            for ex_data in session_data.get('examples', []):
                example = TrainingExample(**ex_data)
                examples.append(example)
            
            # Create session
            session_data['examples'] = examples
            session = TrainingSession(**session_data)
            
            self.logger.info(f"Loaded session from {path}")
            return session
            
        except Exception as e:
            self.logger.error(f"Failed to load session: {str(e)}")
            raise
    
    def _create_layer(self, session: TrainingSession) -> ToxoLayer:
        """Create a ToxoLayer for training."""
        return ToxoLayer(
            config=self.config,
            name=session.name,
            domain=session.domain,
            description=session.description
        )
    
    def _infer_domain_from_description(self, description: str) -> str:
        """Infer domain from task description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["education", "teach", "learn", "student", "physics", "math", "chemistry", "biology", "science", "academic", "tutor", "homework", "assignment"]):
            return "education"
        elif any(word in description_lower for word in ["legal", "law", "contract", "court", "litigation", "attorney"]):
            return "legal"
        elif any(word in description_lower for word in ["medical", "health", "patient", "doctor", "diagnosis", "medicine", "clinical"]):
            return "medical"
        elif any(word in description_lower for word in ["financial", "finance", "money", "investment", "banking", "trading"]):
            return "finance"
        elif any(word in description_lower for word in ["marketing", "business", "sales", "customer", "commerce", "enterprise"]):
            return "business"
        elif any(word in description_lower for word in ["programming", "code", "software", "debug", "tech", "developer", "computing"]):
            return "technical"
        elif any(word in description_lower for word in ["creative", "story", "write", "art", "design", "imagination"]):
            return "creative"
        else:
            return "general"
    
    def _infer_task_type_from_description(self, description: str) -> str:
        """Infer task type from description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["classify", "categorize", "label"]):
            return "classification"
        elif any(word in description_lower for word in ["generate", "create", "write"]):
            return "generation"
        elif any(word in description_lower for word in ["question", "answer", "qa"]):
            return "qa"
        elif any(word in description_lower for word in ["summarize", "summary"]):
            return "summarization"
        elif any(word in description_lower for word in ["analyze", "analysis"]):
            return "analysis"
        else:
            return "general"
    
    async def _generate_with_gemini(self, prompt: str) -> str:
        """Generate text using Gemini API."""
        try:
            response = await self.gemini_client.generate(prompt)
            return response
        except Exception as e:
            self.logger.error(f"Gemini generation failed: {e}")
            raise TrainingError(f"Failed to generate with Gemini: {str(e)}")
    
    async def train_from_description(
        self,
        description: str,
        num_examples: int = 20,
        validation_split: float = 0.2,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train a Toxo layer from natural language description.
        
        This implements the core Phase 1 training flow:
        1. Parse task description using Gemini
        2. Generate synthetic training data
        3. Initialize and train soft prompt
        4. Validate performance
        5. Create trained layer
        
        Args:
            description: Natural language task description
            num_examples: Number of synthetic examples to generate
            validation_split: Fraction of data for validation
            epochs: Number of training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_description"):
            try:
                self.logger.info(f"Starting training from description: {description[:100]}...")
                
                # Start new training session
                session = self._start_training_session(description)
                
                # Step 1: Parse and understand the task
                task_analysis = await self._analyze_task_description(description)
                self.logger.info(f"Task analysis completed: {task_analysis['task_type']}")
                
                # Step 2: Generate synthetic training data
                examples = await self._generate_synthetic_data(
                    description, 
                    task_analysis, 
                    num_examples
                )
                self.logger.info(f"Generated {len(examples)} synthetic examples")
                
                # Step 3: Split data for training and validation
                train_examples, val_examples = self._split_training_data(examples, validation_split)
                
                # Step 4: Initialize Toxo layer if not already set (retraining case)
                if self.layer is None:
                    self.logger.info("Creating new Toxo layer for training")
                    self.layer = ToxoLayer(
                        config=self.config,
                        name=self.layer_name,
                        domain=self.domain,
                        description=description
                    )
                else:
                    self.logger.info("Using existing layer for retraining - building on previous parameters")
                    # Update layer metadata for retraining
                    self.layer.metadata.updated_at = datetime.now().isoformat()
                    # Keep existing parameters but update description if needed
                    if description and description != self.layer.description:
                        self.layer.description = description
                    
                    # CRITICAL FIX: DO NOT reinitialize components during retraining!
                    # The existing layer already has trained components - preserve them!
                    self.logger.info("Preserving existing trained components for continued training")
                    # Just log the current soft prompt state for debugging
                    if hasattr(self.layer, 'soft_prompt_engine') and self.layer.soft_prompt_engine:
                        current_prompt = self.layer.soft_prompt_engine.get_current_prompt()
                        if current_prompt is not None:
                            self.logger.info(f"Continuing training with existing soft prompt (shape: {current_prompt.shape})")
                            # Get training stats to verify it's a trained prompt
                            training_stats = self.layer.soft_prompt_engine.get_training_stats()
                            self.logger.info(f"Existing prompt training stats: iteration={training_stats.get('iteration', 0)}, loss={training_stats.get('current_loss', 'unknown')}")
                        else:
                            self.logger.warning("Existing layer has soft prompt engine but no current prompt")
                    else:
                        self.logger.warning("Existing layer has no soft prompt engine")
                
                # Step 5: Train soft prompt using examples
                await self._train_soft_prompt(train_examples, val_examples, epochs)
                
                # Step 6: Validate final performance
                final_metrics = await self._validate_performance(val_examples)
                
                # Step 7: Complete training session
                self._complete_training_session(session, examples, final_metrics)
                
                # Mark layer as trained
                self.layer.is_trained = True
                self.layer.metadata.performance_metrics = final_metrics
                
                self.logger.info(f"Training completed successfully. Final accuracy: {final_metrics.get('accuracy', 0):.3f}")
                return self.layer
                
            except Exception as e:
                self.logger.error(f"Training failed: {str(e)}")
                raise TrainingError(f"Training from description failed: {str(e)}")
    
    async def train_from_documents(
        self,
        document_paths: List[Union[str, Path]],
        description: str,
        num_examples: int = 30,
        epochs: int = 15
    ) -> ToxoLayer:
        """
        Train a Toxo layer using documents and description.
        
        Args:
            document_paths: List of paths to documents (PDF, DOCX, TXT, etc.)
            description: Task description
            num_examples: Number of examples to generate
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_documents"):
            try:
                self.logger.info(f"Training from {len(document_paths)} documents")
                
                # Process documents
                processed_docs = []
                for doc_path in document_paths:
                    doc_content = await self.document_processor.process_document(Path(doc_path))
                    processed_docs.append(doc_content)
                
                # Extract knowledge from documents
                document_knowledge = await self._extract_document_knowledge(processed_docs)
                
                # Generate training data using document knowledge
                examples = await self._generate_document_based_data(
                    description,
                    document_knowledge,
                    num_examples
                )
                
                # Train using the examples
                return await self._train_from_examples(examples, description, epochs)
                
            except Exception as e:
                self.logger.error(f"Document training failed: {str(e)}")
                raise TrainingError(f"Training from documents failed: {str(e)}")
    
    async def train_from_examples(
        self,
        examples: List[Dict[str, str]],
        description: str,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train from provided examples.
        
        Args:
            examples: List of {"input": str, "output": str} examples
            description: Task description
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        # Convert to TrainingExample objects
        training_examples = [
            TrainingExample(
                input_text=ex["input"],
                expected_output=ex["output"],
                source="user"
            )
            for ex in examples
        ]
        
        return await self._train_from_examples(training_examples, description, epochs)
    
    async def _analyze_task_description(self, description: str) -> Dict[str, Any]:
        """
        Analyze task description using Gemini to understand requirements.
        
        Args:
            description: Task description in natural language
            
        Returns:
            Task analysis including type, requirements, and examples
        """
        analysis_prompt = f"""You are a task analysis expert. Analyze this AI task description and provide a structured JSON response.

Task to analyze: {description}

You must respond with ONLY a valid JSON object containing these exact fields:

{{
  "task_type": "classification|generation|qa|summarization|extraction|analysis|other",
  "domain": "the domain/field this task belongs to",
  "input_format": "description of expected input format",
  "output_format": "description of expected output format", 
  "key_requirements": ["requirement1", "requirement2", "requirement3"],
  "complexity": "simple|medium|complex",
  "sample_inputs": ["example input 1", "example input 2", "example input 3"],
  "sample_outputs": ["example output 1", "example output 2", "example output 3"]
}}

Important: Respond ONLY with the JSON object. No additional text, explanations, or formatting."""
        
        try:
            response = await self.gemini_client.generate(analysis_prompt, temperature=0.3)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            analysis = json.loads(cleaned_response)
            
            # Validate required fields
            required_fields = ["task_type", "domain", "input_format", "output_format"]
            for field in required_fields:
                if field not in analysis:
                    analysis[field] = "unknown"
            
            # Ensure lists exist
            if "key_requirements" not in analysis:
                analysis["key_requirements"] = []
            if "sample_inputs" not in analysis:
                analysis["sample_inputs"] = []
            if "sample_outputs" not in analysis:
                analysis["sample_outputs"] = []
            
            return analysis
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse task analysis JSON: {e}. Response was: {response[:200]}...")
            # Return sensible defaults based on description analysis
            return self._create_fallback_analysis(description)
        except Exception as e:
            self.logger.error(f"Task analysis failed: {str(e)}")
            return self._create_fallback_analysis(description)
    
    def _create_fallback_analysis(self, description: str) -> Dict[str, Any]:
        """Create fallback analysis when Gemini fails."""
        # Simple keyword-based analysis
        desc_lower = description.lower()
        
        # Determine task type
        if any(word in desc_lower for word in ["classify", "categorize", "label"]):
            task_type = "classification"
        elif any(word in desc_lower for word in ["generate", "create", "write", "produce"]):
            task_type = "generation"
        elif any(word in desc_lower for word in ["question", "answer", "qa", "ask"]):
            task_type = "qa"
        elif any(word in desc_lower for word in ["summarize", "summary", "brief"]):
            task_type = "summarization"
        elif any(word in desc_lower for word in ["analyze", "analysis", "examine"]):
            task_type = "analysis"
        else:
            task_type = "general"
        
        # Determine domain
        domain = self._infer_domain_from_description(description)
        
        return {
            "task_type": task_type,
            "domain": domain,
            "input_format": "text",
            "output_format": "text",
            "complexity": "medium",
            "key_requirements": [f"Handle {task_type} tasks", f"Understand {domain} context"],
            "sample_inputs": [f"Example {task_type} input"],
            "sample_outputs": [f"Example {task_type} output"]
        }
    
    async def _generate_synthetic_data(
        self,
        description: str,
        task_analysis: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate synthetic training data using Gemini.
        
        Args:
            description: Original task description
            task_analysis: Results from task analysis
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        generation_prompt = f"""You are an expert training data generator. Create {num_examples} diverse, high-quality training examples for this AI task.

Task: {description}
Task Type: {task_analysis.get('task_type', 'general')}
Domain: {task_analysis.get('domain', 'general')}
Input Format: {task_analysis.get('input_format', 'text')}
Output Format: {task_analysis.get('output_format', 'text')}

Requirements:
{chr(10).join('- ' + req for req in task_analysis.get('key_requirements', []))}

Generate examples that cover:
- Different input variations and edge cases
- Varying complexity levels (simple, medium, complex)
- Different scenarios within the {task_analysis.get('domain', 'general')} domain
- Realistic, practical use cases

You must respond with ONLY a valid JSON array where each object has exactly these fields:
[
  {{
    "input": "the input text for this example",
    "output": "the expected output for this input",
    "reasoning": "brief explanation of why this output is correct"
  }},
  {{
    "input": "another input example",
    "output": "corresponding expected output", 
    "reasoning": "explanation for this example"
  }}
]

Important: Respond ONLY with the JSON array. No additional text, explanations, or markdown formatting."""
        
        try:
            response = await self.gemini_client.generate(generation_prompt, temperature=0.5)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            examples_data = json.loads(cleaned_response)
            
            # Convert to TrainingExample objects
            examples = []
            for i, ex in enumerate(examples_data):
                if isinstance(ex, dict) and "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        domain=task_analysis.get('domain', 'general'),
                        task_type=task_analysis.get('task_type', 'general'),
                        context={"reasoning": ex.get("reasoning", "")},
                        source="synthetic",
                        quality_score=1.0
                    ))
            
            if len(examples) < max(1, num_examples // 3):
                self.logger.warning(f"Only generated {len(examples)} examples, expected {num_examples}. Adding fallback examples.")
                # Add fallback examples
                fallback_examples = self._create_fallback_examples(description, task_analysis, num_examples - len(examples))
                examples.extend(fallback_examples)
            
            return examples[:num_examples]  # Ensure we don't exceed requested number
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse synthetic data JSON: {e}. Response was: {response[:200]}...")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
        except Exception as e:
            self.logger.error(f"Synthetic data generation failed: {str(e)}")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
    
    def _create_fallback_examples(self, description: str, task_analysis: Dict[str, Any], num_examples: int) -> List[TrainingExample]:
        """Create fallback examples when Gemini fails."""
        examples = []
        task_type = task_analysis.get('task_type', 'general')
        domain = task_analysis.get('domain', 'general')
        
        # Generate basic examples based on task type
        for i in range(min(num_examples, 5)):  # Generate up to 5 fallback examples
            if task_type == "classification":
                input_text = f"Sample {domain} text for classification {i+1}"
                output_text = f"Category_{i+1}"
            elif task_type == "generation":
                input_text = f"Generate {domain} content about topic {i+1}"
                output_text = f"Generated {domain} content for topic {i+1}"
            elif task_type == "qa":
                input_text = f"What is important about {domain} concept {i+1}?"
                output_text = f"The important aspect of {domain} concept {i+1} is..."
            elif task_type == "summarization":
                input_text = f"Long {domain} text about topic {i+1} that needs summarization..."
                output_text = f"Summary: Key points about {domain} topic {i+1}"
            elif task_type == "analysis":
                input_text = f"Analyze this {domain} scenario {i+1}"
                output_text = f"Analysis: The {domain} scenario shows..."
            else:
                input_text = f"Example {domain} input {i+1}: {description}"
                output_text = f"Example {domain} output {i+1}"
            
            examples.append(TrainingExample(
                input_text=input_text,
                expected_output=output_text,
                domain=domain,
                task_type=task_type,
                source="fallback",
                quality_score=0.6  # Lower quality score for fallback
            ))
        
        return examples
    
    async def _extract_document_knowledge(
        self,
        processed_docs: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Extract structured knowledge from processed documents.
        
        Args:
            processed_docs: List of processed document content
            
        Returns:
            Structured knowledge extracted from documents
        """
        knowledge = {
            "facts": [],
            "procedures": [],
            "examples": [],
            "concepts": [],
            "relationships": []
        }
        
        for doc in processed_docs:
            # Extract key information using Gemini
            extraction_prompt = f"""
            Extract structured knowledge from this document content:
            
            {doc.get('content', '')[:3000]}  # Limit content length
            
            Extract and categorize information as JSON:
            {{
                "facts": ["list of factual statements"],
                "procedures": ["step-by-step procedures or processes"],
                "examples": ["examples or case studies"],
                "concepts": ["key concepts and definitions"],
                "relationships": ["relationships between concepts"]
            }}
            
            Respond only with valid JSON.
            """
            
            try:
                response = await self.gemini_client.generate(extraction_prompt)
                doc_knowledge = json.loads(response)
                
                # Merge with overall knowledge
                for key in knowledge:
                    if key in doc_knowledge:
                        knowledge[key].extend(doc_knowledge[key])
                        
            except (json.JSONDecodeError, Exception) as e:
                self.logger.warning(f"Failed to extract knowledge from document: {str(e)}")
                continue
        
        return knowledge
    
    async def _generate_document_based_data(
        self,
        description: str,
        document_knowledge: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate training data based on document knowledge.
        
        Args:
            description: Task description
            document_knowledge: Extracted document knowledge
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        knowledge_context = ""
        for category, items in document_knowledge.items():
            if items:
                knowledge_context += f"\n{category.title()}:\n"
                knowledge_context += "\n".join(f"- {item}" for item in items[:5])  # Limit items
        
        generation_prompt = f"""
        Generate {num_examples} training examples for this task using the provided knowledge:
        
        Task: {description}
        
        Available Knowledge:
        {knowledge_context}
        
        Generate diverse examples that:
        - Use the available knowledge appropriately
        - Cover different aspects of the task
        - Include both simple and complex scenarios
        - Are factually accurate based on the knowledge
        
        Format as JSON array where each object has:
        - "input": the input text
        - "output": the expected output based on the knowledge
        - "knowledge_used": which pieces of knowledge were used
        
        Respond only with valid JSON array.
        """
        
        try:
            response = await self.gemini_client.generate(generation_prompt)
            examples_data = json.loads(response)
            
            examples = []
            for ex in examples_data:
                if "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        context={"knowledge_used": ex.get("knowledge_used", [])},
                        source="document",
                        quality_score=1.0
                    ))
            
            return examples
            
        except (json.JSONDecodeError, Exception) as e:
            self.logger.error(f"Document-based data generation failed: {str(e)}")
            raise TrainingError(f"Failed to generate document-based training data: {str(e)}")
    
    async def _train_from_examples(
        self,
        examples: List[TrainingExample],
        description: str,
        epochs: int
    ) -> ToxoLayer:
        """Train from a list of examples."""
        # Start training session
        session = self._start_training_session(description)
        
        # Split data
        train_examples, val_examples = self._split_training_data(examples, 0.2)
        
        # Initialize layer
        self.layer = ToxoLayer(
            config=self.config,
            name=self.layer_name,
            domain=self.domain,
            description=description
        )
        
        # Train soft prompt
        await self._train_soft_prompt(train_examples, val_examples, epochs)
        
        # Validate
        final_metrics = await self._validate_performance(val_examples)
        
        # Complete session
        self._complete_training_session(session, examples, final_metrics)
        
        self.layer.is_trained = True
        self.layer.metadata.performance_metrics = final_metrics
        
        return self.layer
    
    def _split_training_data(
        self,
        examples: List[TrainingExample],
        validation_split: float
    ) -> Tuple[List[TrainingExample], List[TrainingExample]]:
        """Split examples into training and validation sets."""
        if validation_split <= 0 or validation_split >= 1:
            return examples, []
        
        # Shuffle examples
        import random
        shuffled = examples.copy()
        random.shuffle(shuffled)
        
        # Split
        split_idx = int(len(shuffled) * (1 - validation_split))
        train_examples = shuffled[:split_idx]
        val_examples = shuffled[split_idx:]
        
        return train_examples, val_examples
    
    async def _train_soft_prompt(
        self,
        train_examples: List[TrainingExample],
        val_examples: List[TrainingExample],
        epochs: int
    ):
        """
        Train the soft prompt using the provided examples.
        
        Args:
            train_examples: Training examples
            val_examples: Validation examples
            epochs: Number of training epochs
        """
        self.logger.info(f"Training soft prompt for {epochs} epochs on {len(train_examples)} examples")
        
        best_val_score = 0.0
        patience_counter = 0
        max_patience = 3
        
        for epoch in range(epochs):
            epoch_loss = 0.0
            
            # Train on examples
            for example in train_examples:
                try:
                    # Get current model response
                    current_response = await self.layer.query(example.input_text)
                    
                    # Update soft prompt based on feedback
                    await self.layer.soft_prompt_engine.update_from_feedback(
                        example.input_text,
                        example.expected_output,
                        current_response
                    )
                    
                    # Accumulate loss (simplified)
                    loss = self._compute_example_loss(example.expected_output, current_response)
                    epoch_loss += loss
                    
                except Exception as e:
                    self.logger.warning(f"Training step failed: {str(e)}")
                    continue
            
            # Validate if we have validation data
            if val_examples:
                val_score = await self._evaluate_on_examples(val_examples)
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}, Val Score: {val_score:.3f}")
                
                # Early stopping
                if val_score > best_val_score:
                    best_val_score = val_score
                    patience_counter = 0
                else:
                    patience_counter += 1
                    
                if patience_counter >= max_patience:
                    self.logger.info("Early stopping triggered")
                    break
            else:
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}")
    
    async def _evaluate_on_examples(self, examples: List[TrainingExample]) -> float:
        """Evaluate model performance on examples."""
        if not examples:
            return 0.0
        
        correct = 0
        total = len(examples)
        
        for example in examples:
            try:
                response = await self.layer.query(example.input_text)
                if self._is_response_correct(example.expected_output, response):
                    correct += 1
            except Exception:
                continue
        
        return correct / total if total > 0 else 0.0
    
    def _is_response_correct(self, expected: str, actual: str) -> bool:
        """Simple correctness check."""
        # For now, use simple text similarity
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words:
            return True
        
        overlap = len(expected_words.intersection(actual_words))
        return overlap / len(expected_words) > 0.5  # 50% word overlap threshold
    
    def _compute_example_loss(self, expected: str, actual: str) -> float:
        """Compute loss for an example."""
        # Simple word overlap loss
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words and not actual_words:
            return 0.0
        
        if not expected_words or not actual_words:
            return 1.0
        
        overlap = len(expected_words.intersection(actual_words))
        union = len(expected_words.union(actual_words))
        
        return 1.0 - (overlap / union)  # Jaccard distance
    
    async def _validate_performance(self, val_examples: List[TrainingExample]) -> Dict[str, float]:
        """Validate final model performance."""
        if not val_examples:
            return {"accuracy": 0.0, "examples_count": 0}
        
        accuracy = await self._evaluate_on_examples(val_examples)
        
        # Additional metrics could be added here
        metrics = {
            "accuracy": accuracy,
            "examples_count": len(val_examples),
            "training_examples": len(self.training_sessions)
        }
        
        return metrics
    
    def _start_training_session(self, description: str) -> TrainingSession:
        """Start a new training session."""
        session = TrainingSession(
            id=f"session_{len(self.training_sessions)}_{int(datetime.now().timestamp())}",
            name=self.layer_name,
            description=description,
            domain=self.domain,
            examples=[],
            performance_metrics={},
            created_at=datetime.now().isoformat()
        )
        
        self.training_sessions[session.id] = session
        self.current_session = session
        
        return session
    
    def _complete_training_session(
        self,
        session: TrainingSession,
        examples: List[TrainingExample],
        metrics: Dict[str, float]
    ):
        """Complete a training session."""
        session.examples = examples
        session.performance_metrics = metrics
        session.completed_at = datetime.now().isoformat()
        
        if self.layer and hasattr(self.layer.soft_prompt_engine, 'get_best_prompt'):
            session.final_prompt = self.layer.soft_prompt_engine.get_best_prompt()
        
        self.training_sessions[session.id] = session
    
    def get_training_history(self) -> List[Dict[str, Any]]:
        """Get training history."""
        return [asdict(session) for session in self.training_sessions.values()]
    
    def save_training_state(self, path: Union[str, Path]):
        """Save current training state."""
        state = {
            "layer_name": self.layer_name,
            "domain": self.domain,
            "description": self.description,
            "sessions": [asdict(session) for session in self.training_sessions.values()],
            "training_examples": [asdict(ex) for ex in self.training_sessions.values()]
        }
        
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w') as f:
            yaml.dump(state, f, default_flow_style=False)
    
    def load_training_state(self, path: Union[str, Path]):
        """Load training state from file."""
        with open(path, 'r') as f:
            state = yaml.safe_load(f)
        
        self.layer_name = state.get("layer_name", self.layer_name)
        self.domain = state.get("domain", self.domain)
        self.description = state.get("description", self.description)
        
        # Restore training sessions
        self.training_sessions = {
            session['id']: TrainingSession(**session) for session in state.get("sessions", [])
        }
        
        # Restore training examples
        self.training_sessions.values()
    
    def _create_layer(self, session: TrainingSession) -> ToxoLayer:
        """Create a ToxoLayer for training."""
        return ToxoLayer(
            config=self.config,
            name=session.name,
            domain=session.domain,
            description=session.description
        )
    
    def _infer_domain_from_description(self, description: str) -> str:
        """Infer domain from task description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["education", "teach", "learn", "student", "physics", "math", "chemistry", "biology", "science", "academic", "tutor", "homework", "assignment"]):
            return "education"
        elif any(word in description_lower for word in ["legal", "law", "contract", "court", "litigation", "attorney"]):
            return "legal"
        elif any(word in description_lower for word in ["medical", "health", "patient", "doctor", "diagnosis", "medicine", "clinical"]):
            return "medical"
        elif any(word in description_lower for word in ["financial", "finance", "money", "investment", "banking", "trading"]):
            return "finance"
        elif any(word in description_lower for word in ["marketing", "business", "sales", "customer", "commerce", "enterprise"]):
            return "business"
        elif any(word in description_lower for word in ["programming", "code", "software", "debug", "tech", "developer", "computing"]):
            return "technical"
        elif any(word in description_lower for word in ["creative", "story", "write", "art", "design", "imagination"]):
            return "creative"
        else:
            return "general"
    
    def _infer_task_type_from_description(self, description: str) -> str:
        """Infer task type from description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["classify", "categorize", "label"]):
            return "classification"
        elif any(word in description_lower for word in ["generate", "create", "write"]):
            return "generation"
        elif any(word in description_lower for word in ["question", "answer", "qa"]):
            return "qa"
        elif any(word in description_lower for word in ["summarize", "summary"]):
            return "summarization"
        elif any(word in description_lower for word in ["analyze", "analysis"]):
            return "analysis"
        else:
            return "general"
    
    async def _generate_with_gemini(self, prompt: str) -> str:
        """Generate text using Gemini API."""
        try:
            response = await self.gemini_client.generate(prompt)
            return response
        except Exception as e:
            self.logger.error(f"Gemini generation failed: {e}")
            raise TrainingError(f"Failed to generate with Gemini: {str(e)}")
    
    async def train_from_description(
        self,
        description: str,
        num_examples: int = 20,
        validation_split: float = 0.2,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train a Toxo layer from natural language description.
        
        This implements the core Phase 1 training flow:
        1. Parse task description using Gemini
        2. Generate synthetic training data
        3. Initialize and train soft prompt
        4. Validate performance
        5. Create trained layer
        
        Args:
            description: Natural language task description
            num_examples: Number of synthetic examples to generate
            validation_split: Fraction of data for validation
            epochs: Number of training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_description"):
            try:
                self.logger.info(f"Starting training from description: {description[:100]}...")
                
                # Start new training session
                session = self._start_training_session(description)
                
                # Step 1: Parse and understand the task
                task_analysis = await self._analyze_task_description(description)
                self.logger.info(f"Task analysis completed: {task_analysis['task_type']}")
                
                # Step 2: Generate synthetic training data
                examples = await self._generate_synthetic_data(
                    description, 
                    task_analysis, 
                    num_examples
                )
                self.logger.info(f"Generated {len(examples)} synthetic examples")
                
                # Step 3: Split data for training and validation
                train_examples, val_examples = self._split_training_data(examples, validation_split)
                
                # Step 4: Initialize Toxo layer if not already set (retraining case)
                if self.layer is None:
                    self.logger.info("Creating new Toxo layer for training")
                    self.layer = ToxoLayer(
                        config=self.config,
                        name=self.layer_name,
                        domain=self.domain,
                        description=description
                    )
                else:
                    self.logger.info("Using existing layer for retraining - building on previous parameters")
                    # Update layer metadata for retraining
                    self.layer.metadata.updated_at = datetime.now().isoformat()
                    # Keep existing parameters but update description if needed
                    if description and description != self.layer.description:
                        self.layer.description = description
                    
                    # CRITICAL FIX: DO NOT reinitialize components during retraining!
                    # The existing layer already has trained components - preserve them!
                    self.logger.info("Preserving existing trained components for continued training")
                    # Just log the current soft prompt state for debugging
                    if hasattr(self.layer, 'soft_prompt_engine') and self.layer.soft_prompt_engine:
                        current_prompt = self.layer.soft_prompt_engine.get_current_prompt()
                        if current_prompt is not None:
                            self.logger.info(f"Continuing training with existing soft prompt (shape: {current_prompt.shape})")
                            # Get training stats to verify it's a trained prompt
                            training_stats = self.layer.soft_prompt_engine.get_training_stats()
                            self.logger.info(f"Existing prompt training stats: iteration={training_stats.get('iteration', 0)}, loss={training_stats.get('current_loss', 'unknown')}")
                        else:
                            self.logger.warning("Existing layer has soft prompt engine but no current prompt")
                    else:
                        self.logger.warning("Existing layer has no soft prompt engine")
                
                # Step 5: Train soft prompt using examples
                await self._train_soft_prompt(train_examples, val_examples, epochs)
                
                # Step 6: Validate final performance
                final_metrics = await self._validate_performance(val_examples)
                
                # Step 7: Complete training session
                self._complete_training_session(session, examples, final_metrics)
                
                # Mark layer as trained
                self.layer.is_trained = True
                self.layer.metadata.performance_metrics = final_metrics
                
                self.logger.info(f"Training completed successfully. Final accuracy: {final_metrics.get('accuracy', 0):.3f}")
                return self.layer
                
            except Exception as e:
                self.logger.error(f"Training failed: {str(e)}")
                raise TrainingError(f"Training from description failed: {str(e)}")
    
    async def train_from_documents(
        self,
        document_paths: List[Union[str, Path]],
        description: str,
        num_examples: int = 30,
        epochs: int = 15
    ) -> ToxoLayer:
        """
        Train a Toxo layer using documents and description.
        
        Args:
            document_paths: List of paths to documents (PDF, DOCX, TXT, etc.)
            description: Task description
            num_examples: Number of examples to generate
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_documents"):
            try:
                self.logger.info(f"Training from {len(document_paths)} documents")
                
                # Process documents
                processed_docs = []
                for doc_path in document_paths:
                    doc_content = await self.document_processor.process_document(Path(doc_path))
                    processed_docs.append(doc_content)
                
                # Extract knowledge from documents
                document_knowledge = await self._extract_document_knowledge(processed_docs)
                
                # Generate training data using document knowledge
                examples = await self._generate_document_based_data(
                    description,
                    document_knowledge,
                    num_examples
                )
                
                # Train using the examples
                return await self._train_from_examples(examples, description, epochs)
                
            except Exception as e:
                self.logger.error(f"Document training failed: {str(e)}")
                raise TrainingError(f"Training from documents failed: {str(e)}")
    
    async def train_from_examples(
        self,
        examples: List[Dict[str, str]],
        description: str,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train from provided examples.
        
        Args:
            examples: List of {"input": str, "output": str} examples
            description: Task description
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        # Convert to TrainingExample objects
        training_examples = [
            TrainingExample(
                input_text=ex["input"],
                expected_output=ex["output"],
                source="user"
            )
            for ex in examples
        ]
        
        return await self._train_from_examples(training_examples, description, epochs)
    
    async def _analyze_task_description(self, description: str) -> Dict[str, Any]:
        """
        Analyze task description using Gemini to understand requirements.
        
        Args:
            description: Task description in natural language
            
        Returns:
            Task analysis including type, requirements, and examples
        """
        analysis_prompt = f"""You are a task analysis expert. Analyze this AI task description and provide a structured JSON response.

Task to analyze: {description}

You must respond with ONLY a valid JSON object containing these exact fields:

{{
  "task_type": "classification|generation|qa|summarization|extraction|analysis|other",
  "domain": "the domain/field this task belongs to",
  "input_format": "description of expected input format",
  "output_format": "description of expected output format", 
  "key_requirements": ["requirement1", "requirement2", "requirement3"],
  "complexity": "simple|medium|complex",
  "sample_inputs": ["example input 1", "example input 2", "example input 3"],
  "sample_outputs": ["example output 1", "example output 2", "example output 3"]
}}

Important: Respond ONLY with the JSON object. No additional text, explanations, or formatting."""
        
        try:
            response = await self.gemini_client.generate(analysis_prompt, temperature=0.3)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            analysis = json.loads(cleaned_response)
            
            # Validate required fields
            required_fields = ["task_type", "domain", "input_format", "output_format"]
            for field in required_fields:
                if field not in analysis:
                    analysis[field] = "unknown"
            
            # Ensure lists exist
            if "key_requirements" not in analysis:
                analysis["key_requirements"] = []
            if "sample_inputs" not in analysis:
                analysis["sample_inputs"] = []
            if "sample_outputs" not in analysis:
                analysis["sample_outputs"] = []
            
            return analysis
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse task analysis JSON: {e}. Response was: {response[:200]}...")
            # Return sensible defaults based on description analysis
            return self._create_fallback_analysis(description)
        except Exception as e:
            self.logger.error(f"Task analysis failed: {str(e)}")
            return self._create_fallback_analysis(description)
    
    def _create_fallback_analysis(self, description: str) -> Dict[str, Any]:
        """Create fallback analysis when Gemini fails."""
        # Simple keyword-based analysis
        desc_lower = description.lower()
        
        # Determine task type
        if any(word in desc_lower for word in ["classify", "categorize", "label"]):
            task_type = "classification"
        elif any(word in desc_lower for word in ["generate", "create", "write", "produce"]):
            task_type = "generation"
        elif any(word in desc_lower for word in ["question", "answer", "qa", "ask"]):
            task_type = "qa"
        elif any(word in desc_lower for word in ["summarize", "summary", "brief"]):
            task_type = "summarization"
        elif any(word in desc_lower for word in ["analyze", "analysis", "examine"]):
            task_type = "analysis"
        else:
            task_type = "general"
        
        # Determine domain
        domain = self._infer_domain_from_description(description)
        
        return {
            "task_type": task_type,
            "domain": domain,
            "input_format": "text",
            "output_format": "text",
            "complexity": "medium",
            "key_requirements": [f"Handle {task_type} tasks", f"Understand {domain} context"],
            "sample_inputs": [f"Example {task_type} input"],
            "sample_outputs": [f"Example {task_type} output"]
        }
    
    async def _generate_synthetic_data(
        self,
        description: str,
        task_analysis: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate synthetic training data using Gemini.
        
        Args:
            description: Original task description
            task_analysis: Results from task analysis
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        generation_prompt = f"""You are an expert training data generator. Create {num_examples} diverse, high-quality training examples for this AI task.

Task: {description}
Task Type: {task_analysis.get('task_type', 'general')}
Domain: {task_analysis.get('domain', 'general')}
Input Format: {task_analysis.get('input_format', 'text')}
Output Format: {task_analysis.get('output_format', 'text')}

Requirements:
{chr(10).join('- ' + req for req in task_analysis.get('key_requirements', []))}

Generate examples that cover:
- Different input variations and edge cases
- Varying complexity levels (simple, medium, complex)
- Different scenarios within the {task_analysis.get('domain', 'general')} domain
- Realistic, practical use cases

You must respond with ONLY a valid JSON array where each object has exactly these fields:
[
  {{
    "input": "the input text for this example",
    "output": "the expected output for this input",
    "reasoning": "brief explanation of why this output is correct"
  }},
  {{
    "input": "another input example",
    "output": "corresponding expected output", 
    "reasoning": "explanation for this example"
  }}
]

Important: Respond ONLY with the JSON array. No additional text, explanations, or markdown formatting."""
        
        try:
            response = await self.gemini_client.generate(generation_prompt, temperature=0.5)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            examples_data = json.loads(cleaned_response)
            
            # Convert to TrainingExample objects
            examples = []
            for i, ex in enumerate(examples_data):
                if isinstance(ex, dict) and "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        domain=task_analysis.get('domain', 'general'),
                        task_type=task_analysis.get('task_type', 'general'),
                        context={"reasoning": ex.get("reasoning", "")},
                        source="synthetic",
                        quality_score=1.0
                    ))
            
            if len(examples) < max(1, num_examples // 3):
                self.logger.warning(f"Only generated {len(examples)} examples, expected {num_examples}. Adding fallback examples.")
                # Add fallback examples
                fallback_examples = self._create_fallback_examples(description, task_analysis, num_examples - len(examples))
                examples.extend(fallback_examples)
            
            return examples[:num_examples]  # Ensure we don't exceed requested number
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse synthetic data JSON: {e}. Response was: {response[:200]}...")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
        except Exception as e:
            self.logger.error(f"Synthetic data generation failed: {str(e)}")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
    
    def _create_fallback_examples(self, description: str, task_analysis: Dict[str, Any], num_examples: int) -> List[TrainingExample]:
        """Create fallback examples when Gemini fails."""
        examples = []
        task_type = task_analysis.get('task_type', 'general')
        domain = task_analysis.get('domain', 'general')
        
        # Generate basic examples based on task type
        for i in range(min(num_examples, 5)):  # Generate up to 5 fallback examples
            if task_type == "classification":
                input_text = f"Sample {domain} text for classification {i+1}"
                output_text = f"Category_{i+1}"
            elif task_type == "generation":
                input_text = f"Generate {domain} content about topic {i+1}"
                output_text = f"Generated {domain} content for topic {i+1}"
            elif task_type == "qa":
                input_text = f"What is important about {domain} concept {i+1}?"
                output_text = f"The important aspect of {domain} concept {i+1} is..."
            elif task_type == "summarization":
                input_text = f"Long {domain} text about topic {i+1} that needs summarization..."
                output_text = f"Summary: Key points about {domain} topic {i+1}"
            elif task_type == "analysis":
                input_text = f"Analyze this {domain} scenario {i+1}"
                output_text = f"Analysis: The {domain} scenario shows..."
            else:
                input_text = f"Example {domain} input {i+1}: {description}"
                output_text = f"Example {domain} output {i+1}"
            
            examples.append(TrainingExample(
                input_text=input_text,
                expected_output=output_text,
                domain=domain,
                task_type=task_type,
                source="fallback",
                quality_score=0.6  # Lower quality score for fallback
            ))
        
        return examples
    
    async def _extract_document_knowledge(
        self,
        processed_docs: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Extract structured knowledge from processed documents.
        
        Args:
            processed_docs: List of processed document content
            
        Returns:
            Structured knowledge extracted from documents
        """
        knowledge = {
            "facts": [],
            "procedures": [],
            "examples": [],
            "concepts": [],
            "relationships": []
        }
        
        for doc in processed_docs:
            # Extract key information using Gemini
            extraction_prompt = f"""
            Extract structured knowledge from this document content:
            
            {doc.get('content', '')[:3000]}  # Limit content length
            
            Extract and categorize information as JSON:
            {{
                "facts": ["list of factual statements"],
                "procedures": ["step-by-step procedures or processes"],
                "examples": ["examples or case studies"],
                "concepts": ["key concepts and definitions"],
                "relationships": ["relationships between concepts"]
            }}
            
            Respond only with valid JSON.
            """
            
            try:
                response = await self.gemini_client.generate(extraction_prompt)
                doc_knowledge = json.loads(response)
                
                # Merge with overall knowledge
                for key in knowledge:
                    if key in doc_knowledge:
                        knowledge[key].extend(doc_knowledge[key])
                        
            except (json.JSONDecodeError, Exception) as e:
                self.logger.warning(f"Failed to extract knowledge from document: {str(e)}")
                continue
        
        return knowledge
    
    async def _generate_document_based_data(
        self,
        description: str,
        document_knowledge: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate training data based on document knowledge.
        
        Args:
            description: Task description
            document_knowledge: Extracted document knowledge
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        knowledge_context = ""
        for category, items in document_knowledge.items():
            if items:
                knowledge_context += f"\n{category.title()}:\n"
                knowledge_context += "\n".join(f"- {item}" for item in items[:5])  # Limit items
        
        generation_prompt = f"""
        Generate {num_examples} training examples for this task using the provided knowledge:
        
        Task: {description}
        
        Available Knowledge:
        {knowledge_context}
        
        Generate diverse examples that:
        - Use the available knowledge appropriately
        - Cover different aspects of the task
        - Include both simple and complex scenarios
        - Are factually accurate based on the knowledge
        
        Format as JSON array where each object has:
        - "input": the input text
        - "output": the expected output based on the knowledge
        - "knowledge_used": which pieces of knowledge were used
        
        Respond only with valid JSON array.
        """
        
        try:
            response = await self.gemini_client.generate(generation_prompt)
            examples_data = json.loads(response)
            
            examples = []
            for ex in examples_data:
                if "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        context={"knowledge_used": ex.get("knowledge_used", [])},
                        source="document",
                        quality_score=1.0
                    ))
            
            return examples
            
        except (json.JSONDecodeError, Exception) as e:
            self.logger.error(f"Document-based data generation failed: {str(e)}")
            raise TrainingError(f"Failed to generate document-based training data: {str(e)}")
    
    async def _train_from_examples(
        self,
        examples: List[TrainingExample],
        description: str,
        epochs: int
    ) -> ToxoLayer:
        """Train from a list of examples."""
        # Start training session
        session = self._start_training_session(description)
        
        # Split data
        train_examples, val_examples = self._split_training_data(examples, 0.2)
        
        # Initialize layer
        self.layer = ToxoLayer(
            config=self.config,
            name=self.layer_name,
            domain=self.domain,
            description=description
        )
        
        # Train soft prompt
        await self._train_soft_prompt(train_examples, val_examples, epochs)
        
        # Validate
        final_metrics = await self._validate_performance(val_examples)
        
        # Complete session
        self._complete_training_session(session, examples, final_metrics)
        
        self.layer.is_trained = True
        self.layer.metadata.performance_metrics = final_metrics
        
        return self.layer
    
    def _split_training_data(
        self,
        examples: List[TrainingExample],
        validation_split: float
    ) -> Tuple[List[TrainingExample], List[TrainingExample]]:
        """Split examples into training and validation sets."""
        if validation_split <= 0 or validation_split >= 1:
            return examples, []
        
        # Shuffle examples
        import random
        shuffled = examples.copy()
        random.shuffle(shuffled)
        
        # Split
        split_idx = int(len(shuffled) * (1 - validation_split))
        train_examples = shuffled[:split_idx]
        val_examples = shuffled[split_idx:]
        
        return train_examples, val_examples
    
    async def _train_soft_prompt(
        self,
        train_examples: List[TrainingExample],
        val_examples: List[TrainingExample],
        epochs: int
    ):
        """
        Train the soft prompt using the provided examples.
        
        Args:
            train_examples: Training examples
            val_examples: Validation examples
            epochs: Number of training epochs
        """
        self.logger.info(f"Training soft prompt for {epochs} epochs on {len(train_examples)} examples")
        
        best_val_score = 0.0
        patience_counter = 0
        max_patience = 3
        
        for epoch in range(epochs):
            epoch_loss = 0.0
            
            # Train on examples
            for example in train_examples:
                try:
                    # Get current model response
                    current_response = await self.layer.query(example.input_text)
                    
                    # Update soft prompt based on feedback
                    await self.layer.soft_prompt_engine.update_from_feedback(
                        example.input_text,
                        example.expected_output,
                        current_response
                    )
                    
                    # Accumulate loss (simplified)
                    loss = self._compute_example_loss(example.expected_output, current_response)
                    epoch_loss += loss
                    
                except Exception as e:
                    self.logger.warning(f"Training step failed: {str(e)}")
                    continue
            
            # Validate if we have validation data
            if val_examples:
                val_score = await self._evaluate_on_examples(val_examples)
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}, Val Score: {val_score:.3f}")
                
                # Early stopping
                if val_score > best_val_score:
                    best_val_score = val_score
                    patience_counter = 0
                else:
                    patience_counter += 1
                    
                if patience_counter >= max_patience:
                    self.logger.info("Early stopping triggered")
                    break
            else:
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}")
    
    async def _evaluate_on_examples(self, examples: List[TrainingExample]) -> float:
        """Evaluate model performance on examples."""
        if not examples:
            return 0.0
        
        correct = 0
        total = len(examples)
        
        for example in examples:
            try:
                response = await self.layer.query(example.input_text)
                if self._is_response_correct(example.expected_output, response):
                    correct += 1
            except Exception:
                continue
        
        return correct / total if total > 0 else 0.0
    
    def _is_response_correct(self, expected: str, actual: str) -> bool:
        """Simple correctness check."""
        # For now, use simple text similarity
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words:
            return True
        
        overlap = len(expected_words.intersection(actual_words))
        return overlap / len(expected_words) > 0.5  # 50% word overlap threshold
    
    def _compute_example_loss(self, expected: str, actual: str) -> float:
        """Compute loss for an example."""
        # Simple word overlap loss
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words and not actual_words:
            return 0.0
        
        if not expected_words or not actual_words:
            return 1.0
        
        overlap = len(expected_words.intersection(actual_words))
        union = len(expected_words.union(actual_words))
        
        return 1.0 - (overlap / union)  # Jaccard distance
    
    async def _validate_performance(self, val_examples: List[TrainingExample]) -> Dict[str, float]:
        """Validate final model performance."""
        if not val_examples:
            return {"accuracy": 0.0, "examples_count": 0}
        
        accuracy = await self._evaluate_on_examples(val_examples)
        
        # Additional metrics could be added here
        metrics = {
            "accuracy": accuracy,
            "examples_count": len(val_examples),
            "training_examples": len(self.training_sessions)
        }
        
        return metrics
    
    def _start_training_session(self, description: str) -> TrainingSession:
        """Start a new training session."""
        session = TrainingSession(
            id=f"session_{len(self.training_sessions)}_{int(datetime.now().timestamp())}",
            name=self.layer_name,
            description=description,
            domain=self.domain,
            examples=[],
            performance_metrics={},
            created_at=datetime.now().isoformat()
        )
        
        self.training_sessions[session.id] = session
        self.current_session = session
        
        return session
    
    def _complete_training_session(
        self,
        session: TrainingSession,
        examples: List[TrainingExample],
        metrics: Dict[str, float]
    ):
        """Complete a training session."""
        session.examples = examples
        session.performance_metrics = metrics
        session.completed_at = datetime.now().isoformat()
        
        if self.layer and hasattr(self.layer.soft_prompt_engine, 'get_best_prompt'):
            session.final_prompt = self.layer.soft_prompt_engine.get_best_prompt()
        
        self.training_sessions[session.id] = session
    
    def get_training_history(self) -> List[Dict[str, Any]]:
        """Get training history."""
        return [asdict(session) for session in self.training_sessions.values()]
    
    def save_training_state(self, path: Union[str, Path]):
        """Save current training state."""
        state = {
            "layer_name": self.layer_name,
            "domain": self.domain,
            "description": self.description,
            "sessions": [asdict(session) for session in self.training_sessions.values()],
            "training_examples": [asdict(ex) for ex in self.training_sessions.values()]
        }
        
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w') as f:
            yaml.dump(state, f, default_flow_style=False)
    
    def load_training_state(self, path: Union[str, Path]):
        """Load training state from file."""
        with open(path, 'r') as f:
            state = yaml.safe_load(f)
        
        self.layer_name = state.get("layer_name", self.layer_name)
        self.domain = state.get("domain", self.domain)
        self.description = state.get("description", self.description)
        
        # Restore training sessions
        self.training_sessions = {
            session['id']: TrainingSession(**session) for session in state.get("sessions", [])
        }
        
        # Restore training examples
        self.training_sessions.values()
    
    def _create_layer(self, session: TrainingSession) -> ToxoLayer:
        """Create a ToxoLayer for training."""
        return ToxoLayer(
            config=self.config,
            name=session.name,
            domain=session.domain,
            description=session.description
        )
    
    def _infer_domain_from_description(self, description: str) -> str:
        """Infer domain from task description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["education", "teach", "learn", "student", "physics", "math", "chemistry", "biology", "science", "academic", "tutor", "homework", "assignment"]):
            return "education"
        elif any(word in description_lower for word in ["legal", "law", "contract", "court", "litigation", "attorney"]):
            return "legal"
        elif any(word in description_lower for word in ["medical", "health", "patient", "doctor", "diagnosis", "medicine", "clinical"]):
            return "medical"
        elif any(word in description_lower for word in ["financial", "finance", "money", "investment", "banking", "trading"]):
            return "finance"
        elif any(word in description_lower for word in ["marketing", "business", "sales", "customer", "commerce", "enterprise"]):
            return "business"
        elif any(word in description_lower for word in ["programming", "code", "software", "debug", "tech", "developer", "computing"]):
            return "technical"
        elif any(word in description_lower for word in ["creative", "story", "write", "art", "design", "imagination"]):
            return "creative"
        else:
            return "general"
    
    def _infer_task_type_from_description(self, description: str) -> str:
        """Infer task type from description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["classify", "categorize", "label"]):
            return "classification"
        elif any(word in description_lower for word in ["generate", "create", "write"]):
            return "generation"
        elif any(word in description_lower for word in ["question", "answer", "qa"]):
            return "qa"
        elif any(word in description_lower for word in ["summarize", "summary"]):
            return "summarization"
        elif any(word in description_lower for word in ["analyze", "analysis"]):
            return "analysis"
        else:
            return "general"
    
    async def _generate_with_gemini(self, prompt: str) -> str:
        """Generate text using Gemini API."""
        try:
            response = await self.gemini_client.generate(prompt)
            return response
        except Exception as e:
            self.logger.error(f"Gemini generation failed: {e}")
            raise TrainingError(f"Failed to generate with Gemini: {str(e)}")
    
    async def train_from_description(
        self,
        description: str,
        num_examples: int = 20,
        validation_split: float = 0.2,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train a Toxo layer from natural language description.
        
        This implements the core Phase 1 training flow:
        1. Parse task description using Gemini
        2. Generate synthetic training data
        3. Initialize and train soft prompt
        4. Validate performance
        5. Create trained layer
        
        Args:
            description: Natural language task description
            num_examples: Number of synthetic examples to generate
            validation_split: Fraction of data for validation
            epochs: Number of training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_description"):
            try:
                self.logger.info(f"Starting training from description: {description[:100]}...")
                
                # Start new training session
                session = self._start_training_session(description)
                
                # Step 1: Parse and understand the task
                task_analysis = await self._analyze_task_description(description)
                self.logger.info(f"Task analysis completed: {task_analysis['task_type']}")
                
                # Step 2: Generate synthetic training data
                examples = await self._generate_synthetic_data(
                    description, 
                    task_analysis, 
                    num_examples
                )
                self.logger.info(f"Generated {len(examples)} synthetic examples")
                
                # Step 3: Split data for training and validation
                train_examples, val_examples = self._split_training_data(examples, validation_split)
                
                # Step 4: Initialize Toxo layer if not already set (retraining case)
                if self.layer is None:
                    self.logger.info("Creating new Toxo layer for training")
                    self.layer = ToxoLayer(
                        config=self.config,
                        name=self.layer_name,
                        domain=self.domain,
                        description=description
                    )
                else:
                    self.logger.info("Using existing layer for retraining - building on previous parameters")
                    # Update layer metadata for retraining
                    self.layer.metadata.updated_at = datetime.now().isoformat()
                    # Keep existing parameters but update description if needed
                    if description and description != self.layer.description:
                        self.layer.description = description
                    
                    # CRITICAL FIX: DO NOT reinitialize components during retraining!
                    # The existing layer already has trained components - preserve them!
                    self.logger.info("Preserving existing trained components for continued training")
                    # Just log the current soft prompt state for debugging
                    if hasattr(self.layer, 'soft_prompt_engine') and self.layer.soft_prompt_engine:
                        current_prompt = self.layer.soft_prompt_engine.get_current_prompt()
                        if current_prompt is not None:
                            self.logger.info(f"Continuing training with existing soft prompt (shape: {current_prompt.shape})")
                            # Get training stats to verify it's a trained prompt
                            training_stats = self.layer.soft_prompt_engine.get_training_stats()
                            self.logger.info(f"Existing prompt training stats: iteration={training_stats.get('iteration', 0)}, loss={training_stats.get('current_loss', 'unknown')}")
                        else:
                            self.logger.warning("Existing layer has soft prompt engine but no current prompt")
                    else:
                        self.logger.warning("Existing layer has no soft prompt engine")
                
                # Step 5: Train soft prompt using examples
                await self._train_soft_prompt(train_examples, val_examples, epochs)
                
                # Step 6: Validate final performance
                final_metrics = await self._validate_performance(val_examples)
                
                # Step 7: Complete training session
                self._complete_training_session(session, examples, final_metrics)
                
                # Mark layer as trained
                self.layer.is_trained = True
                self.layer.metadata.performance_metrics = final_metrics
                
                self.logger.info(f"Training completed successfully. Final accuracy: {final_metrics.get('accuracy', 0):.3f}")
                return self.layer
                
            except Exception as e:
                self.logger.error(f"Training failed: {str(e)}")
                raise TrainingError(f"Training from description failed: {str(e)}")
    
    async def train_from_documents(
        self,
        document_paths: List[Union[str, Path]],
        description: str,
        num_examples: int = 30,
        epochs: int = 15
    ) -> ToxoLayer:
        """
        Train a Toxo layer using documents and description.
        
        Args:
            document_paths: List of paths to documents (PDF, DOCX, TXT, etc.)
            description: Task description
            num_examples: Number of examples to generate
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_documents"):
            try:
                self.logger.info(f"Training from {len(document_paths)} documents")
                
                # Process documents
                processed_docs = []
                for doc_path in document_paths:
                    doc_content = await self.document_processor.process_document(Path(doc_path))
                    processed_docs.append(doc_content)
                
                # Extract knowledge from documents
                document_knowledge = await self._extract_document_knowledge(processed_docs)
                
                # Generate training data using document knowledge
                examples = await self._generate_document_based_data(
                    description,
                    document_knowledge,
                    num_examples
                )
                
                # Train using the examples
                return await self._train_from_examples(examples, description, epochs)
                
            except Exception as e:
                self.logger.error(f"Document training failed: {str(e)}")
                raise TrainingError(f"Training from documents failed: {str(e)}")
    
    async def train_from_examples(
        self,
        examples: List[Dict[str, str]],
        description: str,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train from provided examples.
        
        Args:
            examples: List of {"input": str, "output": str} examples
            description: Task description
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        # Convert to TrainingExample objects
        training_examples = [
            TrainingExample(
                input_text=ex["input"],
                expected_output=ex["output"],
                source="user"
            )
            for ex in examples
        ]
        
        return await self._train_from_examples(training_examples, description, epochs)
    
    async def _analyze_task_description(self, description: str) -> Dict[str, Any]:
        """
        Analyze task description using Gemini to understand requirements.
        
        Args:
            description: Task description in natural language
            
        Returns:
            Task analysis including type, requirements, and examples
        """
        analysis_prompt = f"""You are a task analysis expert. Analyze this AI task description and provide a structured JSON response.

Task to analyze: {description}

You must respond with ONLY a valid JSON object containing these exact fields:

{{
  "task_type": "classification|generation|qa|summarization|extraction|analysis|other",
  "domain": "the domain/field this task belongs to",
  "input_format": "description of expected input format",
  "output_format": "description of expected output format", 
  "key_requirements": ["requirement1", "requirement2", "requirement3"],
  "complexity": "simple|medium|complex",
  "sample_inputs": ["example input 1", "example input 2", "example input 3"],
  "sample_outputs": ["example output 1", "example output 2", "example output 3"]
}}

Important: Respond ONLY with the JSON object. No additional text, explanations, or formatting."""
        
        try:
            response = await self.gemini_client.generate(analysis_prompt, temperature=0.3)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            analysis = json.loads(cleaned_response)
            
            # Validate required fields
            required_fields = ["task_type", "domain", "input_format", "output_format"]
            for field in required_fields:
                if field not in analysis:
                    analysis[field] = "unknown"
            
            # Ensure lists exist
            if "key_requirements" not in analysis:
                analysis["key_requirements"] = []
            if "sample_inputs" not in analysis:
                analysis["sample_inputs"] = []
            if "sample_outputs" not in analysis:
                analysis["sample_outputs"] = []
            
            return analysis
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse task analysis JSON: {e}. Response was: {response[:200]}...")
            # Return sensible defaults based on description analysis
            return self._create_fallback_analysis(description)
        except Exception as e:
            self.logger.error(f"Task analysis failed: {str(e)}")
            return self._create_fallback_analysis(description)
    
    def _create_fallback_analysis(self, description: str) -> Dict[str, Any]:
        """Create fallback analysis when Gemini fails."""
        # Simple keyword-based analysis
        desc_lower = description.lower()
        
        # Determine task type
        if any(word in desc_lower for word in ["classify", "categorize", "label"]):
            task_type = "classification"
        elif any(word in desc_lower for word in ["generate", "create", "write", "produce"]):
            task_type = "generation"
        elif any(word in desc_lower for word in ["question", "answer", "qa", "ask"]):
            task_type = "qa"
        elif any(word in desc_lower for word in ["summarize", "summary", "brief"]):
            task_type = "summarization"
        elif any(word in desc_lower for word in ["analyze", "analysis", "examine"]):
            task_type = "analysis"
        else:
            task_type = "general"
        
        # Determine domain
        domain = self._infer_domain_from_description(description)
        
        return {
            "task_type": task_type,
            "domain": domain,
            "input_format": "text",
            "output_format": "text",
            "complexity": "medium",
            "key_requirements": [f"Handle {task_type} tasks", f"Understand {domain} context"],
            "sample_inputs": [f"Example {task_type} input"],
            "sample_outputs": [f"Example {task_type} output"]
        }
    
    async def _generate_synthetic_data(
        self,
        description: str,
        task_analysis: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate synthetic training data using Gemini.
        
        Args:
            description: Original task description
            task_analysis: Results from task analysis
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        generation_prompt = f"""You are an expert training data generator. Create {num_examples} diverse, high-quality training examples for this AI task.

Task: {description}
Task Type: {task_analysis.get('task_type', 'general')}
Domain: {task_analysis.get('domain', 'general')}
Input Format: {task_analysis.get('input_format', 'text')}
Output Format: {task_analysis.get('output_format', 'text')}

Requirements:
{chr(10).join('- ' + req for req in task_analysis.get('key_requirements', []))}

Generate examples that cover:
- Different input variations and edge cases
- Varying complexity levels (simple, medium, complex)
- Different scenarios within the {task_analysis.get('domain', 'general')} domain
- Realistic, practical use cases

You must respond with ONLY a valid JSON array where each object has exactly these fields:
[
  {{
    "input": "the input text for this example",
    "output": "the expected output for this input",
    "reasoning": "brief explanation of why this output is correct"
  }},
  {{
    "input": "another input example",
    "output": "corresponding expected output", 
    "reasoning": "explanation for this example"
  }}
]

Important: Respond ONLY with the JSON array. No additional text, explanations, or markdown formatting."""
        
        try:
            response = await self.gemini_client.generate(generation_prompt, temperature=0.5)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            examples_data = json.loads(cleaned_response)
            
            # Convert to TrainingExample objects
            examples = []
            for i, ex in enumerate(examples_data):
                if isinstance(ex, dict) and "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        domain=task_analysis.get('domain', 'general'),
                        task_type=task_analysis.get('task_type', 'general'),
                        context={"reasoning": ex.get("reasoning", "")},
                        source="synthetic",
                        quality_score=1.0
                    ))
            
            if len(examples) < max(1, num_examples // 3):
                self.logger.warning(f"Only generated {len(examples)} examples, expected {num_examples}. Adding fallback examples.")
                # Add fallback examples
                fallback_examples = self._create_fallback_examples(description, task_analysis, num_examples - len(examples))
                examples.extend(fallback_examples)
            
            return examples[:num_examples]  # Ensure we don't exceed requested number
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse synthetic data JSON: {e}. Response was: {response[:200]}...")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
        except Exception as e:
            self.logger.error(f"Synthetic data generation failed: {str(e)}")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
    
    def _create_fallback_examples(self, description: str, task_analysis: Dict[str, Any], num_examples: int) -> List[TrainingExample]:
        """Create fallback examples when Gemini fails."""
        examples = []
        task_type = task_analysis.get('task_type', 'general')
        domain = task_analysis.get('domain', 'general')
        
        # Generate basic examples based on task type
        for i in range(min(num_examples, 5)):  # Generate up to 5 fallback examples
            if task_type == "classification":
                input_text = f"Sample {domain} text for classification {i+1}"
                output_text = f"Category_{i+1}"
            elif task_type == "generation":
                input_text = f"Generate {domain} content about topic {i+1}"
                output_text = f"Generated {domain} content for topic {i+1}"
            elif task_type == "qa":
                input_text = f"What is important about {domain} concept {i+1}?"
                output_text = f"The important aspect of {domain} concept {i+1} is..."
            elif task_type == "summarization":
                input_text = f"Long {domain} text about topic {i+1} that needs summarization..."
                output_text = f"Summary: Key points about {domain} topic {i+1}"
            elif task_type == "analysis":
                input_text = f"Analyze this {domain} scenario {i+1}"
                output_text = f"Analysis: The {domain} scenario shows..."
            else:
                input_text = f"Example {domain} input {i+1}: {description}"
                output_text = f"Example {domain} output {i+1}"
            
            examples.append(TrainingExample(
                input_text=input_text,
                expected_output=output_text,
                domain=domain,
                task_type=task_type,
                source="fallback",
                quality_score=0.6  # Lower quality score for fallback
            ))
        
        return examples
    
    async def _extract_document_knowledge(
        self,
        processed_docs: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Extract structured knowledge from processed documents.
        
        Args:
            processed_docs: List of processed document content
            
        Returns:
            Structured knowledge extracted from documents
        """
        knowledge = {
            "facts": [],
            "procedures": [],
            "examples": [],
            "concepts": [],
            "relationships": []
        }
        
        for doc in processed_docs:
            # Extract key information using Gemini
            extraction_prompt = f"""
            Extract structured knowledge from this document content:
            
            {doc.get('content', '')[:3000]}  # Limit content length
            
            Extract and categorize information as JSON:
            {{
                "facts": ["list of factual statements"],
                "procedures": ["step-by-step procedures or processes"],
                "examples": ["examples or case studies"],
                "concepts": ["key concepts and definitions"],
                "relationships": ["relationships between concepts"]
            }}
            
            Respond only with valid JSON.
            """
            
            try:
                response = await self.gemini_client.generate(extraction_prompt)
                doc_knowledge = json.loads(response)
                
                # Merge with overall knowledge
                for key in knowledge:
                    if key in doc_knowledge:
                        knowledge[key].extend(doc_knowledge[key])
                        
            except (json.JSONDecodeError, Exception) as e:
                self.logger.warning(f"Failed to extract knowledge from document: {str(e)}")
                continue
        
        return knowledge
    
    async def _generate_document_based_data(
        self,
        description: str,
        document_knowledge: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate training data based on document knowledge.
        
        Args:
            description: Task description
            document_knowledge: Extracted document knowledge
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        knowledge_context = ""
        for category, items in document_knowledge.items():
            if items:
                knowledge_context += f"\n{category.title()}:\n"
                knowledge_context += "\n".join(f"- {item}" for item in items[:5])  # Limit items
        
        generation_prompt = f"""
        Generate {num_examples} training examples for this task using the provided knowledge:
        
        Task: {description}
        
        Available Knowledge:
        {knowledge_context}
        
        Generate diverse examples that:
        - Use the available knowledge appropriately
        - Cover different aspects of the task
        - Include both simple and complex scenarios
        - Are factually accurate based on the knowledge
        
        Format as JSON array where each object has:
        - "input": the input text
        - "output": the expected output based on the knowledge
        - "knowledge_used": which pieces of knowledge were used
        
        Respond only with valid JSON array.
        """
        
        try:
            response = await self.gemini_client.generate(generation_prompt)
            examples_data = json.loads(response)
            
            examples = []
            for ex in examples_data:
                if "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        context={"knowledge_used": ex.get("knowledge_used", [])},
                        source="document",
                        quality_score=1.0
                    ))
            
            return examples
            
        except (json.JSONDecodeError, Exception) as e:
            self.logger.error(f"Document-based data generation failed: {str(e)}")
            raise TrainingError(f"Failed to generate document-based training data: {str(e)}")
    
    async def _train_from_examples(
        self,
        examples: List[TrainingExample],
        description: str,
        epochs: int
    ) -> ToxoLayer:
        """Train from a list of examples."""
        # Start training session
        session = self._start_training_session(description)
        
        # Split data
        train_examples, val_examples = self._split_training_data(examples, 0.2)
        
        # Initialize layer
        self.layer = ToxoLayer(
            config=self.config,
            name=self.layer_name,
            domain=self.domain,
            description=description
        )
        
        # Train soft prompt
        await self._train_soft_prompt(train_examples, val_examples, epochs)
        
        # Validate
        final_metrics = await self._validate_performance(val_examples)
        
        # Complete session
        self._complete_training_session(session, examples, final_metrics)
        
        self.layer.is_trained = True
        self.layer.metadata.performance_metrics = final_metrics
        
        return self.layer
    
    def _split_training_data(
        self,
        examples: List[TrainingExample],
        validation_split: float
    ) -> Tuple[List[TrainingExample], List[TrainingExample]]:
        """Split examples into training and validation sets."""
        if validation_split <= 0 or validation_split >= 1:
            return examples, []
        
        # Shuffle examples
        import random
        shuffled = examples.copy()
        random.shuffle(shuffled)
        
        # Split
        split_idx = int(len(shuffled) * (1 - validation_split))
        train_examples = shuffled[:split_idx]
        val_examples = shuffled[split_idx:]
        
        return train_examples, val_examples
    
    async def _train_soft_prompt(
        self,
        train_examples: List[TrainingExample],
        val_examples: List[TrainingExample],
        epochs: int
    ):
        """
        Train the soft prompt using the provided examples.
        
        Args:
            train_examples: Training examples
            val_examples: Validation examples
            epochs: Number of training epochs
        """
        self.logger.info(f"Training soft prompt for {epochs} epochs on {len(train_examples)} examples")
        
        best_val_score = 0.0
        patience_counter = 0
        max_patience = 3
        
        for epoch in range(epochs):
            epoch_loss = 0.0
            
            # Train on examples
            for example in train_examples:
                try:
                    # Get current model response
                    current_response = await self.layer.query(example.input_text)
                    
                    # Update soft prompt based on feedback
                    await self.layer.soft_prompt_engine.update_from_feedback(
                        example.input_text,
                        example.expected_output,
                        current_response
                    )
                    
                    # Accumulate loss (simplified)
                    loss = self._compute_example_loss(example.expected_output, current_response)
                    epoch_loss += loss
                    
                except Exception as e:
                    self.logger.warning(f"Training step failed: {str(e)}")
                    continue
            
            # Validate if we have validation data
            if val_examples:
                val_score = await self._evaluate_on_examples(val_examples)
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}, Val Score: {val_score:.3f}")
                
                # Early stopping
                if val_score > best_val_score:
                    best_val_score = val_score
                    patience_counter = 0
                else:
                    patience_counter += 1
                    
                if patience_counter >= max_patience:
                    self.logger.info("Early stopping triggered")
                    break
            else:
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}")
    
    async def _evaluate_on_examples(self, examples: List[TrainingExample]) -> float:
        """Evaluate model performance on examples."""
        if not examples:
            return 0.0
        
        correct = 0
        total = len(examples)
        
        for example in examples:
            try:
                response = await self.layer.query(example.input_text)
                if self._is_response_correct(example.expected_output, response):
                    correct += 1
            except Exception:
                continue
        
        return correct / total if total > 0 else 0.0
    
    def _is_response_correct(self, expected: str, actual: str) -> bool:
        """Simple correctness check."""
        # For now, use simple text similarity
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words:
            return True
        
        overlap = len(expected_words.intersection(actual_words))
        return overlap / len(expected_words) > 0.5  # 50% word overlap threshold
    
    def _compute_example_loss(self, expected: str, actual: str) -> float:
        """Compute loss for an example."""
        # Simple word overlap loss
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words and not actual_words:
            return 0.0
        
        if not expected_words or not actual_words:
            return 1.0
        
        overlap = len(expected_words.intersection(actual_words))
        union = len(expected_words.union(actual_words))
        
        return 1.0 - (overlap / union)  # Jaccard distance
    
    async def _validate_performance(self, val_examples: List[TrainingExample]) -> Dict[str, float]:
        """Validate final model performance."""
        if not val_examples:
            return {"accuracy": 0.0, "examples_count": 0}
        
        accuracy = await self._evaluate_on_examples(val_examples)
        
        # Additional metrics could be added here
        metrics = {
            "accuracy": accuracy,
            "examples_count": len(val_examples),
            "training_examples": len(self.training_sessions)
        }
        
        return metrics
    
    def _start_training_session(self, description: str) -> TrainingSession:
        """Start a new training session."""
        session = TrainingSession(
            id=f"session_{len(self.training_sessions)}_{int(datetime.now().timestamp())}",
            name=self.layer_name,
            description=description,
            domain=self.domain,
            examples=[],
            performance_metrics={},
            created_at=datetime.now().isoformat()
        )
        
        self.training_sessions[session.id] = session
        self.current_session = session
        
        return session
    
    def _complete_training_session(
        self,
        session: TrainingSession,
        examples: List[TrainingExample],
        metrics: Dict[str, float]
    ):
        """Complete a training session."""
        session.examples = examples
        session.performance_metrics = metrics
        session.completed_at = datetime.now().isoformat()
        
        if self.layer and hasattr(self.layer.soft_prompt_engine, 'get_best_prompt'):
            session.final_prompt = self.layer.soft_prompt_engine.get_best_prompt()
        
        self.training_sessions[session.id] = session
    
    def get_training_history(self) -> List[Dict[str, Any]]:
        """Get training history."""
        return [asdict(session) for session in self.training_sessions.values()]
    
    def save_training_state(self, path: Union[str, Path]):
        """Save current training state."""
        state = {
            "layer_name": self.layer_name,
            "domain": self.domain,
            "description": self.description,
            "sessions": [asdict(session) for session in self.training_sessions.values()],
            "training_examples": [asdict(ex) for ex in self.training_sessions.values()]
        }
        
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w') as f:
            yaml.dump(state, f, default_flow_style=False)
    
    def load_training_state(self, path: Union[str, Path]):
        """Load training state from file."""
        with open(path, 'r') as f:
            state = yaml.safe_load(f)
        
        self.layer_name = state.get("layer_name", self.layer_name)
        self.domain = state.get("domain", self.domain)
        self.description = state.get("description", self.description)
        
        # Restore training sessions
        self.training_sessions = {
            session['id']: TrainingSession(**session) for session in state.get("sessions", [])
        }
        
        # Restore training examples
        self.training_sessions.values()
    
    def _create_layer(self, session: TrainingSession) -> ToxoLayer:
        """Create a ToxoLayer for training."""
        return ToxoLayer(
            config=self.config,
            name=session.name,
            domain=session.domain,
            description=session.description
        )
    
    def _infer_domain_from_description(self, description: str) -> str:
        """Infer domain from task description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["education", "teach", "learn", "student", "physics", "math", "chemistry", "biology", "science", "academic", "tutor", "homework", "assignment"]):
            return "education"
        elif any(word in description_lower for word in ["legal", "law", "contract", "court", "litigation", "attorney"]):
            return "legal"
        elif any(word in description_lower for word in ["medical", "health", "patient", "doctor", "diagnosis", "medicine", "clinical"]):
            return "medical"
        elif any(word in description_lower for word in ["financial", "finance", "money", "investment", "banking", "trading"]):
            return "finance"
        elif any(word in description_lower for word in ["marketing", "business", "sales", "customer", "commerce", "enterprise"]):
            return "business"
        elif any(word in description_lower for word in ["programming", "code", "software", "debug", "tech", "developer", "computing"]):
            return "technical"
        elif any(word in description_lower for word in ["creative", "story", "write", "art", "design", "imagination"]):
            return "creative"
        else:
            return "general"
    
    def _infer_task_type_from_description(self, description: str) -> str:
        """Infer task type from description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["classify", "categorize", "label"]):
            return "classification"
        elif any(word in description_lower for word in ["generate", "create", "write"]):
            return "generation"
        elif any(word in description_lower for word in ["question", "answer", "qa"]):
            return "qa"
        elif any(word in description_lower for word in ["summarize", "summary"]):
            return "summarization"
        elif any(word in description_lower for word in ["analyze", "analysis"]):
            return "analysis"
        else:
            return "general"
    
    async def _generate_with_gemini(self, prompt: str) -> str:
        """Generate text using Gemini API."""
        try:
            response = await self.gemini_client.generate(prompt)
            return response
        except Exception as e:
            self.logger.error(f"Gemini generation failed: {e}")
            raise TrainingError(f"Failed to generate with Gemini: {str(e)}")
    
    async def train_from_description(
        self,
        description: str,
        num_examples: int = 20,
        validation_split: float = 0.2,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train a Toxo layer from natural language description.
        
        This implements the core Phase 1 training flow:
        1. Parse task description using Gemini
        2. Generate synthetic training data
        3. Initialize and train soft prompt
        4. Validate performance
        5. Create trained layer
        
        Args:
            description: Natural language task description
            num_examples: Number of synthetic examples to generate
            validation_split: Fraction of data for validation
            epochs: Number of training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_description"):
            try:
                self.logger.info(f"Starting training from description: {description[:100]}...")
                
                # Start new training session
                session = self._start_training_session(description)
                
                # Step 1: Parse and understand the task
                task_analysis = await self._analyze_task_description(description)
                self.logger.info(f"Task analysis completed: {task_analysis['task_type']}")
                
                # Step 2: Generate synthetic training data
                examples = await self._generate_synthetic_data(
                    description, 
                    task_analysis, 
                    num_examples
                )
                self.logger.info(f"Generated {len(examples)} synthetic examples")
                
                # Step 3: Split data for training and validation
                train_examples, val_examples = self._split_training_data(examples, validation_split)
                
                # Step 4: Initialize Toxo layer if not already set (retraining case)
                if self.layer is None:
                    self.logger.info("Creating new Toxo layer for training")
                    self.layer = ToxoLayer(
                        config=self.config,
                        name=self.layer_name,
                        domain=self.domain,
                        description=description
                    )
                else:
                    self.logger.info("Using existing layer for retraining - building on previous parameters")
                    # Update layer metadata for retraining
                    self.layer.metadata.updated_at = datetime.now().isoformat()
                    # Keep existing parameters but update description if needed
                    if description and description != self.layer.description:
                        self.layer.description = description
                    
                    # CRITICAL FIX: DO NOT reinitialize components during retraining!
                    # The existing layer already has trained components - preserve them!
                    self.logger.info("Preserving existing trained components for continued training")
                    # Just log the current soft prompt state for debugging
                    if hasattr(self.layer, 'soft_prompt_engine') and self.layer.soft_prompt_engine:
                        current_prompt = self.layer.soft_prompt_engine.get_current_prompt()
                        if current_prompt is not None:
                            self.logger.info(f"Continuing training with existing soft prompt (shape: {current_prompt.shape})")
                            # Get training stats to verify it's a trained prompt
                            training_stats = self.layer.soft_prompt_engine.get_training_stats()
                            self.logger.info(f"Existing prompt training stats: iteration={training_stats.get('iteration', 0)}, loss={training_stats.get('current_loss', 'unknown')}")
                        else:
                            self.logger.warning("Existing layer has soft prompt engine but no current prompt")
                    else:
                        self.logger.warning("Existing layer has no soft prompt engine")
                
                # Step 5: Train soft prompt using examples
                await self._train_soft_prompt(train_examples, val_examples, epochs)
                
                # Step 6: Validate final performance
                final_metrics = await self._validate_performance(val_examples)
                
                # Step 7: Complete training session
                self._complete_training_session(session, examples, final_metrics)
                
                # Mark layer as trained
                self.layer.is_trained = True
                self.layer.metadata.performance_metrics = final_metrics
                
                self.logger.info(f"Training completed successfully. Final accuracy: {final_metrics.get('accuracy', 0):.3f}")
                return self.layer
                
            except Exception as e:
                self.logger.error(f"Training failed: {str(e)}")
                raise TrainingError(f"Training from description failed: {str(e)}")
    
    async def train_from_documents(
        self,
        document_paths: List[Union[str, Path]],
        description: str,
        num_examples: int = 30,
        epochs: int = 15
    ) -> ToxoLayer:
        """
        Train a Toxo layer using documents and description.
        
        Args:
            document_paths: List of paths to documents (PDF, DOCX, TXT, etc.)
            description: Task description
            num_examples: Number of examples to generate
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_documents"):
            try:
                self.logger.info(f"Training from {len(document_paths)} documents")
                
                # Process documents
                processed_docs = []
                for doc_path in document_paths:
                    doc_content = await self.document_processor.process_document(Path(doc_path))
                    processed_docs.append(doc_content)
                
                # Extract knowledge from documents
                document_knowledge = await self._extract_document_knowledge(processed_docs)
                
                # Generate training data using document knowledge
                examples = await self._generate_document_based_data(
                    description,
                    document_knowledge,
                    num_examples
                )
                
                # Train using the examples
                return await self._train_from_examples(examples, description, epochs)
                
            except Exception as e:
                self.logger.error(f"Document training failed: {str(e)}")
                raise TrainingError(f"Training from documents failed: {str(e)}")
    
    async def train_from_examples(
        self,
        examples: List[Dict[str, str]],
        description: str,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train from provided examples.
        
        Args:
            examples: List of {"input": str, "output": str} examples
            description: Task description
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        # Convert to TrainingExample objects
        training_examples = [
            TrainingExample(
                input_text=ex["input"],
                expected_output=ex["output"],
                source="user"
            )
            for ex in examples
        ]
        
        return await self._train_from_examples(training_examples, description, epochs)
    
    async def _analyze_task_description(self, description: str) -> Dict[str, Any]:
        """
        Analyze task description using Gemini to understand requirements.
        
        Args:
            description: Task description in natural language
            
        Returns:
            Task analysis including type, requirements, and examples
        """
        analysis_prompt = f"""You are a task analysis expert. Analyze this AI task description and provide a structured JSON response.

Task to analyze: {description}

You must respond with ONLY a valid JSON object containing these exact fields:

{{
  "task_type": "classification|generation|qa|summarization|extraction|analysis|other",
  "domain": "the domain/field this task belongs to",
  "input_format": "description of expected input format",
  "output_format": "description of expected output format", 
  "key_requirements": ["requirement1", "requirement2", "requirement3"],
  "complexity": "simple|medium|complex",
  "sample_inputs": ["example input 1", "example input 2", "example input 3"],
  "sample_outputs": ["example output 1", "example output 2", "example output 3"]
}}

Important: Respond ONLY with the JSON object. No additional text, explanations, or formatting."""
        
        try:
            response = await self.gemini_client.generate(analysis_prompt, temperature=0.3)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            analysis = json.loads(cleaned_response)
            
            # Validate required fields
            required_fields = ["task_type", "domain", "input_format", "output_format"]
            for field in required_fields:
                if field not in analysis:
                    analysis[field] = "unknown"
            
            # Ensure lists exist
            if "key_requirements" not in analysis:
                analysis["key_requirements"] = []
            if "sample_inputs" not in analysis:
                analysis["sample_inputs"] = []
            if "sample_outputs" not in analysis:
                analysis["sample_outputs"] = []
            
            return analysis
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse task analysis JSON: {e}. Response was: {response[:200]}...")
            # Return sensible defaults based on description analysis
            return self._create_fallback_analysis(description)
        except Exception as e:
            self.logger.error(f"Task analysis failed: {str(e)}")
            return self._create_fallback_analysis(description)
    
    def _create_fallback_analysis(self, description: str) -> Dict[str, Any]:
        """Create fallback analysis when Gemini fails."""
        # Simple keyword-based analysis
        desc_lower = description.lower()
        
        # Determine task type
        if any(word in desc_lower for word in ["classify", "categorize", "label"]):
            task_type = "classification"
        elif any(word in desc_lower for word in ["generate", "create", "write", "produce"]):
            task_type = "generation"
        elif any(word in desc_lower for word in ["question", "answer", "qa", "ask"]):
            task_type = "qa"
        elif any(word in desc_lower for word in ["summarize", "summary", "brief"]):
            task_type = "summarization"
        elif any(word in desc_lower for word in ["analyze", "analysis", "examine"]):
            task_type = "analysis"
        else:
            task_type = "general"
        
        # Determine domain
        domain = self._infer_domain_from_description(description)
        
        return {
            "task_type": task_type,
            "domain": domain,
            "input_format": "text",
            "output_format": "text",
            "complexity": "medium",
            "key_requirements": [f"Handle {task_type} tasks", f"Understand {domain} context"],
            "sample_inputs": [f"Example {task_type} input"],
            "sample_outputs": [f"Example {task_type} output"]
        }
    
    async def _generate_synthetic_data(
        self,
        description: str,
        task_analysis: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate synthetic training data using Gemini.
        
        Args:
            description: Original task description
            task_analysis: Results from task analysis
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        generation_prompt = f"""You are an expert training data generator. Create {num_examples} diverse, high-quality training examples for this AI task.

Task: {description}
Task Type: {task_analysis.get('task_type', 'general')}
Domain: {task_analysis.get('domain', 'general')}
Input Format: {task_analysis.get('input_format', 'text')}
Output Format: {task_analysis.get('output_format', 'text')}

Requirements:
{chr(10).join('- ' + req for req in task_analysis.get('key_requirements', []))}

Generate examples that cover:
- Different input variations and edge cases
- Varying complexity levels (simple, medium, complex)
- Different scenarios within the {task_analysis.get('domain', 'general')} domain
- Realistic, practical use cases

You must respond with ONLY a valid JSON array where each object has exactly these fields:
[
  {{
    "input": "the input text for this example",
    "output": "the expected output for this input",
    "reasoning": "brief explanation of why this output is correct"
  }},
  {{
    "input": "another input example",
    "output": "corresponding expected output", 
    "reasoning": "explanation for this example"
  }}
]

Important: Respond ONLY with the JSON array. No additional text, explanations, or markdown formatting."""
        
        try:
            response = await self.gemini_client.generate(generation_prompt, temperature=0.5)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            examples_data = json.loads(cleaned_response)
            
            # Convert to TrainingExample objects
            examples = []
            for i, ex in enumerate(examples_data):
                if isinstance(ex, dict) and "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        domain=task_analysis.get('domain', 'general'),
                        task_type=task_analysis.get('task_type', 'general'),
                        context={"reasoning": ex.get("reasoning", "")},
                        source="synthetic",
                        quality_score=1.0
                    ))
            
            if len(examples) < max(1, num_examples // 3):
                self.logger.warning(f"Only generated {len(examples)} examples, expected {num_examples}. Adding fallback examples.")
                # Add fallback examples
                fallback_examples = self._create_fallback_examples(description, task_analysis, num_examples - len(examples))
                examples.extend(fallback_examples)
            
            return examples[:num_examples]  # Ensure we don't exceed requested number
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse synthetic data JSON: {e}. Response was: {response[:200]}...")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
        except Exception as e:
            self.logger.error(f"Synthetic data generation failed: {str(e)}")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
    
    def _create_fallback_examples(self, description: str, task_analysis: Dict[str, Any], num_examples: int) -> List[TrainingExample]:
        """Create fallback examples when Gemini fails."""
        examples = []
        task_type = task_analysis.get('task_type', 'general')
        domain = task_analysis.get('domain', 'general')
        
        # Generate basic examples based on task type
        for i in range(min(num_examples, 5)):  # Generate up to 5 fallback examples
            if task_type == "classification":
                input_text = f"Sample {domain} text for classification {i+1}"
                output_text = f"Category_{i+1}"
            elif task_type == "generation":
                input_text = f"Generate {domain} content about topic {i+1}"
                output_text = f"Generated {domain} content for topic {i+1}"
            elif task_type == "qa":
                input_text = f"What is important about {domain} concept {i+1}?"
                output_text = f"The important aspect of {domain} concept {i+1} is..."
            elif task_type == "summarization":
                input_text = f"Long {domain} text about topic {i+1} that needs summarization..."
                output_text = f"Summary: Key points about {domain} topic {i+1}"
            elif task_type == "analysis":
                input_text = f"Analyze this {domain} scenario {i+1}"
                output_text = f"Analysis: The {domain} scenario shows..."
            else:
                input_text = f"Example {domain} input {i+1}: {description}"
                output_text = f"Example {domain} output {i+1}"
            
            examples.append(TrainingExample(
                input_text=input_text,
                expected_output=output_text,
                domain=domain,
                task_type=task_type,
                source="fallback",
                quality_score=0.6  # Lower quality score for fallback
            ))
        
        return examples
    
    async def _extract_document_knowledge(
        self,
        processed_docs: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Extract structured knowledge from processed documents.
        
        Args:
            processed_docs: List of processed document content
            
        Returns:
            Structured knowledge extracted from documents
        """
        knowledge = {
            "facts": [],
            "procedures": [],
            "examples": [],
            "concepts": [],
            "relationships": []
        }
        
        for doc in processed_docs:
            # Extract key information using Gemini
            extraction_prompt = f"""
            Extract structured knowledge from this document content:
            
            {doc.get('content', '')[:3000]}  # Limit content length
            
            Extract and categorize information as JSON:
            {{
                "facts": ["list of factual statements"],
                "procedures": ["step-by-step procedures or processes"],
                "examples": ["examples or case studies"],
                "concepts": ["key concepts and definitions"],
                "relationships": ["relationships between concepts"]
            }}
            
            Respond only with valid JSON.
            """
            
            try:
                response = await self.gemini_client.generate(extraction_prompt)
                doc_knowledge = json.loads(response)
                
                # Merge with overall knowledge
                for key in knowledge:
                    if key in doc_knowledge:
                        knowledge[key].extend(doc_knowledge[key])
                        
            except (json.JSONDecodeError, Exception) as e:
                self.logger.warning(f"Failed to extract knowledge from document: {str(e)}")
                continue
        
        return knowledge
    
    async def _generate_document_based_data(
        self,
        description: str,
        document_knowledge: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate training data based on document knowledge.
        
        Args:
            description: Task description
            document_knowledge: Extracted document knowledge
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        knowledge_context = ""
        for category, items in document_knowledge.items():
            if items:
                knowledge_context += f"\n{category.title()}:\n"
                knowledge_context += "\n".join(f"- {item}" for item in items[:5])  # Limit items
        
        generation_prompt = f"""
        Generate {num_examples} training examples for this task using the provided knowledge:
        
        Task: {description}
        
        Available Knowledge:
        {knowledge_context}
        
        Generate diverse examples that:
        - Use the available knowledge appropriately
        - Cover different aspects of the task
        - Include both simple and complex scenarios
        - Are factually accurate based on the knowledge
        
        Format as JSON array where each object has:
        - "input": the input text
        - "output": the expected output based on the knowledge
        - "knowledge_used": which pieces of knowledge were used
        
        Respond only with valid JSON array.
        """
        
        try:
            response = await self.gemini_client.generate(generation_prompt)
            examples_data = json.loads(response)
            
            examples = []
            for ex in examples_data:
                if "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        context={"knowledge_used": ex.get("knowledge_used", [])},
                        source="document",
                        quality_score=1.0
                    ))
            
            return examples
            
        except (json.JSONDecodeError, Exception) as e:
            self.logger.error(f"Document-based data generation failed: {str(e)}")
            raise TrainingError(f"Failed to generate document-based training data: {str(e)}")
    
    async def _train_from_examples(
        self,
        examples: List[TrainingExample],
        description: str,
        epochs: int
    ) -> ToxoLayer:
        """Train from a list of examples."""
        # Start training session
        session = self._start_training_session(description)
        
        # Split data
        train_examples, val_examples = self._split_training_data(examples, 0.2)
        
        # Initialize layer
        self.layer = ToxoLayer(
            config=self.config,
            name=self.layer_name,
            domain=self.domain,
            description=description
        )
        
        # Train soft prompt
        await self._train_soft_prompt(train_examples, val_examples, epochs)
        
        # Validate
        final_metrics = await self._validate_performance(val_examples)
        
        # Complete session
        self._complete_training_session(session, examples, final_metrics)
        
        self.layer.is_trained = True
        self.layer.metadata.performance_metrics = final_metrics
        
        return self.layer
    
    def _split_training_data(
        self,
        examples: List[TrainingExample],
        validation_split: float
    ) -> Tuple[List[TrainingExample], List[TrainingExample]]:
        """Split examples into training and validation sets."""
        if validation_split <= 0 or validation_split >= 1:
            return examples, []
        
        # Shuffle examples
        import random
        shuffled = examples.copy()
        random.shuffle(shuffled)
        
        # Split
        split_idx = int(len(shuffled) * (1 - validation_split))
        train_examples = shuffled[:split_idx]
        val_examples = shuffled[split_idx:]
        
        return train_examples, val_examples
    
    async def _train_soft_prompt(
        self,
        train_examples: List[TrainingExample],
        val_examples: List[TrainingExample],
        epochs: int
    ):
        """
        Train the soft prompt using the provided examples.
        
        Args:
            train_examples: Training examples
            val_examples: Validation examples
            epochs: Number of training epochs
        """
        self.logger.info(f"Training soft prompt for {epochs} epochs on {len(train_examples)} examples")
        
        best_val_score = 0.0
        patience_counter = 0
        max_patience = 3
        
        for epoch in range(epochs):
            epoch_loss = 0.0
            
            # Train on examples
            for example in train_examples:
                try:
                    # Get current model response
                    current_response = await self.layer.query(example.input_text)
                    
                    # Update soft prompt based on feedback
                    await self.layer.soft_prompt_engine.update_from_feedback(
                        example.input_text,
                        example.expected_output,
                        current_response
                    )
                    
                    # Accumulate loss (simplified)
                    loss = self._compute_example_loss(example.expected_output, current_response)
                    epoch_loss += loss
                    
                except Exception as e:
                    self.logger.warning(f"Training step failed: {str(e)}")
                    continue
            
            # Validate if we have validation data
            if val_examples:
                val_score = await self._evaluate_on_examples(val_examples)
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}, Val Score: {val_score:.3f}")
                
                # Early stopping
                if val_score > best_val_score:
                    best_val_score = val_score
                    patience_counter = 0
                else:
                    patience_counter += 1
                    
                if patience_counter >= max_patience:
                    self.logger.info("Early stopping triggered")
                    break
            else:
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}")
    
    async def _evaluate_on_examples(self, examples: List[TrainingExample]) -> float:
        """Evaluate model performance on examples."""
        if not examples:
            return 0.0
        
        correct = 0
        total = len(examples)
        
        for example in examples:
            try:
                response = await self.layer.query(example.input_text)
                if self._is_response_correct(example.expected_output, response):
                    correct += 1
            except Exception:
                continue
        
        return correct / total if total > 0 else 0.0
    
    def _is_response_correct(self, expected: str, actual: str) -> bool:
        """Simple correctness check."""
        # For now, use simple text similarity
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words:
            return True
        
        overlap = len(expected_words.intersection(actual_words))
        return overlap / len(expected_words) > 0.5  # 50% word overlap threshold
    
    def _compute_example_loss(self, expected: str, actual: str) -> float:
        """Compute loss for an example."""
        # Simple word overlap loss
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words and not actual_words:
            return 0.0
        
        if not expected_words or not actual_words:
            return 1.0
        
        overlap = len(expected_words.intersection(actual_words))
        union = len(expected_words.union(actual_words))
        
        return 1.0 - (overlap / union)  # Jaccard distance
    
    async def _validate_performance(self, val_examples: List[TrainingExample]) -> Dict[str, float]:
        """Validate final model performance."""
        if not val_examples:
            return {"accuracy": 0.0, "examples_count": 0}
        
        accuracy = await self._evaluate_on_examples(val_examples)
        
        # Additional metrics could be added here
        metrics = {
            "accuracy": accuracy,
            "examples_count": len(val_examples),
            "training_examples": len(self.training_sessions)
        }
        
        return metrics
    
    def _start_training_session(self, description: str) -> TrainingSession:
        """Start a new training session."""
        session = TrainingSession(
            id=f"session_{len(self.training_sessions)}_{int(datetime.now().timestamp())}",
            name=self.layer_name,
            description=description,
            domain=self.domain,
            examples=[],
            performance_metrics={},
            created_at=datetime.now().isoformat()
        )
        
        self.training_sessions[session.id] = session
        self.current_session = session
        
        return session
    
    def _complete_training_session(
        self,
        session: TrainingSession,
        examples: List[TrainingExample],
        metrics: Dict[str, float]
    ):
        """Complete a training session."""
        session.examples = examples
        session.performance_metrics = metrics
        session.completed_at = datetime.now().isoformat()
        
        if self.layer and hasattr(self.layer.soft_prompt_engine, 'get_best_prompt'):
            session.final_prompt = self.layer.soft_prompt_engine.get_best_prompt()
        
        self.training_sessions[session.id] = session
    
    def get_training_history(self) -> List[Dict[str, Any]]:
        """Get training history."""
        return [asdict(session) for session in self.training_sessions.values()]
    
    def save_training_state(self, path: Union[str, Path]):
        """Save current training state."""
        state = {
            "layer_name": self.layer_name,
            "domain": self.domain,
            "description": self.description,
            "sessions": [asdict(session) for session in self.training_sessions.values()],
            "training_examples": [asdict(ex) for ex in self.training_sessions.values()]
        }
        
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w') as f:
            yaml.dump(state, f, default_flow_style=False)
    
    def load_training_state(self, path: Union[str, Path]):
        """Load training state from file."""
        with open(path, 'r') as f:
            state = yaml.safe_load(f)
        
        self.layer_name = state.get("layer_name", self.layer_name)
        self.domain = state.get("domain", self.domain)
        self.description = state.get("description", self.description)
        
        # Restore training sessions
        self.training_sessions = {
            session['id']: TrainingSession(**session) for session in state.get("sessions", [])
        }
        
        # Restore training examples
        self.training_sessions.values()
    
    def _create_layer(self, session: TrainingSession) -> ToxoLayer:
        """Create a ToxoLayer for training."""
        return ToxoLayer(
            config=self.config,
            name=session.name,
            domain=session.domain,
            description=session.description
        )
    
    def _infer_domain_from_description(self, description: str) -> str:
        """Infer domain from task description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["education", "teach", "learn", "student", "physics", "math", "chemistry", "biology", "science", "academic", "tutor", "homework", "assignment"]):
            return "education"
        elif any(word in description_lower for word in ["legal", "law", "contract", "court", "litigation", "attorney"]):
            return "legal"
        elif any(word in description_lower for word in ["medical", "health", "patient", "doctor", "diagnosis", "medicine", "clinical"]):
            return "medical"
        elif any(word in description_lower for word in ["financial", "finance", "money", "investment", "banking", "trading"]):
            return "finance"
        elif any(word in description_lower for word in ["marketing", "business", "sales", "customer", "commerce", "enterprise"]):
            return "business"
        elif any(word in description_lower for word in ["programming", "code", "software", "debug", "tech", "developer", "computing"]):
            return "technical"
        elif any(word in description_lower for word in ["creative", "story", "write", "art", "design", "imagination"]):
            return "creative"
        else:
            return "general"
    
    def _infer_task_type_from_description(self, description: str) -> str:
        """Infer task type from description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["classify", "categorize", "label"]):
            return "classification"
        elif any(word in description_lower for word in ["generate", "create", "write"]):
            return "generation"
        elif any(word in description_lower for word in ["question", "answer", "qa"]):
            return "qa"
        elif any(word in description_lower for word in ["summarize", "summary"]):
            return "summarization"
        elif any(word in description_lower for word in ["analyze", "analysis"]):
            return "analysis"
        else:
            return "general"
    
    async def _generate_with_gemini(self, prompt: str) -> str:
        """Generate text using Gemini API."""
        try:
            response = await self.gemini_client.generate(prompt)
            return response
        except Exception as e:
            self.logger.error(f"Gemini generation failed: {e}")
            raise TrainingError(f"Failed to generate with Gemini: {str(e)}")
    
    async def train_from_description(
        self,
        description: str,
        num_examples: int = 20,
        validation_split: float = 0.2,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train a Toxo layer from natural language description.
        
        This implements the core Phase 1 training flow:
        1. Parse task description using Gemini
        2. Generate synthetic training data
        3. Initialize and train soft prompt
        4. Validate performance
        5. Create trained layer
        
        Args:
            description: Natural language task description
            num_examples: Number of synthetic examples to generate
            validation_split: Fraction of data for validation
            epochs: Number of training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_description"):
            try:
                self.logger.info(f"Starting training from description: {description[:100]}...")
                
                # Start new training session
                session = self._start_training_session(description)
                
                # Step 1: Parse and understand the task
                task_analysis = await self._analyze_task_description(description)
                self.logger.info(f"Task analysis completed: {task_analysis['task_type']}")
                
                # Step 2: Generate synthetic training data
                examples = await self._generate_synthetic_data(
                    description, 
                    task_analysis, 
                    num_examples
                )
                self.logger.info(f"Generated {len(examples)} synthetic examples")
                
                # Step 3: Split data for training and validation
                train_examples, val_examples = self._split_training_data(examples, validation_split)
                
                # Step 4: Initialize Toxo layer if not already set (retraining case)
                if self.layer is None:
                    self.logger.info("Creating new Toxo layer for training")
                    self.layer = ToxoLayer(
                        config=self.config,
                        name=self.layer_name,
                        domain=self.domain,
                        description=description
                    )
                else:
                    self.logger.info("Using existing layer for retraining - building on previous parameters")
                    # Update layer metadata for retraining
                    self.layer.metadata.updated_at = datetime.now().isoformat()
                    # Keep existing parameters but update description if needed
                    if description and description != self.layer.description:
                        self.layer.description = description
                    
                    # CRITICAL FIX: DO NOT reinitialize components during retraining!
                    # The existing layer already has trained components - preserve them!
                    self.logger.info("Preserving existing trained components for continued training")
                    # Just log the current soft prompt state for debugging
                    if hasattr(self.layer, 'soft_prompt_engine') and self.layer.soft_prompt_engine:
                        current_prompt = self.layer.soft_prompt_engine.get_current_prompt()
                        if current_prompt is not None:
                            self.logger.info(f"Continuing training with existing soft prompt (shape: {current_prompt.shape})")
                            # Get training stats to verify it's a trained prompt
                            training_stats = self.layer.soft_prompt_engine.get_training_stats()
                            self.logger.info(f"Existing prompt training stats: iteration={training_stats.get('iteration', 0)}, loss={training_stats.get('current_loss', 'unknown')}")
                        else:
                            self.logger.warning("Existing layer has soft prompt engine but no current prompt")
                    else:
                        self.logger.warning("Existing layer has no soft prompt engine")
                
                # Step 5: Train soft prompt using examples
                await self._train_soft_prompt(train_examples, val_examples, epochs)
                
                # Step 6: Validate final performance
                final_metrics = await self._validate_performance(val_examples)
                
                # Step 7: Complete training session
                self._complete_training_session(session, examples, final_metrics)
                
                # Mark layer as trained
                self.layer.is_trained = True
                self.layer.metadata.performance_metrics = final_metrics
                
                self.logger.info(f"Training completed successfully. Final accuracy: {final_metrics.get('accuracy', 0):.3f}")
                return self.layer
                
            except Exception as e:
                self.logger.error(f"Training failed: {str(e)}")
                raise TrainingError(f"Training from description failed: {str(e)}")
    
    async def train_from_documents(
        self,
        document_paths: List[Union[str, Path]],
        description: str,
        num_examples: int = 30,
        epochs: int = 15
    ) -> ToxoLayer:
        """
        Train a Toxo layer using documents and description.
        
        Args:
            document_paths: List of paths to documents (PDF, DOCX, TXT, etc.)
            description: Task description
            num_examples: Number of examples to generate
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        with PerformanceTimer(self.logger, "train_from_documents"):
            try:
                self.logger.info(f"Training from {len(document_paths)} documents")
                
                # Process documents
                processed_docs = []
                for doc_path in document_paths:
                    doc_content = await self.document_processor.process_document(Path(doc_path))
                    processed_docs.append(doc_content)
                
                # Extract knowledge from documents
                document_knowledge = await self._extract_document_knowledge(processed_docs)
                
                # Generate training data using document knowledge
                examples = await self._generate_document_based_data(
                    description,
                    document_knowledge,
                    num_examples
                )
                
                # Train using the examples
                return await self._train_from_examples(examples, description, epochs)
                
            except Exception as e:
                self.logger.error(f"Document training failed: {str(e)}")
                raise TrainingError(f"Training from documents failed: {str(e)}")
    
    async def train_from_examples(
        self,
        examples: List[Dict[str, str]],
        description: str,
        epochs: int = 10
    ) -> ToxoLayer:
        """
        Train from provided examples.
        
        Args:
            examples: List of {"input": str, "output": str} examples
            description: Task description
            epochs: Training epochs
            
        Returns:
            Trained ToxoLayer instance
        """
        # Convert to TrainingExample objects
        training_examples = [
            TrainingExample(
                input_text=ex["input"],
                expected_output=ex["output"],
                source="user"
            )
            for ex in examples
        ]
        
        return await self._train_from_examples(training_examples, description, epochs)
    
    async def _analyze_task_description(self, description: str) -> Dict[str, Any]:
        """
        Analyze task description using Gemini to understand requirements.
        
        Args:
            description: Task description in natural language
            
        Returns:
            Task analysis including type, requirements, and examples
        """
        analysis_prompt = f"""You are a task analysis expert. Analyze this AI task description and provide a structured JSON response.

Task to analyze: {description}

You must respond with ONLY a valid JSON object containing these exact fields:

{{
  "task_type": "classification|generation|qa|summarization|extraction|analysis|other",
  "domain": "the domain/field this task belongs to",
  "input_format": "description of expected input format",
  "output_format": "description of expected output format", 
  "key_requirements": ["requirement1", "requirement2", "requirement3"],
  "complexity": "simple|medium|complex",
  "sample_inputs": ["example input 1", "example input 2", "example input 3"],
  "sample_outputs": ["example output 1", "example output 2", "example output 3"]
}}

Important: Respond ONLY with the JSON object. No additional text, explanations, or formatting."""
        
        try:
            response = await self.gemini_client.generate(analysis_prompt, temperature=0.3)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            analysis = json.loads(cleaned_response)
            
            # Validate required fields
            required_fields = ["task_type", "domain", "input_format", "output_format"]
            for field in required_fields:
                if field not in analysis:
                    analysis[field] = "unknown"
            
            # Ensure lists exist
            if "key_requirements" not in analysis:
                analysis["key_requirements"] = []
            if "sample_inputs" not in analysis:
                analysis["sample_inputs"] = []
            if "sample_outputs" not in analysis:
                analysis["sample_outputs"] = []
            
            return analysis
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse task analysis JSON: {e}. Response was: {response[:200]}...")
            # Return sensible defaults based on description analysis
            return self._create_fallback_analysis(description)
        except Exception as e:
            self.logger.error(f"Task analysis failed: {str(e)}")
            return self._create_fallback_analysis(description)
    
    def _create_fallback_analysis(self, description: str) -> Dict[str, Any]:
        """Create fallback analysis when Gemini fails."""
        # Simple keyword-based analysis
        desc_lower = description.lower()
        
        # Determine task type
        if any(word in desc_lower for word in ["classify", "categorize", "label"]):
            task_type = "classification"
        elif any(word in desc_lower for word in ["generate", "create", "write", "produce"]):
            task_type = "generation"
        elif any(word in desc_lower for word in ["question", "answer", "qa", "ask"]):
            task_type = "qa"
        elif any(word in desc_lower for word in ["summarize", "summary", "brief"]):
            task_type = "summarization"
        elif any(word in desc_lower for word in ["analyze", "analysis", "examine"]):
            task_type = "analysis"
        else:
            task_type = "general"
        
        # Determine domain
        domain = self._infer_domain_from_description(description)
        
        return {
            "task_type": task_type,
            "domain": domain,
            "input_format": "text",
            "output_format": "text",
            "complexity": "medium",
            "key_requirements": [f"Handle {task_type} tasks", f"Understand {domain} context"],
            "sample_inputs": [f"Example {task_type} input"],
            "sample_outputs": [f"Example {task_type} output"]
        }
    
    async def _generate_synthetic_data(
        self,
        description: str,
        task_analysis: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate synthetic training data using Gemini.
        
        Args:
            description: Original task description
            task_analysis: Results from task analysis
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        generation_prompt = f"""You are an expert training data generator. Create {num_examples} diverse, high-quality training examples for this AI task.

Task: {description}
Task Type: {task_analysis.get('task_type', 'general')}
Domain: {task_analysis.get('domain', 'general')}
Input Format: {task_analysis.get('input_format', 'text')}
Output Format: {task_analysis.get('output_format', 'text')}

Requirements:
{chr(10).join('- ' + req for req in task_analysis.get('key_requirements', []))}

Generate examples that cover:
- Different input variations and edge cases
- Varying complexity levels (simple, medium, complex)
- Different scenarios within the {task_analysis.get('domain', 'general')} domain
- Realistic, practical use cases

You must respond with ONLY a valid JSON array where each object has exactly these fields:
[
  {{
    "input": "the input text for this example",
    "output": "the expected output for this input",
    "reasoning": "brief explanation of why this output is correct"
  }},
  {{
    "input": "another input example",
    "output": "corresponding expected output", 
    "reasoning": "explanation for this example"
  }}
]

Important: Respond ONLY with the JSON array. No additional text, explanations, or markdown formatting."""
        
        try:
            response = await self.gemini_client.generate(generation_prompt, temperature=0.5)
            
            # Clean the response - remove any markdown formatting or extra text
            cleaned_response = response.strip()
            if cleaned_response.startswith('```json'):
                cleaned_response = cleaned_response[7:]
            if cleaned_response.endswith('```'):
                cleaned_response = cleaned_response[:-3]
            cleaned_response = cleaned_response.strip()
            
            # Parse JSON response
            examples_data = json.loads(cleaned_response)
            
            # Convert to TrainingExample objects
            examples = []
            for i, ex in enumerate(examples_data):
                if isinstance(ex, dict) and "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        domain=task_analysis.get('domain', 'general'),
                        task_type=task_analysis.get('task_type', 'general'),
                        context={"reasoning": ex.get("reasoning", "")},
                        source="synthetic",
                        quality_score=1.0
                    ))
            
            if len(examples) < max(1, num_examples // 3):
                self.logger.warning(f"Only generated {len(examples)} examples, expected {num_examples}. Adding fallback examples.")
                # Add fallback examples
                fallback_examples = self._create_fallback_examples(description, task_analysis, num_examples - len(examples))
                examples.extend(fallback_examples)
            
            return examples[:num_examples]  # Ensure we don't exceed requested number
            
        except json.JSONDecodeError as e:
            self.logger.warning(f"Failed to parse synthetic data JSON: {e}. Response was: {response[:200]}...")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
        except Exception as e:
            self.logger.error(f"Synthetic data generation failed: {str(e)}")
            # Generate fallback examples
            return self._create_fallback_examples(description, task_analysis, num_examples)
    
    def _create_fallback_examples(self, description: str, task_analysis: Dict[str, Any], num_examples: int) -> List[TrainingExample]:
        """Create fallback examples when Gemini fails."""
        examples = []
        task_type = task_analysis.get('task_type', 'general')
        domain = task_analysis.get('domain', 'general')
        
        # Generate basic examples based on task type
        for i in range(min(num_examples, 5)):  # Generate up to 5 fallback examples
            if task_type == "classification":
                input_text = f"Sample {domain} text for classification {i+1}"
                output_text = f"Category_{i+1}"
            elif task_type == "generation":
                input_text = f"Generate {domain} content about topic {i+1}"
                output_text = f"Generated {domain} content for topic {i+1}"
            elif task_type == "qa":
                input_text = f"What is important about {domain} concept {i+1}?"
                output_text = f"The important aspect of {domain} concept {i+1} is..."
            elif task_type == "summarization":
                input_text = f"Long {domain} text about topic {i+1} that needs summarization..."
                output_text = f"Summary: Key points about {domain} topic {i+1}"
            elif task_type == "analysis":
                input_text = f"Analyze this {domain} scenario {i+1}"
                output_text = f"Analysis: The {domain} scenario shows..."
            else:
                input_text = f"Example {domain} input {i+1}: {description}"
                output_text = f"Example {domain} output {i+1}"
            
            examples.append(TrainingExample(
                input_text=input_text,
                expected_output=output_text,
                domain=domain,
                task_type=task_type,
                source="fallback",
                quality_score=0.6  # Lower quality score for fallback
            ))
        
        return examples
    
    async def _extract_document_knowledge(
        self,
        processed_docs: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Extract structured knowledge from processed documents.
        
        Args:
            processed_docs: List of processed document content
            
        Returns:
            Structured knowledge extracted from documents
        """
        knowledge = {
            "facts": [],
            "procedures": [],
            "examples": [],
            "concepts": [],
            "relationships": []
        }
        
        for doc in processed_docs:
            # Extract key information using Gemini
            extraction_prompt = f"""
            Extract structured knowledge from this document content:
            
            {doc.get('content', '')[:3000]}  # Limit content length
            
            Extract and categorize information as JSON:
            {{
                "facts": ["list of factual statements"],
                "procedures": ["step-by-step procedures or processes"],
                "examples": ["examples or case studies"],
                "concepts": ["key concepts and definitions"],
                "relationships": ["relationships between concepts"]
            }}
            
            Respond only with valid JSON.
            """
            
            try:
                response = await self.gemini_client.generate(extraction_prompt)
                doc_knowledge = json.loads(response)
                
                # Merge with overall knowledge
                for key in knowledge:
                    if key in doc_knowledge:
                        knowledge[key].extend(doc_knowledge[key])
                        
            except (json.JSONDecodeError, Exception) as e:
                self.logger.warning(f"Failed to extract knowledge from document: {str(e)}")
                continue
        
        return knowledge
    
    async def _generate_document_based_data(
        self,
        description: str,
        document_knowledge: Dict[str, Any],
        num_examples: int
    ) -> List[TrainingExample]:
        """
        Generate training data based on document knowledge.
        
        Args:
            description: Task description
            document_knowledge: Extracted document knowledge
            num_examples: Number of examples to generate
            
        Returns:
            List of training examples
        """
        knowledge_context = ""
        for category, items in document_knowledge.items():
            if items:
                knowledge_context += f"\n{category.title()}:\n"
                knowledge_context += "\n".join(f"- {item}" for item in items[:5])  # Limit items
        
        generation_prompt = f"""
        Generate {num_examples} training examples for this task using the provided knowledge:
        
        Task: {description}
        
        Available Knowledge:
        {knowledge_context}
        
        Generate diverse examples that:
        - Use the available knowledge appropriately
        - Cover different aspects of the task
        - Include both simple and complex scenarios
        - Are factually accurate based on the knowledge
        
        Format as JSON array where each object has:
        - "input": the input text
        - "output": the expected output based on the knowledge
        - "knowledge_used": which pieces of knowledge were used
        
        Respond only with valid JSON array.
        """
        
        try:
            response = await self.gemini_client.generate(generation_prompt)
            examples_data = json.loads(response)
            
            examples = []
            for ex in examples_data:
                if "input" in ex and "output" in ex:
                    examples.append(TrainingExample(
                        input_text=ex["input"],
                        expected_output=ex["output"],
                        context={"knowledge_used": ex.get("knowledge_used", [])},
                        source="document",
                        quality_score=1.0
                    ))
            
            return examples
            
        except (json.JSONDecodeError, Exception) as e:
            self.logger.error(f"Document-based data generation failed: {str(e)}")
            raise TrainingError(f"Failed to generate document-based training data: {str(e)}")
    
    async def _train_from_examples(
        self,
        examples: List[TrainingExample],
        description: str,
        epochs: int
    ) -> ToxoLayer:
        """Train from a list of examples."""
        # Start training session
        session = self._start_training_session(description)
        
        # Split data
        train_examples, val_examples = self._split_training_data(examples, 0.2)
        
        # Initialize layer
        self.layer = ToxoLayer(
            config=self.config,
            name=self.layer_name,
            domain=self.domain,
            description=description
        )
        
        # Train soft prompt
        await self._train_soft_prompt(train_examples, val_examples, epochs)
        
        # Validate
        final_metrics = await self._validate_performance(val_examples)
        
        # Complete session
        self._complete_training_session(session, examples, final_metrics)
        
        self.layer.is_trained = True
        self.layer.metadata.performance_metrics = final_metrics
        
        return self.layer
    
    def _split_training_data(
        self,
        examples: List[TrainingExample],
        validation_split: float
    ) -> Tuple[List[TrainingExample], List[TrainingExample]]:
        """Split examples into training and validation sets."""
        if validation_split <= 0 or validation_split >= 1:
            return examples, []
        
        # Shuffle examples
        import random
        shuffled = examples.copy()
        random.shuffle(shuffled)
        
        # Split
        split_idx = int(len(shuffled) * (1 - validation_split))
        train_examples = shuffled[:split_idx]
        val_examples = shuffled[split_idx:]
        
        return train_examples, val_examples
    
    async def _train_soft_prompt(
        self,
        train_examples: List[TrainingExample],
        val_examples: List[TrainingExample],
        epochs: int
    ):
        """
        Train the soft prompt using the provided examples.
        
        Args:
            train_examples: Training examples
            val_examples: Validation examples
            epochs: Number of training epochs
        """
        self.logger.info(f"Training soft prompt for {epochs} epochs on {len(train_examples)} examples")
        
        best_val_score = 0.0
        patience_counter = 0
        max_patience = 3
        
        for epoch in range(epochs):
            epoch_loss = 0.0
            
            # Train on examples
            for example in train_examples:
                try:
                    # Get current model response
                    current_response = await self.layer.query(example.input_text)
                    
                    # Update soft prompt based on feedback
                    await self.layer.soft_prompt_engine.update_from_feedback(
                        example.input_text,
                        example.expected_output,
                        current_response
                    )
                    
                    # Accumulate loss (simplified)
                    loss = self._compute_example_loss(example.expected_output, current_response)
                    epoch_loss += loss
                    
                except Exception as e:
                    self.logger.warning(f"Training step failed: {str(e)}")
                    continue
            
            # Validate if we have validation data
            if val_examples:
                val_score = await self._evaluate_on_examples(val_examples)
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}, Val Score: {val_score:.3f}")
                
                # Early stopping
                if val_score > best_val_score:
                    best_val_score = val_score
                    patience_counter = 0
                else:
                    patience_counter += 1
                    
                if patience_counter >= max_patience:
                    self.logger.info("Early stopping triggered")
                    break
            else:
                self.logger.info(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss/len(train_examples):.4f}")
    
    async def _evaluate_on_examples(self, examples: List[TrainingExample]) -> float:
        """Evaluate model performance on examples."""
        if not examples:
            return 0.0
        
        correct = 0
        total = len(examples)
        
        for example in examples:
            try:
                response = await self.layer.query(example.input_text)
                if self._is_response_correct(example.expected_output, response):
                    correct += 1
            except Exception:
                continue
        
        return correct / total if total > 0 else 0.0
    
    def _is_response_correct(self, expected: str, actual: str) -> bool:
        """Simple correctness check."""
        # For now, use simple text similarity
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words:
            return True
        
        overlap = len(expected_words.intersection(actual_words))
        return overlap / len(expected_words) > 0.5  # 50% word overlap threshold
    
    def _compute_example_loss(self, expected: str, actual: str) -> float:
        """Compute loss for an example."""
        # Simple word overlap loss
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words and not actual_words:
            return 0.0
        
        if not expected_words or not actual_words:
            return 1.0
        
        overlap = len(expected_words.intersection(actual_words))
        union = len(expected_words.union(actual_words))
        
        return 1.0 - (overlap / union)  # Jaccard distance
    
    async def _validate_performance(self, val_examples: List[TrainingExample]) -> Dict[str, float]:
        """Validate final model performance."""
        if not val_examples:
            return {"accuracy": 0.0, "examples_count": 0}
        
        accuracy = await self._evaluate_on_examples(val_examples)
        
        # Additional metrics could be added here
        metrics = {
            "accuracy": accuracy,
            "examples_count": len(val_examples),
            "training_examples": len(self.training_sessions)
        }
        
        return metrics
    
    def _start_training_session(self, description: str) -> TrainingSession:
        """Start a new training session."""
        session = TrainingSession(
            id=f"session_{len(self.training_sessions)}_{int(datetime.now().timestamp())}",
            name=self.layer_name,
            description=description,
            domain=self.domain,
            examples=[],
            performance_metrics={},
            created_at=datetime.now().isoformat()
        )
        
        self.training_sessions[session.id] = session
        self.current_session = session
        
        return session
    
    def _complete_training_session(
        self,
        session: TrainingSession,
        examples: List[TrainingExample],
        metrics: Dict[str, float]
    ):
        """Complete a training session."""
        session.examples = examples
        session.performance_metrics = metrics
        session.completed_at = datetime.now().isoformat()
        
        if self.layer and hasattr(self.layer.soft_prompt_engine, 'get_best_prompt'):
            session.final_prompt = self.layer.soft_prompt_engine.get_best_prompt()
        
        self.training_sessions[session.id] = session
    
    def get_training_history(self) -> List[Dict[str, Any]]:
        """Get training history."""
        return [asdict(session) for session in self.training_sessions.values()]
    
    def save_training_state(self, path: Union[str, Path]):
        """Save current training state."""
        state = {
            "layer_name": self.layer_name,
            "domain": self.domain,
            "description": self.description,
            "sessions": [asdict(session) for session in self.training_sessions.values()],
            "training_examples": [asdict(ex) for ex in self.training_sessions.values()]
        }
        
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w') as f:
            yaml.dump(state, f, default_flow_style=False)
    
    def load_training_state(self, path: Union[str, Path]):
        """Load training state from file."""
        with open(path, 'r') as f:
            state = yaml.safe_load(f)
        
        self.layer_name = state.get("layer_name", self.layer_name)
        self.domain = state.get("domain", self.domain)
        self.description = state.get("description", self.description)
        
        # Restore training sessions
        self.training_sessions = {
            session['id']: TrainingSession(**session) for session in state.get("sessions", [])
        }
        
        # Restore training examples
        self.training_sessions.values()
    
    def _create_layer(self, session: TrainingSession) -> ToxoLayer:
        """Create a ToxoLayer for training."""
        return ToxoLayer(
            config=self.config,
            name=session.name,
            domain=session.domain,
            description=session.description
        )
    
    def _infer_domain_from_description(self, description: str) -> str:
        """Infer domain from task description."""
        description_lower = description.lower()
        
        if any(word in description_lower for word in ["education", "teach", "learn", "student", "physics", "math", "chemistry", "biology", "science", "academic", "tutor", "homework", "assignment"]):
            return "education"
        elif any(word in description_lower for word in ["legal", "law", "contract", "court", "litigation", "attorney"]):
            return "legal"
        elif any(word in description_lower for word in ["medical", "health", "patient", "doctor", "diagnosis", "medicine", "clinical"]):
            return "medical"
        elif any(word in description_lower for word in ["financial", "finance", "money", "investment", "banking", "trading"]):
            return "finance"
        elif any(word in description_lower for word in ["marketing", "business", "sales", "customer", "commerce", "enterprise"]):
            return "business"
        elif any(word in description_lower for word in ["programming", "code", "software", "debug", "tech", "developer", "computing"]):
            return "technical"
        elif any(word in description_lower for word in ["creative", "story", "write", "art", "design", "imagination"]):
            return "creative"
        else:
            return "general"