# This is an event-driven simulator. A list of future events is maintained by an EventList object.
# - add(t,callback), where t is a number and callback is a function taking no arguments
#   This schedules the function callback to be invoked at simulation time t
# - execute()
#   This advances the simulation time to the next scheduled event, and invokes it.


# The simulator consists of Queues, Links and traffic Sources. These objects shift around
# Packet objects. Each of these objects can receive a packet; if a Queue or Link receives
# a packet then it will schedule a time in the future for the packet to be sent on.
#
# Routing is static. When a Source object is created, it is given a route, and all the
# packets it creates will follow that route. An example route is
#      [myqueue.receive,mylink.receive,thissource.receive]
# which means: after the packet is generated then call myqueue.receive(pkt), then when the
# queue is ready to send on the packet call mylink.receive, then call thissource.receive.
# There are no explicit objects for flow destinations or for ACKs. The packet is simply
# considered to circulate all the way round until it returns to its Source.
#
# This simulator has a simplified notion of drops. In order to make drop detection
# easier, packets are never truly dropped -- they are simply marked as 'ghost packets'
# by calling markdropped(pkt). Ghost packets continue to circulate through the network,
# in sequence with other packets, though they take up no service time and no buffer space
# at the queue. When the source receives a ghost packet, the ghost packet is treated as
# a signal that that packet was dropped.
#
# This simulator does not bother about actual data. It is purely a model for congestion
# control, and it doesn't bother if a given packet contains retransmitted old
# data, or if it contains fresh data. In either case, it's just data which takes up
# space in the network.


# Methods to be used by the simulation harness:
#
# Link(delay,evlist)
# - create a new link with specified propagation delay (numeric), which can schedule itself
#   for future wakeups using evlist.add
# Link.receive(pkt)
# - make this packet enter the pipe; at time 'delay' in the future the link will wake up and
#   send the packet on to its next hop.
#
# Queue(servicerate,bufsize,evlist)
# - create a queue with specified service rate (numeric), buffer size (integer), which can
#   schedule itself for future wakeups using evlist.add. The queue actually uses random exponentially
#   distributed per-packet service times with mean 1/servicerate. This is in order to mimic
#   random jitter/noise, and to break up synchronization artefacts.
# Queue.receive(pkt)
# - make this packet enter the queue (or, if the queue is full, turn the packet into a ghost packet).
#   At some time in the future, the queue will wake up and send the packet on to its next hop.
#
# Source(congcont,route,filesize,startafter,evlist)
# - create a new traffic source using the specified congestion controller, route.
#   It will wait until currenttime+startafter before transmissint its first packet.
#   It will send packets until it has received filesize ACKs (i.e. non-ghost packets).
#   - route should be a list of callbacks which each take a single argument, namely a packet.
#     Each packet from this source will visit each object in the list in turn.
#     The source is automatically added to the end of the list; do not include it in the route.
#   - congcont is a congestion controller. It should have two methods:
#     - initwnd(currenttime), called before sending the first packet
#     - receive(isAck, metadata, rtt, currenttime), called whenever an ACK is received (or a drop detected).
#     Both of these functions should return either an integer or a list.
#     If they return an integer N, the source will immediately send out N packets.
#     If they return a list [a,b,...] of length N, the source will immediately send out N packets,
#     and the packets will contain metadata a, b, ... This metadata is passed to future calls of receive.
#     The source keeps track of the time that each packet was sent, so it knows the rtt experienced by each
#     packet, and this is the rtt argument to receive().
# Source.stats()
# - returns a dictionary with keys 'sendrate', 'goodput'. These give summary statistics over the
#   lifetime of the traffic flow.


import math
import collections
import heapq
import random
# Generate a random number with Exp(lambd) distribution, mean 1/lambd.
def rexp(lambd): return -1.0/lambd * math.log(random.random())

class EventList:
  def __init__(self):
    self.time,self.events = 0,[]
  def add(self,t,callback):
    heapq.heappush(self.events, (t,callback))
  def execute(self):
    if self.events:
      self.time,callback = heapq.heappop(self.events)
      callback()

