"""
    Ensimag 2018 TP Perf.
"""

from datetime import datetime
import simpy
import math
import pandas as pd
import numpy as np
from random import expovariate, randint, seed, random
import matplotlib.pyplot as plt
from enum import Enum

class State(Enum):
	IDLE = 0
	BUSY = 1

class Policy(Enum):
	PURE_ALOHA=0
	SLOTTED_ALOHA=1

class Packet(object):
	""" Packet structure

    Attributes:
    	id (int):
    		Packet identifier
        size (int): 
        	Packet size in Bytes.
		generation_timestamp (float): 
			Timestamp (simulated time) of packet generation
		output_timestamp (float): 
			Timestamp (simulated time) when packet leaves the 
		src_name (str):
			Name of source
	"""
	def __init__(self, id, size, generation_timestamp, src_name=""):
		self.id=id
		self.size=size
		self.generation_timestamp=generation_timestamp
		self.output_timestamp=0
		self.src_name=src_name
	def __str__(self):
		msg="Packet "+str(self.id)+" of size "+str(self.size)+ " sent at " \
		+str(self.generation_timestamp)+ " and arrived at "+str(self.output_timestamp)
		return msg

class Source(object):
	""" Packet generator

	Attributes:
		env (simpy.Environment):
			    Simulation environment
		name (str):
		    Name of the source
		gen_distribution (callable):
		    Function that returns the successive inter-arrival times of the packets
		size_distribution (callable):
		    Function that returns the successive sizes of the packets
		init_delay (int):
		    Starts generation after an initial delay. Default = 0
		destination (object):
		    Entity that receives the packets from the generator
	    stop_time (int):
	    	Number of simulation time steps starting from which the packets are not generated anymore
		debug (bool):
		    Set to true to activate verbose debug
	"""
	
	def __init__(self, env, name, init_delay=0, gen_distribution=lambda:1, size_distribution=lambda:1000,stop_time=math.inf,debug=False):
		self.env=env
		self.name=name
		self.init_delay=init_delay
		self.gen_distribution=gen_distribution
		self.size_distribution=size_distribution
		self.packet_count=0
		self.destination= None
		self.action= env.process(self.run())
		self.stop_time=stop_time
		self.debug=debug

	def run(self):
		""" Packet generation loop

		"""
		# Initial waiting time
		yield self.env.timeout(self.init_delay)
		while True:
			if self.env.now > self.stop_time:
				return
			interarrival_duration= self.gen_distribution()
			packet_size= self.size_distribution()
			yield self.env.timeout(interarrival_duration)
			generated_packet= Packet(id=self.packet_count,size=packet_size,generation_timestamp=env.now,src_name=self.name)
			if self.debug:
				print("Packet (id=%d,size=%d) generated by %s at %f" %(self.packet_count, packet_size, self.name,
					generated_packet.generation_timestamp))
			if self.destination is not None:
				if self.debug:
					print("%s => %s" % (self.name, self.destination.name))	
				env.process(self.destination.put(generated_packet)) # Packet is dropped if destination buffer is full
			self.packet_count+=1

	def attach(self, destination):
		""" Method to set a destination for the generated packets
		
		Args:
			destination (QueuedServer || PacketDestination): 
		"""
		self.destination= destination


