dudong-v2/utils/session_manager.py

"""
Session Manager

Tracks and manages analysis session data.
"""

from datetime import datetime
from typing import List, Dict, Optional
from dataclasses import dataclass, field


@dataclass
class TestResult:
    """Data class for a single test result."""
    test_id: int
    timestamp: datetime
    classification: str
    confidence: float  # 0-100
    probabilities: Dict[str, float]  # class_name -> probability
    processing_time: float
    file_path: Optional[str] = None
    
    
class SessionManager:
    """
    Manager for tracking ripeness testing session data.
    
    Tracks test results, calculates statistics, and manages session state.
    """
    
    def __init__(self, max_history: int = 100):
        """
        Initialize session manager.
        
        Args:
            max_history: Maximum number of results to keep in history
        """
        self.max_history = max_history
        self.session_start = datetime.now()
        self.results: List[TestResult] = []
        self.test_counter = 0
        
    def add_result(self, classification: str, confidence: float,
                   probabilities: Dict[str, float], processing_time: float,
                   file_path: Optional[str] = None) -> TestResult:
        """
        Add a new test result to the session.
        
        Args:
            classification: Predicted class name
            confidence: Confidence percentage (0-100)
            probabilities: Dictionary of class probabilities
            processing_time: Processing time in seconds
            file_path: Optional path to the processed file
            
        Returns:
            TestResult object
        """
        self.test_counter += 1
        
        result = TestResult(
            test_id=self.test_counter,
            timestamp=datetime.now(),
            classification=classification,
            confidence=confidence,
            probabilities=probabilities,
            processing_time=processing_time,
            file_path=file_path
        )
        
        self.results.insert(0, result)  # Insert at beginning (most recent first)
        
        # Trim history if needed
        if len(self.results) > self.max_history:
            self.results = self.results[:self.max_history]
            
        return result
        
    def get_total_tests(self) -> int:
        """Get total number of tests in session."""
        return len(self.results)
        
    def get_average_processing_time(self) -> float:
        """Get average processing time across all tests."""
        if not self.results:
            return 0.0
        return sum(r.processing_time for r in self.results) / len(self.results)
        
    def get_classification_counts(self) -> Dict[str, int]:
        """
        Get counts for each classification type.
        
        Returns:
            Dictionary mapping classification names to counts
        """
        counts = {}
        for result in self.results:
            classification = result.classification
            counts[classification] = counts.get(classification, 0) + 1
        return counts
        
    def get_ripe_count(self) -> int:
        """Get number of 'Ripe' classifications."""
        return sum(1 for r in self.results if r.classification == "Ripe")
        
    def get_session_duration(self) -> Dict[str, int]:
        """
        Get session duration.
        
        Returns:
            Dictionary with 'hours' and 'minutes' keys
        """
        duration = datetime.now() - self.session_start
        hours = duration.seconds // 3600
        minutes = (duration.seconds % 3600) // 60
        return {"hours": hours, "minutes": minutes}
        
    def get_recent_results(self, count: int = 10) -> List[TestResult]:
        """
        Get most recent test results.
        
        Args:
            count: Number of results to return
            
        Returns:
            List of TestResult objects
        """
        return self.results[:count]
        
    def get_last_result(self) -> Optional[TestResult]:
        """Get the most recent test result."""
        return self.results[0] if self.results else None
        
    def clear_session(self):
        """Clear all session data and reset."""
        self.results.clear()
        self.test_counter = 0
        self.session_start = datetime.now()
        
    def get_statistics_summary(self) -> Dict[str, any]:
        """
        Get complete statistics summary.
        
        Returns:
            Dictionary containing all session statistics
        """
        counts = self.get_classification_counts()
        duration = self.get_session_duration()
        
        return {
            "total_tests": self.get_total_tests(),
            "avg_processing_time": self.get_average_processing_time(),
            "ripe_count": self.get_ripe_count(),
            "classification_counts": counts,
            "session_duration": duration,
            "session_start": self.session_start
        }