import argparse
import random
import queue
import dataclasses
from dataclasses import dataclass
import matplotlib.pyplot as plt  # type: ignore
import numpy as np
from matplotlib import colors
from typing import cast, List, Union, Tuple

@dataclass
class SimulatorArgs:
    interarrival_time: float
    service_time: float
    simulation_time: float
    num_servers: int
    seed: int
    histogram: bool

args: SimulatorArgs

# The service being modeled. Jobs arrive and wait in the job_queue until it is
# their turn, then they take service_time seconds to be served.
@dataclass
class TheService:
    # The ordered list of jobs waiting to be executed. The float represents the
    # job's original arrival time.
    job_queue: queue.Queue[float] = dataclasses.field(default_factory=queue.Queue)

    # How many jobs are we currently executing? always <= args.num_servers
    servers_in_use = 0

    # track the total latency of each job, so we can compute stats later
    latencies: List[float] = dataclasses.field(default_factory=list)


# The system proceeds through a sequence of events. Each event in this simple
# service is either a job arrival or a job completion/departure.
Event = Union['Arrival', 'Departure']

# Each Event class implements a process() method that simulates that event.

@dataclass(order=True)
class Arrival:
    def process(self, now: float, event_queue: queue.PriorityQueue[Tuple[float, Event]], service: TheService) -> None:
        if service.servers_in_use < args.num_servers:
            service.servers_in_use += 1
            event_queue.put((now + args.service_time, Departure(now)))
        else:
            service.job_queue.put(now)


@dataclass(order=True)
class Departure:
    arrival_time: float = dataclasses.field(compare=False)

    def process(self, now: float, event_queue: queue.PriorityQueue[Tuple[float, Event]], service: TheService) -> None:
        elapsed = now - self.arrival_time
        service.latencies.append(elapsed)

        service.servers_in_use -= 1
        if not service.job_queue.empty():
            arrival_time = service.job_queue.get()
            service.servers_in_use += 1
            event_queue.put((now + args.service_time, Departure(arrival_time)))


def parse_args() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument('--interarrival-time', type=float, default=3.0, help='average time between requests')
    parser.add_argument('--service-time', type=float, default=1.0, help='deterministic time taken by each request')
    parser.add_argument('--simulation-time', type=float, default=60 * 60 * 24.0, help='how much simulated time, default = 1 day')
    parser.add_argument('--num-servers', type=int, default=1, help='how many servers serve the queue in parallel')
    parser.add_argument('--seed', type=int, default=0, help='set the random seed')
    parser.add_argument('--histogram', action='store_true', help='show a histogram of the latencies')
    global args
    args = cast(SimulatorArgs, parser.parse_args())


# Here we implement a little discrete event simulator.
def main() -> None:
    parse_args()
    random.seed(args.seed)
    
    # This queue tracks events that have not been simulated yet. Each element of
    # the queue is a pair of a float (at what simulated time should the event
    # occur) and an Event object that says what to do to simulate that event.
    # The simulator executes events in order of time. Events can schedule other
    # events in the future.
    event_queue: queue.PriorityQueue[Tuple[float, Event]] = queue.PriorityQueue()

    arrival_rate = 1.0 / args.interarrival_time
    service_rate = args.num_servers / args.service_time
    utilization = arrival_rate / service_rate
    print('utilization:', utilization)

    # Generate random arrival events. Each job arrival is uniformly distributed
    # across the simulation window. (This is one way of simulating a Poisson
    # process.)
    n_requests = int(args.simulation_time / args.interarrival_time)
    for i in range(n_requests):
        t = random.uniform(0, args.simulation_time)
        event_queue.put((t, Arrival()))

    # Here is the main simulation loop.
    service = TheService()
    while not event_queue.empty():
        now, event = event_queue.get()
        event.process(now, event_queue, service)

    # Now compute and print statistics.
    latencies = service.latencies
    assert len(latencies) == n_requests
    if args.num_servers == 1:
        lam = 1.0 / args.interarrival_time
        mu = 1.0 / args.service_time
        rho = lam / mu
        print('theoretical avg latency:', args.service_time + rho / (2 * mu * (1 - rho)))
    print('avg latency:', np.mean(latencies))
    print('median latency:', np.median(latencies))
    print('p99 latency:', np.percentile(latencies, q=99))

    # And draw a histogram of the latencies.
    if args.histogram:
        fig, ax = plt.subplots()
        ax.hist(latencies, bins=25)
        plt.show()


if __name__ == '__main__':
    main()