class QueuedServer(object):
	""" Represents a waiting queue and an associated server.

	Attributes:
		env (simpy.Environment):
			Simulation environment
		name (str):
			Name of the source
		buffer (simpy.Store):
			Simpy FIFO queue
		buffer_max_size (int):
			Maximum buffer size in bytes
		buffer_size (int):
			Current size of the buffer in bytes
		service_rate (float):
			Server service rate in byte/sec
		destination (object):
			Entity that receives the packets from the server
		debug (bool):
			Set to true to activate verbose debug
		busy (bool):
			Is set if packet is currently processed by the server
		packet_count (int):
			Number of packet received
		packets_drop (int):
			Number of packets dropped
		has_collided (bool):
			If set the currently transmitted packet has collided
		attempt (int):
			Number of retransmission attempts
		max_attempt (int):
			Maximum number of retransmission before discarding
    """

	def __init__(self, env, name, channel, buffer_max_size=None,debug=False,attempt=0,max_attempt=10):
		self.env= env
		self.name= name
		self.buffer= simpy.Store(self.env,capacity=math.inf) # buffer size is limited by put method
		self.buffer_max_size= buffer_max_size
		self.buffer_size=0  
		self.packet_drop_count= 0
		self.destination=None
		self.debug=debug
		self.busy=0
		self.packet_count=0
		self.packets_drop=0
		self.channel=channel
		self.service_rate=self.channel.rate # override service rate
		self.has_collided=False
		self.action=env.process(self.run())
		self.attempt=attempt
		self.max_attempt=max_attempt

	def run(self):
		""" Packet waiting & service loop

		"""

		while True:
			packet = yield self.buffer.get()
			# If slotted aloha, wait for next slot to send
			if self.channel.policy == Policy.SLOTTED_ALOHA:
				now=self.env.now
				next_top=(round(now/self.channel.time_slot)+1)*self.channel.time_slot
				# Wait for begining of next time slot
				yield self.env.timeout(next_top-now)

			self.channel.add_sender(self)
			self.buffer_size-=packet.size
			service_time= packet.size/self.service_rate
			yield self.env.timeout(service_time)
			packet.output_timestamp= env.now

			if self.destination is not None:
				while self.has_collided and self.attempt <= self.max_attempt:
					if self.channel.policy is None:
						if self.debug:
							print("%.2f: Packet %d dropped because of collision" %(self.env.now,packet.id))
						self.packets_drop += 1
						break
					else:
						self.channel.remove_sender(self)
						self.has_collided=False 
						delay=random()*channel.backoff*channel.time_slot # Pick random backoff
						yield self.env.timeout(delay) # Wait 
						if self.channel.policy == Policy.SLOTTED_ALOHA:
							now=self.env.now
							next_top=(round(now/self.channel.time_slot)+1)*self.channel.time_slot
							# Wait for begining of next time slot
							yield self.env.timeout(next_top-now)

						self.channel.add_sender(self) # Start retransmitting
						yield self.env.timeout(service_time)
						packet.output_timestamp= env.now
						self.attempt += 1

				self.channel.remove_sender(self)
				self.has_collided=False 

				# drop packet if exceed maximum number of attempts 
				if self.attempt > self.max_attempt:
					self.packets_drop +=1
					if self.debug:
							print("%.2f: Packet %d dropped because of maximal retransmission attempts" %(self.env.now,packet.id))
				elif self.channel.policy is not None:
					self.destination.put(packet)

				# Reinitialize attempt counter	
				self.attempt=0
			else:
				self.channel.remove_sender(self)
				print("Packet dropped because no destination specified")

	def put(self, packet):
		self.packet_count += 1
		buffer_futur_size = self.buffer_size + packet.size
		if self.buffer_max_size is None or buffer_futur_size <= self.buffer_max_size:
			self.buffer_size = buffer_futur_size
			yield self.buffer.put(packet) 
			if self.debug and False:
				print("Packet %d added to queue %s." % (packet.id, self.name))
		else:
			self.packets_drop += 1
			if self.debug:
				print("Packet %d is discarded by queue %s. Reason: Buffer overflow." % (packet.id, self.name))

	def attach(self, destination):
		""" Method to set a destination for the serviced packets

		Args:
			destination (QueuedServer || PacketDestination):
		"""
		self.destination=destination


	def collided(self):
			self.has_collided=True


