backend/experiments/experiment_manager.py

"""
Experiment Manager for AegisLM Framework

Production-grade experiment lifecycle management with full tracking.
"""

import uuid
from datetime import datetime
from typing import Dict, Any, Optional, List
from pathlib import Path
import sys

# Add parent directory to path for imports
current_dir = Path(__file__).parent
backend_dir = current_dir.parent
if str(backend_dir) not in sys.path:
    sys.path.insert(0, str(backend_dir))

from schemas.experiment_schema import (
    Experiment, ExperimentCreate, ExperimentUpdate, ExperimentStatus,
    ExperimentPriority, ConfigSnapshot, ResultSummary, ExperimentFilter,
    ExperimentList, ExperimentComparison, ExperimentStats
)
from experiments.experiment_store import ExperimentStore


class ExperimentManager:
    """
    Production-grade experiment manager for lifecycle management.
    
    Handles experiment creation, status updates, result storage, and retrieval.
    """
    
    def __init__(self, store: Optional[ExperimentStore] = None):
        """
        Initialize experiment manager.
        
        Args:
            store: Optional experiment store instance
        """
        self.store = store or ExperimentStore()
    
    def create_experiment(self, config: ExperimentCreate) -> Experiment:
        """
        Create a new experiment with unique run_id.
        
        Args:
            config: Experiment creation configuration
            
        Returns:
            Created experiment with unique run_id
            
        Raises:
            ValueError: If configuration is invalid
        """
        # Generate unique run_id
        run_id = uuid.uuid4()
        
        # Extract quick-filter fields from config snapshot
        config_snapshot = config.config_snapshot
        
        # Create experiment
        experiment = Experiment(
            run_id=run_id,
            experiment_name=config.experiment_name,
            description=config.description,
            config_snapshot=config_snapshot,
            
            # Quick-filter fields
            model_name=config_snapshot.model_name,
            dataset_name=config_snapshot.dataset_name,
            dataset_version=config_snapshot.dataset_version,
            attack_types=config_snapshot.attack_types,
            prompt_count=config_snapshot.prompt_count,
            
            # Status and metadata
            status=ExperimentStatus.PENDING,
            priority=config.priority,
            tags=config.tags,
            created_by=config.created_by,
            parent_experiment_id=config.parent_experiment_id
        )
        
        # Store experiment
        stored_experiment = self.store.create_experiment(experiment)
        
        return stored_experiment
    
    def update_status(self, run_id: str, status: ExperimentStatus) -> Experiment:
        """
        Update experiment status.
        
        Args:
            run_id: Unique experiment identifier
            status: New status
            
        Returns:
            Updated experiment
            
        Raises:
            ValueError: If experiment not found or status transition invalid
        """
        # Get existing experiment
        experiment = self.store.get_experiment(run_id)
        if not experiment:
            raise ValueError(f"Experiment with run_id {run_id} not found")
        
        # Validate status transition
        self._validate_status_transition(experiment.status, status)
        
        # Create update
        update_data = ExperimentUpdate(status=status)
        
        # Handle status-specific logic
        if status == ExperimentStatus.RUNNING and experiment.status != ExperimentStatus.RUNNING:
            experiment.mark_as_running()
            update_data.status = ExperimentStatus.RUNNING
        elif status == ExperimentStatus.COMPLETED:
            if not experiment.result_summary:
                raise ValueError("Cannot mark experiment as completed without results")
            experiment.mark_as_completed(experiment.result_summary, experiment.full_result)
        elif status == ExperimentStatus.FAILED:
            if not experiment.error_message:
                raise ValueError("Cannot mark experiment as failed without error message")
            experiment.mark_as_failed(experiment.error_message)
        elif status == ExperimentStatus.CANCELLED:
            experiment.mark_as_cancelled()
        
        # Update experiment
        updated_experiment = self.store.update_experiment(run_id, update_data)
        
        return updated_experiment
    
    def save_results(self, run_id: str, result_summary: ResultSummary, 
                    full_result: Optional[Dict[str, Any]] = None) -> Experiment:
        """
        Save experiment results and mark as completed.
        
        Args:
            run_id: Unique experiment identifier
            result_summary: Summary of experiment results
            full_result: Complete result data (optional)
            
        Returns:
            Updated experiment with results
            
        Raises:
            ValueError: If experiment not found or not in running state
        """
        # Get existing experiment
        experiment = self.store.get_experiment(run_id)
        if not experiment:
            raise ValueError(f"Experiment with run_id {run_id} not found")
        
        # Validate experiment state
        if experiment.status not in [ExperimentStatus.RUNNING, ExperimentStatus.PENDING]:
            raise ValueError(f"Cannot save results for experiment in {experiment.status} status")
        
        # Create update with results
        update_data = ExperimentUpdate(
            result_summary=result_summary,
            full_result=full_result
        )
        
        # Mark as completed
        experiment.mark_as_completed(result_summary, full_result)
        update_data.status = ExperimentStatus.COMPLETED
        
        # Update experiment
        updated_experiment = self.store.update_experiment(run_id, update_data)
        
        return updated_experiment
    
    def mark_as_failed(self, run_id: str, error_message: str) -> Experiment:
        """
        Mark experiment as failed with error message.
        
        Args:
            run_id: Unique experiment identifier
            error_message: Error message describing failure
            
        Returns:
            Updated experiment marked as failed
            
        Raises:
            ValueError: If experiment not found
        """
        # Get existing experiment
        experiment = self.store.get_experiment(run_id)
        if not experiment:
            raise ValueError(f"Experiment with run_id {run_id} not found")
        
        # Create update with error
        update_data = ExperimentUpdate(error_message=error_message)
        
        # Mark as failed
        experiment.mark_as_failed(error_message)
        update_data.status = ExperimentStatus.FAILED
        
        # Update experiment
        updated_experiment = self.store.update_experiment(run_id, update_data)
        
        return updated_experiment
    
    def get_experiment(self, run_id: str) -> Optional[Experiment]:
        """
        Get experiment by run_id.
        
        Args:
            run_id: Unique experiment identifier
            
        Returns:
            Experiment if found, None otherwise
        """
        return self.store.get_experiment(run_id)
    
    def list_experiments(self, filters: Optional[ExperimentFilter] = None) -> ExperimentList:
        """
        List experiments with optional filtering.
        
        Args:
            filters: Optional filters to apply
            
        Returns:
            List of experiments matching filters
        """
        return self.store.list_experiments(filters or ExperimentFilter())
    
    def delete_experiment(self, run_id: str) -> bool:
        """
        Delete experiment by run_id.
        
        Args:
            run_id: Unique experiment identifier
            
        Returns:
            True if deleted, False if not found
        """
        return self.store.delete_experiment(run_id)
    
    def compare_experiments(self, run_id_1: str, run_id_2: str) -> ExperimentComparison:
        """
        Compare two experiments.
        
        Args:
            run_id_1: First experiment ID
            run_id_2: Second experiment ID
            
        Returns:
            Comparison of the two experiments
            
        Raises:
            ValueError: If either experiment not found
        """
        # Get experiments
        experiment_1 = self.get_experiment(run_id_1)
        experiment_2 = self.get_experiment(run_id_2)
        
        if not experiment_1:
            raise ValueError(f"Experiment with run_id {run_id_1} not found")
        if not experiment_2:
            raise ValueError(f"Experiment with run_id {run_id_2} not found")
        
        # Calculate comparison metrics
        comparison_metrics = self._calculate_comparison_metrics(experiment_1, experiment_2)
        
        return ExperimentComparison(
            experiment_1=experiment_1,
            experiment_2=experiment_2,
            comparison_metrics=comparison_metrics
        )
    
    def get_experiment_stats(self) -> ExperimentStats:
        """
        Get experiment statistics.
        
        Returns:
            Experiment statistics
        """
        return self.store.get_experiment_stats()
    
    def clone_experiment(self, run_id: str, new_name: Optional[str] = None) -> Experiment:
        """
        Clone an experiment with new run_id.
        
        Args:
            run_id: Original experiment ID
            new_name: Optional new name for cloned experiment
            
        Returns:
            Cloned experiment with new run_id
            
        Raises:
            ValueError: If original experiment not found
        """
        # Get original experiment
        original = self.get_experiment(run_id)
        if not original:
            raise ValueError(f"Experiment with run_id {run_id} not found")
        
        # Create cloned experiment
        clone_config = ExperimentCreate(
            experiment_name=new_name or f"{original.experiment_name} (Clone)" if original.experiment_name else None,
            description=f"Clone of experiment {run_id}",
            config_snapshot=original.config_snapshot,
            priority=original.priority,
            tags=original.tags.copy(),
            parent_experiment_id=original.run_id
        )
        
        return self.create_experiment(clone_config)
    
    def _validate_status_transition(self, current_status: ExperimentStatus, new_status: ExperimentStatus):
        """
        Validate that status transition is allowed.
        
        Args:
            current_status: Current experiment status
            new_status: New status to transition to
            
        Raises:
            ValueError: If transition is not allowed
        """
        # Define allowed transitions
        allowed_transitions = {
            ExperimentStatus.PENDING: [ExperimentStatus.RUNNING, ExperimentStatus.CANCELLED],
            ExperimentStatus.RUNNING: [ExperimentStatus.COMPLETED, ExperimentStatus.FAILED, ExperimentStatus.CANCELLED],
            ExperimentStatus.COMPLETED: [],  # Terminal state
            ExperimentStatus.FAILED: [],      # Terminal state
            ExperimentStatus.CANCELLED: []     # Terminal state
        }
        
        if new_status not in allowed_transitions.get(current_status, []):
            raise ValueError(f"Invalid status transition from {current_status} to {new_status}")
    
    def _calculate_comparison_metrics(self, exp1: Experiment, exp2: Experiment) -> Dict[str, Dict[str, Any]]:
        """
        Calculate comparison metrics between two experiments.
        
        Args:
            exp1: First experiment
            exp2: Second experiment
            
        Returns:
            Dictionary of comparison metrics
        """
        metrics = {}
        
        # Compare result summaries if both completed
        if (exp1.result_summary and exp2.result_summary and 
            exp1.status == ExperimentStatus.COMPLETED and 
            exp2.status == ExperimentStatus.COMPLETED):
            
            metrics['results'] = {
                'robustness_score_diff': exp2.result_summary.robustness_score - exp1.result_summary.robustness_score,
                'risk_score_diff': exp2.result_summary.risk_score - exp1.result_summary.risk_score,
                'success_rate_diff': exp2.result_summary.success_rate - exp1.result_summary.success_rate,
                'execution_time_diff': (exp2.result_summary.execution_time_ms or 0) - (exp1.result_summary.execution_time_ms or 0)
            }
        
        # Compare configurations
        metrics['config'] = {
            'same_model': exp1.model_name == exp2.model_name,
            'same_dataset': exp1.dataset_name == exp2.dataset_name and exp1.dataset_version == exp2.dataset_version,
            'same_attack_types': set(exp1.attack_types) == set(exp2.attack_types),
            'prompt_count_diff': exp2.prompt_count - exp1.prompt_count
        }
        
        # Compare timing
        metrics['timing'] = {
            'created_time_diff': (exp2.created_at - exp1.created_at).total_seconds(),
            'execution_time_diff': ((exp2.completed_at or exp2.created_at) - (exp1.completed_at or exp1.created_at)).total_seconds()
        }
        
        return metrics


