import heapq
import random
import math
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

class Task:
    def __init__(self, id, arrival_time, duration, sigma):
        self.id = id
        self.arrival_time = arrival_time
        self.base_duration = duration # Time on 1 GPU
        self.sigma = sigma
        self.start_time = -1
        self.finish_time = -1
        self.assigned_k = 0

    def get_exec_time(self, k):
        # Efficiency = sigma ^ log2(k)
        # Time = base / (k * efficiency)
        if k == 1: return self.base_duration
        eff = self.sigma ** math.log2(k)
        return self.base_duration / (k * eff)

class Event:
    def __init__(self, time, type, data):
        self.time = time
        self.type = type # 'ARRIVAL' or 'FINISH'
        self.data = data
    
    def __lt__(self, other):
        return self.time < other.time

class Cluster:
    def __init__(self, total_gpus=64):
        self.total_gpus = total_gpus
        self.free_gpus = total_gpus
        self.running_tasks = [] # List of (task, finish_time)

class Simulator:
    def __init__(self, strategy_name, arrival_rate, max_tasks=1000):
        self.strategy_name = strategy_name
        self.arrival_rate = arrival_rate
        self.max_tasks = max_tasks
        self.events = []
        self.cluster = Cluster()
        self.queue = []
        self.finished_tasks = []
        self.current_time = 0.0
        self.total_gpu_busy_time = 0.0 # Integral of busy GPUs over time
        self.last_update_time = 0.0

    def schedule_events(self):
        # Generate all arrivals upfront or dynamically
        t = 0
        for i in range(self.max_tasks):
            dt = random.expovariate(self.arrival_rate)
            t += dt
            duration = random.uniform(10, 100)
            sigma = random.uniform(0.75, 0.95)
            task = Task(i, t, duration, sigma)
            heapq.heappush(self.events, Event(t, 'ARRIVAL', task))

    def update_metrics(self):
        dt = self.current_time - self.last_update_time
        busy_gpus = self.cluster.total_gpus - self.cluster.free_gpus
        self.total_gpu_busy_time += busy_gpus * dt
        self.last_update_time = self.current_time

    def get_k(self, task):
        # Strategies
        available = self.cluster.free_gpus
        q_len = len(self.queue)
        
        possible_ks = [1, 2, 4, 8, 16, 32, 64]
        valid_ks = [k for k in possible_ks if k <= available]
        
        if not valid_ks: return 0
        
        if self.strategy_name == 'conservative':
            # Always k=1
            return 1 if 1 in valid_ks else 0

        elif self.strategy_name == 'aggressive':
            # Use max possible parallel, up to limit (e.g. 16 to avoid total waste)
            # Or just max available
            return valid_ks[-1]
            
        elif self.strategy_name == 'adaptive':
            # Based on queue length
            if q_len == 0:
                target = 32 # High speed if no contention
            elif q_len < 5:
                target = 16
            elif q_len < 10:
                target = 8
            elif q_len < 20:
                target = 4
            elif q_len < 50:
                target = 2
            else:
                target = 1
            
            # Find largest valid k <= target
            best_k = 1
            for k in valid_ks:
                if k <= target:
                    best_k = k
            return best_k
        
        return 1

    def run(self):
        self.schedule_events()
        
        while self.events or self.cluster.running_tasks:
            if not self.events and not self.cluster.running_tasks:
                break
                
            # Peek next event
            if not self.events:
                next_time = float('inf')
            else:
                next_time = self.events[0].time
            
            # Jump time
            self.current_time = next_time
            self.update_metrics()
            
            event = heapq.heappop(self.events)
            
            if event.type == 'ARRIVAL':
                task = event.data
                self.queue.append(task)
            elif event.type == 'FINISH':
                task = event.data
                self.cluster.free_gpus += task.assigned_k
                self.finished_tasks.append(task)
            
            # Try to schedule waiting tasks
            # We iterate queue. Note: Standard queue is FIFO.
            # We can't easily remove from middle if we skip, so we only look at head
            # OR we can try to fit small tasks? 
            # Simple FIFO: look at head. If can schedule, do it. Else stop (or continue?)
            # Let's do Strict FIFO for simplicity and fairness
            
            while self.queue:
                head_task = self.queue[0]
                k = self.get_k(head_task)
                if k > 0:
                    # Assign
                    self.queue.pop(0)
                    head_task.assigned_k = k
                    head_task.start_time = self.current_time
                    exec_time = head_task.get_exec_time(k)
                    head_task.finish_time = self.current_time + exec_time
                    
                    self.cluster.free_gpus -= k
                    heapq.heappush(self.events, Event(head_task.finish_time, 'FINISH', head_task))
                else:
                    # Cannot schedule head task
                    break

    def calculate_results(self):
        total_time = self.current_time
        avg_utilization = self.total_gpu_busy_time / (total_time * 64)
        
        delays = []
        for t in self.finished_tasks:
            expected_finish = t.arrival_time + t.base_duration
            if t.finish_time <= expected_finish:
                delta = 0
            else:
                overdue = t.finish_time - expected_finish
                # Normalized penalty
                delta = (overdue / t.base_duration) ** 2
            delays.append(delta)
        
        avg_delay = sum(delays) / len(delays) if delays else 0
        return avg_utilization, avg_delay

def run_simulations():
    strategies = ['conservative', 'aggressive', 'adaptive']
    # Load: Light (0.5), Medium (0.9), Heavy (1.1)
    loads = [0.5, 0.9, 1.1] 
    
    results = []
    
    print("Running GPU Simulations...")
    for load in loads:
        for strat in strategies:
            # Run multiple times to average? Just once for this demo with 1000 tasks
            sim = Simulator(strat, load, max_tasks=1000)
            sim.run()
            eta, delta = sim.calculate_results()
            results.append({
                'Load': load,
                'Strategy': strat,
                'Utilization (eta)': eta,
                'Avg Delay Penalty (delta)': delta,
                'Score (Balanced)': (1-eta) + 0.1 * delta # Example lambda=0.1
            })
            print(f"Load {load}, Strat {strat}: Eta={eta:.3f}, Delta={delta:.3f}")

    df = pd.DataFrame(results)
    df.to_csv("results/gpu_sim_results.csv", index=False)
    
    # Plotting
    plt.figure(figsize=(12, 5))
    
    plt.subplot(1, 2, 1)
    sns.barplot(data=df, x='Load', y='Utilization (eta)', hue='Strategy')
    plt.title('Cluster Utilization')
    
    plt.subplot(1, 2, 2)
    sns.barplot(data=df, x='Load', y='Avg Delay Penalty (delta)', hue='Strategy')
    plt.title('User Average Delay Penalty')
    
    plt.tight_layout()
    plt.savefig('results/gpu_sim_plots.png')
    plt.close()

if __name__ == "__main__":
    run_simulations()