class PacketDestination(object):
	""" Receives packets and collects delay information

		Attributes:
			env (simpy.Environment):
				Simulation environment
			debug (boolean):
				Sets the debug mode
			arrived_packets (list[Packet]):
				List of received packets
			packet_count (int):
				Number of received packets
			byte_count (int):
				Number of received bytes
			avg_latency (float):
				Average latency between pkt source and destination

	"""

	def __init__(self, env, debug=False):
		self.env=env
		self.debug=debug
		self.arrived_packets = list()
		self.packet_count=0
		self.byte_count=0
		self.latency_sum=0
		self.avg_latency=0.0

	def put(self, pkt):
		now = self.env.now
		# update arrival time of packet
		pkt.output_timestamp = now
		self.latency_sum= self.latency_sum + (pkt.output_timestamp-pkt.generation_timestamp)
		self.arrived_packets.append(pkt)
		self.packet_count += 1
		self.byte_count += pkt.size
		self.avg_latency=self.latency_sum/self.packet_count
		if self.debug:
			print(pkt)


class QueuedServerMonitor(object):
	""" A monitor for a QueuedServer. Observes the packets in service and in 
		the queue and records that info in the sizes[] list. The monitor looks at the queued server 
		at time intervals given by the sampling dist.

	    
		Attributes:
		env (simpy.Environment):
			Simulation environment
		queued_server (QueuedServer):
			QueuedServer to monitor
		sample_distribution (callable):
			Function that returns the successive inter-sampling times
		sizes (list[int]):
			List of the successive number of elements in queue. Elements can be packet or bytes 
			depending on the attribute count_bytes 
		count_bytes (bool):
			If set counts number of bytes instead of number of packets 
	"""

	def __init__(self, env, queued_server, sample_distribution=lambda:1, count_bytes=False):
		self.queued_server= queued_server
		self.env= env
		self.sample_distribution= sample_distribution
		self.count_bytes= count_bytes
		self.sizes= []
		self.action= env.process(self.run())
		self.time_count=0

	def run(self):
		while True:
			yield self.env.timeout(self.sample_distribution())
			self.time_count+=1
			if self.count_bytes:
				total= self.queued_server.buffer_size
			else:
				total= len(self.queued_server.buffer.items) + self.queued_server.busy
			self.sizes.append(total)

class Channel(object):
	""" Communication channel

		Attributes:
		sending_list (List[QueuedServer]):
			List of routers currently transmitting
		rate (float):
			channel transmission rate
		collision (bool):
			If set collisions are considered
		debug (bool):
			If set debug messages are displayed
		state (State):
			Describes the state of the channel: idle or busy
		policy (Policy):
			Multiple access management policy
		backoff (float):
			Maximum value of the backoff
		time_slot (float):
			Average transmission window
	"""

	def __init__(self, rate, collision=False,debug=False,policy=None,backoff=0,time_slot=0):
		self.sending_list=list()
		self.rate=rate
		self.collision=collision
		self.debug=debug
		self.state=State.IDLE
		self.policy=policy
		self.backoff=backoff
		self.time_slot=time_slot

	def add_sender(self, sender):
		self.sending_list.append(sender)
		self.state=State.BUSY	
		if self.collision and len(self.sending_list)>=2:
			self.collision_broadcast()
		if self.debug:
			print("Channel is busy because %s is transmitting at %f" %(sender.name,env.now))


	def remove_sender(self, sender):
		self.sending_list.remove(sender)
		if len(self.sending_list)==0:
			self.busy=State.IDLE
			if self.debug:
				print("Channel is idle")

	def collision_broadcast(self):
		for sender in self.sending_list:
			sender.collided()

