"""
Last modified on 07 May 2020 by Ahmet Sekercioglu. 
Based on the original script by:

Copyright 2014 Dr. Greg M. Bernstein
Released under the MIT license
"""
import random
import functools

import simpy

from SimComponents import PacketGenerator, PacketSink, SwitchPort, RandomBrancher

def run_sim():
    # Set up arrival and packet size distributions
    mean_pkt_size = 100.0  # in bytes
    adist1 = functools.partial(random.expovariate, 2.0)
    adist2 = functools.partial(random.expovariate, 0.5)
    adist3 = functools.partial(random.expovariate, 0.6)
    sdist  = functools.partial(random.expovariate, 1.0/mean_pkt_size)
    samp_dist = functools.partial(random.expovariate, 0.50)
    port_rate = 2.2*8*mean_pkt_size  # want a rate of 2.2 packets per second

    # Create the SimPy environment. This is the simulation infrastructure.
    env = simpy.Environment()

    # Create the packet generators and sink
    pk_gen_1 = PacketGenerator(env, "EE", adist1, sdist)
    pk_gen_2 = PacketGenerator(env, "CS", adist2, sdist)

    pk_sink_1 = PacketSink(env, debug=False, rec_arrivals=True)
    pk_sink_2 = PacketSink(env, debug=False, rec_waits=True)

    branch1 = RandomBrancher(env, [0.75, 0.25])

    switch_port1 = SwitchPort(env, port_rate)
    switch_port2 = SwitchPort(env, port_rate)
    switch_port3 = SwitchPort(env, port_rate)

    # Wire packet generators, switch ports, and sinks together
    pk_gen_1.out = switch_port1
    switch_port1.out = branch1
    branch1.outs[0] = switch_port2
    branch1.outs[1] = switch_port3
    switch_port2.out = pk_sink_1
    switch_port3.out = pk_sink_2
    pk_gen_2.out = switch_port3

    # Run it
    env.run(until = 4000)

    print("Average wait pk_gen_1 to pk_sink_1 = {}".format(sum(pk_sink_1.waits)/len(pk_sink_1.waits)))
    print("Total pkts sent = {}".format(pk_gen_1.packets_sent + pk_gen_2.packets_sent))
    print("Total pkts received = {}".format(len(pk_sink_1.waits) + len(pk_sink_2.waits)))
    # print "average system occupancy: {}".format(float(sum(pm.sizes))/len(pm.sizes))

if __name__ == '__main__':
    run_sim()