import time
from functools import wraps
from collections import defaultdict
import numpy as np

class PerformanceMonitor:
    def __init__(self):
        self.response_times = defaultdict(list)
        self.success_counts = defaultdict(int)
        self.total_counts = defaultdict(int)
        self.problem_types = defaultdict(int)
        self.total_problems = 0
    
    def record_response_time(self, model_type: str, time: float):
        """Record response time for a model"""
        self.response_times[model_type].append(time)
    
    def record_success(self, model_type: str, success: bool):
        """Record success/failure for a model"""
        self.total_counts[model_type] += 1
        if success:
            self.success_counts[model_type] += 1
    
    def record_problem_type(self, problem_type: str):
        """Record problem type"""
        self.problem_types[problem_type] += 1
        self.total_problems += 1
    
    def get_statistics(self) -> dict:
        """Get current performance statistics"""
        stats = {}
        
        # Calculate average response times
        for model_type, times in self.response_times.items():
            if times:
                stats[f'{model_type}_avg_response_time'] = np.mean(times)
        
        # Calculate success rates
        for model_type in self.total_counts.keys():
            total = self.total_counts[model_type]
            if total > 0:
                success_rate = (self.success_counts[model_type] / total) * 100
                stats[f'{model_type}_success_rate'] = success_rate
        
        # Calculate problem type distribution
        if self.total_problems > 0:
            distribution = {
                ptype: (count / self.total_problems) * 100
                for ptype, count in self.problem_types.items()
            }
            stats['problem_type_distribution'] = distribution
        
        return stats

def measure_time(func):
    """Decorator to measure function execution time"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        return result, end_time - start_time
    return wrapper