import random
import math
execfile('cw2-sim.py')  # load the simulator classes
# cw2-sim.py is available at http://www.cs.ucl.ac.uk/staff/d.wischik/Teach/NP/problems.html

class TCP:
  def __init__(self,logfunc=None): self.logfunc = logfunc  # a logging function, called every ACK/drop
  def initwnd(self,currenttime):   # called by the simulator, to get the initial window size
    self.w = 1.0       # target window size i.e. number of packets in flight
    self.inflight = 1  # current number of packets in flight
    return 1           # the number of packets to send immediately when the flow starts
  def receive(self,isAck,pktdata,rtt,currenttime):  # called by the simulator, when an ack arrives or drop detected
    # isAck=True for an ACK, False for a drop
    # pktdata can be ignored (see detailed docs in cw2-sim.py for details)
    # rtt is the round trip time experienced by this packet just arrived
    # currenttime is the current simulation time, in seconds
    self.inflight -= 1
    self.w = (self.w+1.0/self.w) if isAck else max(self.w/2.0,1)
    numtosend = max(int(math.floor(self.w))-self.inflight,0)
    self.inflight += numtosend
    if self.logfunc: self.logfunc(currenttime,'ack' if isAck else 'drop',self.w,self.inflight,numtosend)
    return numtosend   # the number of new packets to send out

# A convenience class for logging. with_prefix(p,q,...) returns a logging function that
# can be passed to an instance of TCP; this logging function prefixes each log entry with p,q,...
# To write out the log file as a CSV, call write_to(filename,headerline).
class Logger:
  def __init__(self): self.log = []
  def with_prefix(self,*label):
    def f(*x): self.log.append(label+x)
    return f
  def write_to(self,filename,header):
    with open(filename,'wt') as f:
      f.write(header+'\n')
      for logline in self.log:
        f.write(','.join([str(t) for t in logline])+'\n')


# Simulate a scenario with 32 long-lived TCP flows, for 1000 sec.
# Write traces for some of those flows to cw2-longlived-trace.csv.
# Write flow summary statistics to cw2-longlived-flowstats.csv
evlist = EventList()
logger = Logger()
q = Queue(500,60,evlist)
sources = []
links = []
for i in range(8):
  ix = i/7.0
  propdelay = (1-ix)*0.01 + ix*0.25
  link = Link(propdelay,evlist)
  for j in range(4):
    logfunc = logger.with_prefix(len(sources),propdelay) if j==0 else None
    # Create a traffic source using the TCP congestion controller.
    # Packets should first visit the queue, then traverse a link with the appropriate delay.
    # Flow sizes are infinite (i.e. they never stop).
    # Start them after a random delay (0--1 sec) to reduce synchronization problems.
    sources.append(Source(TCP(logfunc),
                          route=[q.receive,link.receive],
                          filesize=float('inf'),startafter=random.random(),evlist=evlist))
while evlist.time<1000 and evlist.events: evlist.execute()
#
logger.write_to('cw2-longlived-trace.csv',
                header='flowid,propdelay,time,event,w,inflight,tosend')
with open('cw2-longlived-flowstats.csv','wt') as f:
  wantcol = ['rtt','propdelay','lossrate','goodput','sendrate']
  f.write(','.join(wantcol)+'\n')
  for src in sources:
    stats = src.stats()
    f.write(','.join([str(stats[x]) for x in wantcol])+'\n')

    

# Simulate a scenario Poisson flow arrivals, for 1000 sec.
# Write traces for some of those flows to cw2-openloop-trace.csv.
# Write flow summary statistics to cw2-openloop-flowstats.csv
evlist = EventList()
logger = Logger()
q = Queue(500,60,evlist)
sources = []
propdelays = [(1-i/7.0)*0.01 + (i/7.0)*0.25 for i in range(8)]
links = [Link(propdelay,evlist) for propdelay in propdelays]
arrivalrate = 0.5
filesize = 950
class FlowStarter:
  def execute(self):
    link = random.choice(links)
    logfunc = logger.with_prefix(len(sources),link.delay) if len(sources) % 40 == 0 else None
    sources.append(Source(TCP(logfunc),
                          [q.receive,link.receive],
                          filesize=filesize,startafter=0,evlist=evlist))
    evlist.add(evlist.time+rexp(arrivalrate), self.execute)
evlist.add(0,FlowStarter().execute)
while evlist.time<1000 and evlist.events: evlist.execute()
#
logger.write_to('cw2-openloop-trace.csv',
                header='flowid,propdelay,time,event,w,inflight,tosend')
with open('cw2-openloop-flowstats.csv','wt') as f:
  wantcol = ['rtt','propdelay','lossrate','goodput','sendrate','start','duration','finished']
  f.write(','.join(wantcol)+'\n')
  for src in sources:
    stats = src.stats()
    f.write(','.join([str(stats[x]) for x in wantcol])+'\n')