if __name__ == "__main__":



	# If set collisions are considered in the simulation
	collision= True
	# If collision is true, select an management policy (None,Policy.PURE_ALOHA,Policy.SLOTTED_ALOHA)
	policy= Policy.PURE_ALOHA

	# Simulation configuration
	sim_step=1e-3 # in seconds
	sim_duration=1000/sim_step # 1000 time steps
	# Stop packet generation at this time
	stop_generation=sim_duration/100

	# Link capacity 64kbps
	process_rate = (64000/8)*sim_step  # => 8 kBps
	average_packet_size=400
	# Average transmission time
	time_slot=average_packet_size/process_rate 	

	# Packet length exponentially distributed with average 400 bytes
	# dist_size = lambda: expovariate(1/average_packet_size) # comment if not exponential size distribution
	dist_size= lambda: average_packet_size # comment if not constant size distribution
	# Packet inter-arrival time exponentially distributed
	gen_dist1 = lambda: expovariate(7.5*sim_step) 
	gen_dist2 = lambda: expovariate(7.5*sim_step) 
	avg_latency_backoff = {}
	drop_ratio_backoff={}
	backoff_set = np.linspace(1, 100, num=2)
	for backoff in backoff_set:
		env = simpy.Environment()
		# For reproducibility
		seed(1)

		# Simulation entities
		channel = Channel(rate=process_rate,debug=False,collision=collision,\
			policy=policy,backoff=backoff,time_slot=time_slot)
		src1 = Source(env, "Source 1", gen_distribution=gen_dist1, size_distribution=dist_size, debug=False,stop_time=stop_generation)
		src2 = Source(env, "Source 2", gen_distribution=gen_dist2, size_distribution=dist_size, debug=False,stop_time=stop_generation)
		qs1 = QueuedServer(env, "Router 1", buffer_max_size=math.inf, channel=channel, debug=False, max_attempt=1)
		qs2 = QueuedServer(env, "Router 2", buffer_max_size=math.inf, channel=channel, debug=False, max_attempt=1)
		dst = PacketDestination(env,debug=False)

		# Link Source 1 to Router 
		src1.attach(qs1)
		# Link Source 2 to Router 2
		src2.attach(qs2)
		# Channel output
		qs1.attach(dst)
		qs2.attach(dst)
		# Associate a monitor to Router 1
		qs1_monitor = QueuedServerMonitor(env, qs1, sample_distribution=lambda: 1, count_bytes=False)
		# Associate a monitor to Router 2
		qs2_monitor = QueuedServerMonitor(env, qs2, sample_distribution=lambda: 1, count_bytes=False)

		if collision is False:
			print("#### Simulation without collision ####")
		elif policy is None:
			print("#### Simulation with no multiple access management ####")
		elif policy == Policy.PURE_ALOHA:
			print("#### Simulation with Pure Aloha ####")
		elif policy == Policy.SLOTTED_ALOHA:
			print("#### Simulation with Slotted Aloha ####")

		env.run(until=sim_duration)

		# Stats example
		nb_packets=src2.packet_count+src1.packet_count
		print("Sent %d packet(s)" %nb_packets)
		print("Received %d packet(s)" %dst.packet_count)
		drop_ratio=(nb_packets-dst.packet_count)/nb_packets*100.0
		print("The drop ratio is : %.2f%%" %drop_ratio)
		avg_latency=dst.avg_latency*sim_step
		print("The average latency is %.2f s" %avg_latency)
		avg_latency_backoff[backoff]=avg_latency
		drop_ratio_backoff[backoff]=drop_ratio

		if collision is False:
			print("#### End of simulation without collision ####")
		elif policy is None:
			print("#### End of simulation with no multiple access management ####")
		elif policy == Policy.PURE_ALOHA:
			print("#### End of simulation with Pure Aloha ####")
		elif policy == Policy.SLOTTED_ALOHA:
			print("#### End of simulation with Slotted Aloha ####")

	# Stat visualization examples

	# Latenct vs backoff
	# plt.bar(avg_latency_backoff.keys(), avg_latency_backoff.values(), color='g')
	# plt.title('Latency vs Backoff')
	# plt.xlabel('Backoff')
	# plt.ylabel('Latency (s)')
	# plt.show()

	# Loss vs backoff
	plt.bar(drop_ratio_backoff.keys(), drop_ratio_backoff.values(), color='g')
	plt.title('Loss vs Backoff')
	plt.xlabel('Backoff')
	plt.ylabel('Loss (%)')
	plt.show()





