import bisect
import random 
import time

class Epidemic: 
	def __init__(self,infectProb,contactRate,now=0.0,rng=random.Random(1)): 
		Epidemic.allIndividuals=[]
		Epidemic.eventList=Epidemic.EventList()
		Epidemic.contactRate=contactRate
		Epidemic.infectProb=infectProb
		Epidemic.now=now
		Epidemic.times=[now]
		Epidemic.numInfected=[0]
		Epidemic.rng=rng
		 
	def activate(self): 
		for ind in Epidemic.allIndividuals: 
			if ind.infected==0:
				Epidemic.eventList.insert(Epidemic.now+ind.contactInterval(),ind)
		
	def simulate(self,until):
		print str(Epidemic.now)+'\t'+str(Epidemic.numInfected[-1])
		while Epidemic.now <= until:
			try:
				thisEvent=Epidemic.eventList.pop()
			except IndexError:
				break
			Epidemic.now=thisEvent[0]
			thisEvent[1].makeContact()
						
	class EventList:
		def __init__(self):
			self.queue = []
		def insert(self, timestamp,data):
			""" Insert a new element in the queue according to its priority. """
			bisect.insort(self.queue, (timestamp, data))
		def pop(self):
			""" Pop the highest-priority element of the queue. """
			return self.queue.pop(0)
		def remove(self,timestamp,data):
			""" Remove item from queue. """
			item_delete_point = bisect.bisect(self.queue,(timestamp,data))
			del self.queue[item_delete_point-1]
		def is_present(self,timestamp,data):
			""" Determines whether item is present in queue. """
			item_insert_point = bisect.bisect(self.queue, (timestamp,data))
			is_present = self.queue[item_insert_point-1:item_insert_point] == [item]

	class Individual: 
		def __init__(self): 
			Epidemic.allIndividuals.append(self)
			
		def contactInterval(self):
			return Epidemic.rng.expovariate(Epidemic.contactRate)
		
		def chooseContact(self):
			contact = self
			while contact == self:
				contact=Epidemic.rng.choice(Epidemic.allIndividuals)
			return contact
			
		def makeContact(self):
			B=self.chooseContact()
			if B.infected==1 and Epidemic.rng.random() <= Epidemic.infectProb:
				Epidemic.times.append(Epidemic.now)
				Epidemic.numInfected.append(Epidemic.numInfected[-1]+1)
				print str(Epidemic.now)+'\t'+str(Epidemic.numInfected[-1])
				self.infected=1
			else:
				ci=self.contactInterval()
				Epidemic.eventList.insert(Epidemic.now+ci,self)
			

def model(nrIndividuals,contactRate,infectProb,initialInfected,runTime):
	run=Epidemic(infectProb=infectProb,contactRate=contactRate)
	for i in range(nrIndividuals):
		ind=Epidemic.Individual()
		if i >= (nrIndividuals-initialInfected):
			ind.infected=1
			Epidemic.numInfected[0]+=1
		else:
			ind.infected=0
	run.activate()
	run.simulate(until=runTime)

if __name__=="__main__":
	t1=time.clock()
	model(nrIndividuals=100,contactRate=3,infectProb=0.003,initialInfected=1,runTime=1000)            
	print time.clock()-t1