#--------------------------------------------------------------------------------
# Load in some standard functions, and define some basic probability functions

import time
import heapq
import math
from random import random
Inf = float('inf')
def rexp(lambd):
    """Generate an exponential random variable with rate lambd"""
    if lambd==0: return Inf
    u = random()
    return (-1.0/lambd) * math.log(u)
def mean(d):
    """Given a dictionary of value,prob terms representing a distribution, find the mean"""
    return sum(k*v for k,v in d.iteritems())

#----------------------------------------------------------------------------------
## A link with two classes of user
#
# My code here is generic, not specifically about takers and sharers --
# you pass the code a function called throughputFunction, which calculates the throughput that
# users of each class will receive.

class ProcShareLink:
    """A link shared between two classes of users, called A and B, which get throughputs (thA,thB)=func(#A,#B)""" 
    def __init__(self,throughputFunction):
        self.throughputFunction = throughputFunction
        self.waitingA = []
        self.waitingB = []
        self.workdoneperA = 0.0
        self.workdoneperB = 0.0
    def arrivalA(self,flowsize):
        heapq.heappush(self.waitingA,self.workdoneperA+flowsize)
    def arrivalB(self,flowsize):
        heapq.heappush(self.waitingB,self.workdoneperB+flowsize)
    def advancetime(self,by,isdeparture=False):
        nA,nB = len(self.waitingA),len(self.waitingB)
        throughputA,throughputB = self.throughputFunction(nA,nB)
        self.workdoneperA += throughputA*by
        self.workdoneperB += throughputB*by
        if isdeparture:
            depB = nB>0 and (nA==0 or (self.waitingB[0]-self.workdoneperB < self.waitingA[0]-self.workdoneperA))
            if depB: heapq.heappop(self.waitingB)
            else: heapq.heappop(self.waitingA)
    def timeuntilnextcompletion(self):
        if len(self.waitingA)+len(self.waitingB)==0: return Inf
        nA,nB = len(self.waitingA),len(self.waitingB)
        throughputA,throughputB = self.throughputFunction(nA,nB)
        timetodepartA = (self.waitingA[0]-self.workdoneperA)/throughputA if throughputA>0 else Inf
        timetodepartB = (self.waitingB[0]-self.workdoneperB)/throughputB if throughputB>0 else Inf
        return min(timetodepartA,timetodepartB)

def runsim(lambdaA,lambdaB,filesizeA,filesizeB,throughputFunction):
    link = ProcShareLink(throughputFunction)
    def randomInterarrival(): return rexp(lambdaA+lambdaB)
    def randomIsB(): return random() < (0.0+lambdaB)/(lambdaA+lambdaB) 
    def randomFilesizeA(): return rexp(1.0/filesizeA)
    def randomFilesizeB(): return rexp(1.0/filesizeB)
    #
    simtime = 0.0
    nextarrival = simtime + randomInterarrival()
    trace = [(0,0,0)]  # store the trace as a list of (simtime,#A,#B)
    for i in range(10000):
        nextdeparture = simtime + link.timeuntilnextcompletion()
        if nextarrival<nextdeparture:
            link.advancetime(nextarrival-simtime)
            if randomIsB(): link.arrivalB(randomFilesizeB())
            else: link.arrivalA(randomFilesizeA())
            simtime = nextarrival
            nextarrival = simtime + randomInterarrival()
        else:
            link.advancetime(nextdeparture-simtime, isdeparture=True)
            simtime = nextdeparture
        trace.append( (simtime,len(link.waitingA),len(link.waitingB)) )
    #
    # Find the fraction of time spent at each level of #A, and at each level of #B
    # Return the mean #A and the mean #B
    timeWithAat,timeWithBat = {},{}
    tottime = trace[-1][0]
    t,nA,nB = trace[0]
    for t2,nA2,nB2 in trace[1:]:
        timeWithAat[nA] = (t2-t)/tottime + (timeWithAat[nA] if nA in timeWithAat else 0)
        timeWithBat[nB] = (t2-t)/tottime + (timeWithBat[nB] if nB in timeWithBat else 0)
        t,nA,nB = t2,nA2,nB2
    return sum(k*v for k,v in timeWithAat.iteritems()), sum(k*v for k,v in timeWithBat.iteritems())



#----------------------------------------------------------------------------------
## Now we specialize to takers and sharers, with the throughput function specified in the coursework.

# The command ratecap(c,a) returns a function f, where f(nt,ns) is the throughput that takers and sharers
# get respectively when there are nt takers and ns sharers. For example, ratecap(100,10)(4,10)=(10,6)
# (In Python, the lambda command defines an anonymous function.)
def ratecap(c,a): return lambda nt,ns: (min(float(c)/nt,a) if nt>0 else 0, max(c-nt*a,0)/float(ns) if ns>0 else 0)
def procshare(c): return lambda nt,ns: (float(c)/(nt+ns) if nt>0 else 0, float(c)/(nt+ns) if ns>0 else 0)