# Global instance for easy access
_experiment_manager = None


def get_experiment_manager() -> ExperimentManager:
    """
    Get the global experiment manager instance.
    
    Returns:
        Global experiment manager instance
    """
    global _experiment_manager
    if _experiment_manager is None:
        _experiment_manager = ExperimentManager()
    return _experiment_manager


# Convenience functions for direct usage
def create_experiment(config: ExperimentCreate) -> Experiment:
    """
    Create a new experiment.
    
    Args:
        config: Experiment creation configuration
        
    Returns:
        Created experiment
    """
    manager = get_experiment_manager()
    return manager.create_experiment(config)


def update_experiment_status(run_id: str, status: ExperimentStatus) -> Experiment:
    """
    Update experiment status.
    
    Args:
        run_id: Unique experiment identifier
        status: New status
        
    Returns:
        Updated experiment
    """
    manager = get_experiment_manager()
    return manager.update_status(run_id, status)


def save_experiment_results(run_id: str, result_summary: ResultSummary, 
                           full_result: Optional[Dict[str, Any]] = None) -> Experiment:
    """
    Save experiment results.
    
    Args:
        run_id: Unique experiment identifier
        result_summary: Summary of results
        full_result: Complete result data
        
    Returns:
        Updated experiment
    """
    manager = get_experiment_manager()
    return manager.save_results(run_id, result_summary, full_result)


def get_experiment_by_id(run_id: str) -> Optional[Experiment]:
    """
    Get experiment by run_id.
    
    Args:
        run_id: Unique experiment identifier
        
    Returns:
        Experiment if found, None otherwise
    """
    manager = get_experiment_manager()
    return manager.get_experiment(run_id)