class Packet:
  def __init__(self,route,timesent,metadata):
    self.route,self.timesent,self.metadata,self.wasdropped = route,timesent,metadata,False
  def markdropped(self): self.wasdropped=True
  def sendon(self): self.route.pop()(self)

class Link:
  def __init__(self,delay,evlist):
    self.delay = delay
    self.inflight,self.evlist = collections.deque(),evlist
  def receive(self,pkt):
    self.inflight.appendleft( (self.evlist.time+self.delay,pkt) )
    if len(self.inflight)==1: self.evlist.add(self.inflight[-1][0],self.emitnextpkt)
  def emitnextpkt(self):
    time,pkt = self.inflight.pop()
    pkt.sendon()
    if len(self.inflight)>0: self.evlist.add(self.inflight[-1][0],self.emitnextpkt)
   
class Queue:
  def __init__(self,servicerate,bufsize,evlist):
    self.servicerate,self.bufsize,self.queuesize = servicerate,bufsize,0
    self.packetsqueued,self.packetsdropped,self.packetsserved = 0,0,0
    self.queue,self.evlist = collections.deque(),evlist
  def receive(self,pkt):
    if pkt.wasdropped:
      pktservicetime = 0
    elif self.queuesize>=self.bufsize:
      pkt.markdropped()
      pktservicetime = 0
      self.packetsdropped += 1
    else:
      self.queuesize += 1
      pktservicetime = rexp(self.servicerate)
      self.packetsqueued += 1
    self.queue.appendleft( (pktservicetime,pkt) )
    if len(self.queue)==1: self.evlist.add(self.evlist.time+self.queue[-1][0],self.finishservice)
  def finishservice(self):
    servicetime,pkt = self.queue.pop()
    if servicetime>0:
      self.queuesize -= 1
      self.packetsserved += 1
    pkt.sendon()
    if len(self.queue)>0: self.evlist.add(self.evlist.time+self.queue[-1][0],self.finishservice)
  def stats(self):
    return {'utilization':self.packetsserved/(self.servicerate*self.evlist.time),
            'lossrate':self.packetsdropped/float(self.packetsqueued+self.packetsdropped)}

class Source:
  def __init__(self,congcont,route,filesize,startafter,evlist):
    self.congcont,self.route,self.filesize = congcont,route,filesize
    self.packetssent,self.acks,self.drops,self.totrtt = 0,0,0,0
    self.evlist = evlist
    self.evlist.add(self.evlist.time+startafter,self.sendfirstpkts)
  def sendfirstpkts(self):
    self.firstsendtime,self.lastreceivetime = self.evlist.time,float('-inf')
    self.sendpackets(self.congcont.initwnd(self.evlist.time))
  def receive(self,pkt):
    self.lastreceivetime = self.evlist.time
    self.drops,self.acks = (self.drops+1,self.acks) if pkt.wasdropped else (self.drops,self.acks+1)
    rtt = self.evlist.time - pkt.timesent
    self.totrtt += rtt
    tosend = self.congcont.receive(not pkt.wasdropped,pkt.metadata,rtt,self.evlist.time)
    self.sendpackets(tosend)
  def sendpackets(self,tosend):
    if isinstance(tosend,int): tosend = [() for i in range(tosend)]
    numlefttosend = max(0,self.filesize-self.acks)
    if len(tosend)>numlefttosend: tosend = tosend[:numlefttosend]
    for m in tosend:
      newroute = self.route+[self.receive]
      newroute.reverse()
      Packet(newroute,self.evlist.time,m).sendon()
      self.packetssent += 1
  def stats(self):
    duration = self.lastreceivetime - self.firstsendtime
    return {'sendrate':self.packetssent/float(duration),
            'goodput':self.acks/float(duration),
            'lossrate':1-self.acks/float(self.acks+self.drops),
            'propdelay':sum(r.__self__.delay for r in self.route if isinstance(r.__self__,Link)),
            'rtt':self.totrtt/float(self.drops+self.acks),
            'start':self.firstsendtime,
            'duration':duration,
            'finished':self.acks>=self.filesize}