# Validation 1: no takers, so sharers should be pure processor sharing
res1 = [("validate1",0,30,-1,1,100,-1, runsim(0,30,-1,1,ratecap(100,1))) for r in range(15)]
# Validation 2: very small takers, so sharers should be nearly pure processor sharing
res2 = [("validate2",30,30,.01,1,100,1, runsim(30,30,.01,1,ratecap(100,1))) for r in range(15)]
# Validation 3: large rate cap for the takers, so they should be pure processor sharing
res3 = [("validate3",30,30,1.8,1,100,100, runsim(30,30,1.8,1,ratecap(100,100))) for r in range(15)]
# Validation 4: pretty large rate cap for the takers, so they should be nearly pure processor sharing
res4 = [("validate4",30,30,1.8,1,100,95, runsim(30,30,1.8,1,ratecap(100,95))) for r in range(15)]
# Main experiment: 15 runs at a range of values of a, for lambdat=30,lambdas=30,ft=1.8,fs=1,c=100
res5 = [("ratecap",30,30,1.8,1,100,a, runsim(30,30,1.8,1,ratecap(100,a))) for a in range(1,100,8) for r in range(15)]
# Test of pure proc sharing
res6 = [("procshare",30,30,1.8,1,100,-1, runsim(30,30,1.8,1,procshare(100))) for r in range(15)]

allres = [res1,res2,res3,res4,res5,res6]
# Write it all out to a log, to load into R for plotting
f = open('cw2-out.csv','wt')
print >>f, 'expt,lambdat,lambdas,ft,fs,c,a,nt,ns'
for r in allres:
    for expt,lambdat,lambdas,ft,fs,c,a,(nt,ns) in r:
        print >>f, ','.join([str(x) for x in [expt,lambdat,lambdas,ft,fs,c,a,nt,ns]])
f.close()
    

#----------------------------------------------------------------------------------
## R code for plotting & analysis

# Load in a library with plotting routines, and set graphics style
library(djwutils)
gp <- col.whitebg()
gp$superpose.line$col <- c('black','black','red','red')
gp$superpose.line$lwd <- c(2,2,1,1)
gp$superpose.line$lty <- c(1,2)
gp$superpose.symbol$pch <- c(0,19)
gp$superpose.symbol$col <- c('black','black','red','red')
gp$superpose.symbol$cex <- 1.5
trellis.par.set(gp)

# load the data
df <- read.csv('cw2-out.csv')
df$wt <- df$nt/df$lambdat
df$ws <- df$ns/df$lambdas

# a utility function, to show mean and confidence interval
meanci <- function(x) c(n=length(x),
                        lower=mean(x)-1.96*sd(x)/sqrt(length(x)-1),
                        mean=mean(x),
                        upper=mean(x)+1.96*sd(x)/sqrt(length(x)-1))

# Summary of how many observations there are in each experiment, to remind ourself
xtabs(~expt, data=df)
unique(df[,c('expt','lambdat','lambdas','ft','fs','c','a')])

# Validation 1: we expect #sharers = rho/(1-rho), wait=fs/(C-lambdas*fs)
meanci(df$ws[df$expt=='validate1'])
1/(100-30*1)

# Validation 2: we expect a similar answer
meanci(df$ws[df$expt=='validate2'])

# Validation 3: we expect #takers = rho/(1-rho), wait=ft/(c-lambda*ft)
meanci(df$wt[df$expt=='validate3'])
1.8/(100-30*1.8)

# Validation 4: we expect a similar answer
meanci(df$wt[df$expt=='validate4'])

# Plot of main experimental outcome
bwplot2(wt+ws~a, data=df, subset=expt=='ratecap',
        err=.95,type='b', ylim=c(0,0.3), panel=function(...) { panel.abline(h=c(0,.1,.2),col='grey75'); panel.bwplot2(...) },
        auto.key=TRUE)

# Comparison of procshare sim with procshare theory. Should be ft/(c-lambda*f), fs/(c-lambda*f)
meanci(df$wt[df$expt=='procshare'])
meanci(df$ws[df$expt=='procshare'])
1.8/(100-60*1.4)
1/(100-60*1.4)

# Comparison of ratecap sim with ratecap theory.
erlang <- function(rho,C) dpois(C,rho)/ppois(C,rho)
frlang <- function(rho,C) rho*(1-erlang(rho,C))
q4flows <- function(rho,alpha) {
    e <- 1/erlang(rho*alpha,floor(alpha))
    g <- rho/(1-rho)
    e/(e+g) * frlang(rho*alpha,floor(alpha)) + g/(e+g) * (floor(alpha)+1/(1-rho))
    }
dfrc <- df[df$expt=='ratecap',]
dfrc$ntTheory <- q4flows(dfrc$lambdat*dfrc$ft/dfrc$c,dfrc$c/dfrc$a)
dfrc$nsTheory <- dfrc$lambdas*dfrc$fs/(dfrc$c-dfrc$ntTheory*dfrc$a-dfrc$lambdas*dfrc$fs)
dfrc$nsTheory[dfrc$nsTheory<0] <- NA
#
bwplot2(nt+ns+ntTheory+nsTheory~a, data=dfrc,
        err='none', type='l', ylim=c(0,10), 
        auto.key=list(lines=TRUE,points=FALSE))

# Comparison of status quo to ratecap scheme
bwplot2(wt+ws~a, data=df, subset=expt=='ratecap',
        err='none',type='b', ylim=c(0,.3), xlim=c(0,50),
        panel=function(...) {
            panel.abline(h=1/(100-60*1.4),col='red',lty=2,lwd=1) # sharers
            panel.abline(h=1.8/(100-60*1.4),col='red',lty=1,lwd=1) # takers
            panel.bwplot2(...) },
        auto.key=TRUE)