import time
import numpy as np
from typing import List, Dict, Any
from .model_runner import ModelRunner
from .input_loader import InputLoader
from .resource_monitor import ResourceMonitor

class Benchmark:
    def __init__(self, models_config: List[Dict[str, Any]], input_loader: InputLoader):
        self.models_config = models_config
        self.input_loader = input_loader
        self.runners = [ModelRunner(cfg) for cfg in models_config]
        self.monitor = ResourceMonitor(interval=1.0)
        self.resource_results = None

    def warmup(self, iterations=10):
        # Silence warmup logs
        for runner, config in zip(self.runners, self.models_config):
            input_data = self.input_loader.get_input(config)
            for _ in range(iterations):
                runner.run_inference(input_data)

    def run_single_model_benchmark(self, iterations=100):
        results = []
        
        self.monitor.start()
        
        for runner, config in zip(self.runners, self.models_config):
            latencies = []
            input_data = self.input_loader.get_input(config)
            
            for _ in range(iterations):
                start = time.time()
                runner.run_inference(input_data)
                latencies.append((time.time() - start) * 1000) # ms
            
            avg_lat = np.mean(latencies)
            min_lat = np.min(latencies)
            max_lat = np.max(latencies)
            throughput = config['batch_size'] / (avg_lat / 1000)
            
            results.append({
                'model_name': runner.name,
                'batch_size': config['batch_size'],
                'avg_latency_ms': avg_lat,
                'min_latency_ms': min_lat,
                'max_latency_ms': max_lat,
                'throughput': throughput
            })
            
        self.resource_results = self.monitor.stop()
        return results

    def run_pipeline_benchmark(self, iterations=100):
        latencies = []
        
        self.monitor.start()
        
        start_total = time.time()
        for _ in range(iterations):
            start_iter = time.time()
            # Sequential execution for Single NPU
            for runner, config in zip(self.runners, self.models_config):
                input_data = self.input_loader.get_input(config)
                runner.run_inference(input_data)
            latencies.append((time.time() - start_iter) * 1000)
        
        total_time = time.time() - start_total
        if latencies:
            avg_pipe_lat = np.mean(latencies)
            # Use batch_size of the first model as representative for pipeline
            rep_batch_size = self.models_config[0].get('batch_size', 1)
            pipeline_fps = rep_batch_size / (avg_pipe_lat / 1000.0)
        else:
            avg_pipe_lat = 0
            pipeline_fps = 0
            
        self.resource_results = self.monitor.stop()
        
        return {
            'total_iterations': iterations,
            'avg_pipeline_latency_ms': avg_pipe_lat,
            'pipeline_fps': pipeline_fps,
            'total_execution_time_s': total_time
        }

    def close(self):
        for runner in self.runners:
            runner.close()
        self.input_loader.release()
