# Notebook from kadupitiya/RNN-MD
Path: src/md-codes/many_particle_code.ipynb

# Define vector 3D class_____no_output_____

import math
import numpy as np

class Vector3D:
    def __init__(self, initial_x = 0.0, initial_y = 0.0, initial_z = 0.0):
        self.x = initial_x
        self.y = initial_y
        self.z = initial_z

    def magnitude(self):
        return math.sqrt(self.x**2 + self.y**2 + self.z**2)

    def sqd_magnitude(self):
        return self.x**2 + self.y**2 + self.z**2

    # Operator overloading for adding two vecs
    def __add__(self, v):
        return Vector3D(self.x + v.x, self.y + v.y, self.z + v.z)

    def __mul__(self, scalar):
        return Vector3D(self.x*scalar, self.y*scalar, self.z*scalar)

    def __sub__(self, v):
        return Vector3D(self.x - v.x, self.y - v.y, self.z - v.z)

    def __eq__(self, other):
        if isinstance(other, Vector3D):
            return self.x == other.x and self.y == other.y and self.z == other.z

    #printing overloaded
    def __str__(self):
        return "x=" + str(self.x) + ", y=" + str(self.y) + ", z=" + str(self.z)

    def apply_periodic_boundries(self, lx, ly, lz):
        # periodic boundary
        if (self.x > lx/2.0):
            self.x = self.x - lx
        if (self.x < -lx/2.0):
            self.x = self.x + lx
        if (self.y > ly/2.0):
            self.y = self.y - ly
        if (self.y < -ly/2.0):
            self.y = self.y + ly
        if (self.z > lz/2.0):
            self.z = self.z - lz
        if (self.z < -lz/2.0):
            self.z = self.z + lz
        _____no_output_____

# Define Particle class_____no_output_____

import math

class Particle:

    def __init__(self, id_value=0, diameter = 0.0, initial_m = 0.0,  initial_position = Vector3D(0.0, 0.0, 0.0), initial_velocity = Vector3D(0.0, 0.0, 0.0), initial_lx = 0.0, initial_ly = 0.0, initial_lz = 0.0):
        self.id = id_value
        self.m = initial_m
        self.d = diameter
        self.position = initial_position
        self.velocity = initial_velocity
        self.lx = initial_lx;
        self.ly = initial_ly;
        self.lz = initial_lz;
        self.pe = 0.0
        self.ke = 0.0
        self.force = Vector3D(0.0, 0.0, 0.0)

    def volume(self):
        self.volume = (4.0/3.0) * math.pi * ((self.d/2.0)**3)

    #position updated to a full time-step
    def update_position(self, dt, box):
        self.position = self.position + (self.velocity * dt)
        # periodic boundary
        self.position.apply_periodic_boundries(box.lx, box.ly, box.lz)

    #velocity updated to a half time-step
    def update_velocity(self, dt):
        self.velocity = self.velocity + (self.force * ( 0.5 * dt / self.m))

    def kinetic_energy(self):
        self.ke = 0.5 * self.m * (self.velocity.magnitude()**2)
_____no_output_____

# Define Box class_____no_output_____

import random
class Box:

    def __init__(self, initial_position = Vector3D(0.0, 0.0, 0.0), m_lx = 10.0, m_ly = 10.0, m_lz = 10.0):
        self.position = initial_position #origin point of the simulation box
        self.lx = m_lx                        #length of the box in x direction
        self.ly = m_ly                        #length of the box in y direction
        self.lz = m_lz                        #length of the box in z direction

    def put_ljatoms_inside(self, number_ljatom=0, ljatom_diameter=0.0, ljatom_mass=0.0, density=0.0):
        self.ljatom_diameter=ljatom_diameter
        #initialize the positions of the particles on a lattice (not randomly)
        #out_initialize_position = open("many_particle_data/ip.lammpstrj", "w")
        print("# of atoms created in the box = {}".format(number_ljatom))
        a = 1
        print("# lattice spacing (based on crystal of closest packing) is = {}".format(a))

        num_atoms_linear = int(math.ceil(math.pow((float(number_ljatom)),1.0/3.0)))
        print("# in each direction, about {} atoms".format(num_atoms_linear))

        self.ljatom = []
        # initial ion positions
        for i in range(num_atoms_linear):
            for j in range(num_atoms_linear):
                for k in range(num_atoms_linear):
                    # stop if atoms created are equal to the requested number
                    if len(self.ljatom) < number_ljatom:
                        x = (-self.lx/2) + i*a
                        y = (-self.ly/2) + j*a
                        z = (-self.lz/2) + k*a
                        if (z >= (self.lz/2.0 - a/2.0) or y >= (self.ly/2.0 - a/2.0 ) or x >= (self.lx/2.0 - a/2.0)):
                            continue
                        position = Vector3D(x,y,z)
                        velocity = Vector3D(0.0,0.0,0.0)
                        fresh_atom = Particle(len(self.ljatom)+1, self.ljatom_diameter, ljatom_mass, position, velocity, self.lx, self.ly, self.lz)
                        self.ljatom.append(fresh_atom)

        #out_initialize_position.close()

        #make a snapshot movie of initial positions of the particles
        self.make_movie(num=0, file_name="final_ip.lammpstrj")

        return len(self.ljatom)

    def put_ljatoms_inside_random(self, number_ljatom=0, ljatom_diameter=0.0, ljatom_mass=0.0, density=0.0, edge_length=0.0):
        self.ljatom_diameter=ljatom_diameter
        #fake_diameter = 0.01
        #initialize the positions of the particles on a lattice (not randomly)
        #out_initialize_position = open("many_particle_data/ip.lammpstrj", "w")
        print("# of atoms created in the box = {}".format(number_ljatom))
        a = 1.0
        print("# lattice spacing (based on crystal of closest packing) is = {}".format(a))
        try_count_limit = 1000000*number_ljatom
        self.ljatom = []
        # initial ion positions
        try_count = 0
        while len(self.ljatom) < number_ljatom and try_count < try_count_limit:

            xx = np.random.uniform(-edge_length/2 , edge_length/2 ,1)[0]
            yy = np.random.uniform(-edge_length/2 , edge_length/2 ,1)[0]
            zz = np.random.uniform(-edge_length/2 , edge_length/2 ,1)[0]

            position = Vector3D(xx,yy,zz)
            sanity_check = 0

            for i in range(0, len(self.ljatom),1):
                if abs(self.ljatom[i].position.magnitude() - position.magnitude()) >  self.ljatom_diameter:
                    sanity_check += 1
                else:
                    break

            if len(self.ljatom) == sanity_check:
                velocity = Vector3D(0.0,0.0,0.0)
                fresh_atom = Particle(len(self.ljatom)+1, self.ljatom_diameter, ljatom_mass, position, velocity, self.lx, self.ly, self.lz)
                self.ljatom.append(fresh_atom)
                print(position)
                print(len(self.ljatom))

            try_count += 1

            if try_count == try_count_limit:
                print("Cant put ions inside the box ")
                print(try_count)
                self.ljatom = []
                return 0

        #out_initialize_position.close()

        #make a snapshot movie of initial positions of the particles
        self.make_movie(num=0, file_name="final_ip1.lammpstrj")

        return len(self.ljatom)

    #make movie
    def make_movie(self, num=0, file_name = ''):
        if num==0:
            outdump = open("many_particle_data/"+file_name, "w")
        else:
            outdump = open("many_particle_data/"+file_name, "a")
        outdump.write("ITEM: TIMESTEP\n")
        outdump.write("{}\n".format(num - 1))
        outdump.write("ITEM: NUMBER OF ATOMS\n")
        outdump.write("{}\n".format(len(self.ljatom)))
        outdump.write("ITEM: BOX BOUNDS\n")
        outdump.write("{}\t{}\n".format(-0.5*self.lx, 0.5*self.lx))
        outdump.write("{}\t{}\n".format(-0.5*self.ly, 0.5*self.ly))
        outdump.write("{}\t{}\n".format(-0.5*self.lz, 0.5*self.lz))
        outdump.write("ITEM: ATOMS index type x y z v\n")
        for i in range(len(self.ljatom)):
            outdump.write("{}\t{}\t{}\t{}\t{}\t{}\n".format(i+1, "1", self.ljatom[i].position.x, self.ljatom[i].position.y, self.ljatom[i].position.z, self.ljatom[i].velocity.magnitude()))
        outdump.close()

    #Velocity init: Currenly zero
    def initialize_particle_velocities(self):
        print("Velocity init")

    #Force computation
    def update_forces(self, dcut=0.0):
        self.dcut = dcut
        elj = 2.0
        dcut2 = self.dcut*self.dcut
        d = 1 #recall that we are working in reduced units where the unit of length is the diameter of the particle
        d2 = d*d
        d6 = d2 * d2 * d2
        d12 = d6 * d6
        d_cut_off = dcut2 * d2

        for i in range(len(self.ljatom)):
            flj = Vector3D(0.0,0.0,0.0)
            for j in range(len(self.ljatom)):
                if (j == i):
                    continue
                r_vec = self.ljatom[i].position - self.ljatom[j].position
                #Add periodic boundries
                r_vec.apply_periodic_boundries(self.lx, self.ly, self.lz)
                #r = r_vec.magnitude()
                r2 = r_vec.sqd_magnitude()
                #Cut off check
                if (r2 < d_cut_off):
                    r6 = r2 * r2 * r2
                    r12 = r6 * r6
                    flj += ( r_vec * ( 48 * elj * (  (d12 / r12)  - 0.5 *  (d6 / r6) ) * ( 1 / r2 ) ) )
            self.ljatom[i].force = flj

    #Potential Energy computation
    def update_p_energies(self, dcut=0.0):
        self.dcut = dcut
        elj = 2.0
        dcut2 = self.dcut  * self.dcut
        dcut6 = dcut2 * dcut2 * dcut2
        dcut12 = dcut6 * dcut6
        d = 1 #recall that we are working in reduced units where the unit of length is the diameter of the particle
        d2 = d * d
        d6 = d2 * d2 * d2
        d_cut_off = dcut2 * d2

        energy_shift = 4*elj*((1/dcut12) - (1/dcut6))
        total_lj_atom_atom = 0.0

        for i in range(len(self.ljatom)):
            uljpair = 0.0
            for j in range(len(self.ljatom)):
                if (j == i):
                    continue
                r_vec = self.ljatom[i].position - self.ljatom[j].position
                #Add periodic boundries
                r_vec.apply_periodic_boundries( self.lx, self.ly, self.lz)
                r2 = r_vec.sqd_magnitude()
                #Cut off check
                if (r2 < d_cut_off):
                    r6 = r2 * r2 * r2
                    uljpair +=  4 * elj * (d6 / r6) * ( ( d6 / r6 ) - 1 ) - energy_shift

            self.ljatom[i].pe = uljpair
            total_lj_atom_atom += uljpair
        total_lj_atom_atom = 0.5*total_lj_atom_atom
        return total_lj_atom_atom

    #kinatic Energy computation
    def update_k_energies(self):
        total_ke_atom = 0.0
        for i in range(len(self.ljatom)):
            self.ljatom[i].kinetic_energy()
            total_ke_atom += self.ljatom[i].ke
        return total_ke_atom

    #update half time velocity
    def update_velocity(self, delta_t):
        for i in range(len(self.ljatom)):
            self.ljatom[i].update_velocity(delta_t)

    #update full time velocity
    def update_position(self, delta_t):
        for i in range(len(self.ljatom)):
            self.ljatom[i].update_position(delta_t, self)    _____no_output_____

# Other functions_____no_output_____

#display progress
def progressBar(fraction_completed):
    val = int(fraction_completed * 100)
    lpad = int(fraction_completed * 60)
    sys.stdout.write("\r[%-60s] %d%%" % ('='*lpad, val))_____no_output_____

# Velocity verlet Function_____no_output_____

import math
import time
import sys
import os
import shutil
import time

def velocity_verlet(ljatom_density=None, number_ljatom=None):
    #we begin with some fundamental numbers/physical constants
    pi = 3.141593 #Pi
    kB = 1.38e-23 #Joules per Kelvin
    mol = 6.0e23 #Avogadro number

    #reduced units
    unitlength = 0.3405e-9      # m; unit of length is these many meters (diameter of a typical atom)
    unitenergy = 119.8 * kB     #Joules; unit of energy is these many Joules (typical interaction strength)
    unitmass = 0.03994 / mol    #kg; unit of mass (mass of a typical atom)
    unittime = math.sqrt(unitmass * unitlength * unitlength / unitenergy)    #Unit of time
    unittemperature = unitenergy/kB       #unit of temperature

    #essential parameters needed for particles in a box simulation
    #**********************
    #ljatom_density=0.0          #this is taken in as mass density (in reduced units)
    #number_ljatom=0             #total number of particles in your system
    ljatom_diameter=0.0
    ljatom_mass=0.0
    bx=0.0                      #box edge lengths
    by=0.0
    bz=0.0
    temperature=0.0
    #dcut=0.0                    #cutoff distance for potential in reduced units
    #**********************

    #shutil.rmtree('many_particle_data',True)
    if not os.path.exists('many_particle_data'):
        os.makedirs('many_particle_data')

    print("Program starts")
    print("units of length, energy, mass are given near the top of the main function")
    print("unit of length is {} meters".format(unitlength))
    print("unit of mass is {} kilograms".format(unitmass))
    print("unit of energy is {} Joules".format(unitenergy))
    print("unit of time is {} seconds".format(unittime))

    # inputs 1
    #this is taken in as mass density (in reduced units)
    if ljatom_density is None:
        print("enter density (in LJ reduced units: recall unit of length is diameter of your lj particle); tested for rho = 0.8442  ")
        time.sleep(0.1) # This sleep is not needed, just added to get input box below the print statements
        ljatom_density = float(input())
    # inputs 2
    #total number of particles in your system
    if number_ljatom is None:
        print("enter total number of ljatoms (108 is a typical value; tested for this value) ")
        time.sleep(0.1)
        number_ljatom = float(input())
    dcut=2.5
    ljatom_diameter = 1.0         #in reduced units
    ljatom_mass = 1.0

    edge_length = math.pow(number_ljatom/ljatom_density,1.0/3.0)
    edge_length_ = edge_length
    #edge_length_ = 20
    print("edge length calculated to be: {}".format(edge_length))

    bx = edge_length_
    by = edge_length_
    bz = edge_length_
    #place where the action happens, periodic boundaries
    simulation_box = Box(Vector3D(0.0, 0.0, 0.0), bx, by, bz)
    ljatom_size = simulation_box.put_ljatoms_inside(number_ljatom, ljatom_diameter, ljatom_mass, ljatom_density)

    multi = 4
    simulation_box.lx = edge_length_*multi                        #length of the box in x direction
    simulation_box.ly = edge_length_*multi                        #length of the box in y direction
    simulation_box.lz = edge_length_*multi

    #ljatom_size = simulation_box.put_ljatoms_inside_random(number_ljatom, ljatom_diameter, ljatom_mass, ljatom_density, 4.0)

    #output to screen the parameters of the problem
    print("Box dimensions x = {} | y = {} | z = {} ".format(simulation_box.lx, simulation_box.ly, simulation_box.lz))
    print("Number of LJ atoms inside the box = {} ".format(ljatom_size))
    print("LJ atom diameter (in SI units; meter) = {} ".format(ljatom_diameter*unitlength))
    print("LJ atom mass (in SI units; kg) = {} ".format(ljatom_mass*unitmass))
    print("cut off distance (in reduced units) = {} ".format(dcut))

    #simulation_box.initialize_particle_velocities()
    #Inital force computation
    simulation_box.update_forces(dcut)

    totaltime = 20
    #number of time discretizations (slices)
    steps = 20000
    #the code deterimes the time-step delta_t. choose steps carefully, make sure you have a fine discretization
    delta_t = totaltime/steps
    print("timestep (reduced units) = {}".format(delta_t))
    print("timestep (in SI units; second) = {}".format(delta_t*unittime))
    print("total simulation time (in real units) = {}".format(totaltime*unittime))

    movie_file_name = "propagation1.lammpstrj"
    #energy_filename = "energy_rho={}_N={}.out".format(ljatom_density, number_ljatom)
    energy_filename = "energy1.dat"

    #Initial energies
    totalpe = simulation_box.update_p_energies(dcut)
    totalke = simulation_box.update_k_energies()

    with open("many_particle_data/"+energy_filename, "w") as output_energy:
        output_energy.write("{}\t{}\t{}\t{}\n".format(0, totalke/ljatom_size, totalpe/ljatom_size, (totalke+totalpe)/ljatom_size))

    print("initial total energy per lj particle is = {}".format( (totalke+totalpe)/ljatom_size))

    #Initial movie dump
    simulation_box.make_movie(num=0, file_name="propagation1.lammpstrj")

    movie_frequency = 100
    average_pe = 0.0
    average_ke = 0.0
    data_collect_frequency = 1000
    samples = 0;
    hit_eqm = 3000               #this is your choice of where you think the system hit equilibrium

    #Molecular Dynamics
    print("progress...")

    for num in range(1, steps+1):
        #velocity-Verlet
        #update velocity half timestep
        simulation_box.update_velocity(delta_t)
        #update position full timestep
        simulation_box.update_position(delta_t)
        #expensive step
        simulation_box.update_forces(dcut)
        totalpe = simulation_box.update_p_energies(dcut)
        #update other half velocity
        simulation_box.update_velocity(delta_t)
        totalke = simulation_box.update_k_energies()

        #calling a movie function to get a movie of the simulation
        if (num%movie_frequency == 0):
            simulation_box.make_movie(num=num, file_name=movie_file_name)

        #outputting the energy to make sure simulation can be trusted
        if (num%data_collect_frequency == 0):
            with open("many_particle_data/"+energy_filename, "a") as output_energy:
                output_energy.write("{}\t{}\t{}\t{}\n".format(num, totalke/ljatom_size, totalpe/ljatom_size, (totalke+totalpe)/ljatom_size))

        if (num > hit_eqm and num%data_collect_frequency == 0):
            average_pe += totalpe
            average_ke += totalke
            samples += 1

        if num%100==0:
            fraction_completed = num/steps
            progressBar(fraction_completed)

    print("")
    print("average pe per particle is = {}".format((average_pe/samples)/ljatom_size))
    print("average ke per particle is = {}".format((average_ke/samples)/ljatom_size))
    print("average T per particle is = {}".format(2*(average_ke/samples)/(3*number_ljatom)))
    print("Simulation is over.")_____no_output_____

# Run the code_____no_output_____

start = time.time()
# Run the program
#name__ = (0.8442, 108)
#velocity_verlet(*name__)
velocity_verlet(ljatom_density=0.5, number_ljatom=2)

end = time.time()
print("Time: "+str(end - start))Program starts
units of length, energy, mass are given near the top of the main function
unit of length is 3.405e-10 meters
unit of mass is 6.656666666666667e-26 kilograms
unit of energy is 1.65324e-21 Joules
unit of time is 2.1606159067248715e-12 seconds
edge length calculated to be: 1.5874010519681994
# of atoms created in the box = 2
# lattice spacing (based on crystal of closest packing) is = 1
# in each direction, about 2 atoms
Box dimensions x = 6.3496042078727974 | y = 6.3496042078727974 | z = 6.3496042078727974
Number of LJ atoms inside the box = 2
LJ atom diameter (in SI units; meter) = 3.405e-10
LJ atom mass (in SI units; kg) = 6.656666666666667e-26
cut off distance (in reduced units) = 2.5
timestep (reduced units) = 0.001
timestep (in SI units; second) = 2.1606159067248715e-15
total simulation time (in real units) = 4.321231813449743e-11
initial total energy per lj particle is = 0.016316891136
progress...
[============================================================] 100%
average pe per particle is = -0.030603360206919582
average ke per particle is = 0.04681120344724756
average T per particle is = 0.031207468964831706
Simulation is over.
Time: 0.6707022190093994

# Plot the energy graphs_____no_output_____

# Visualize the data
name__ = (0.8442, 3)
import matplotlib.pyplot as plt
#%matplotlib notebook
%matplotlib inline

import numpy as np
#file_name = 'energy_rho='+str(name__[0])+'_N='+str(name__[1])+'.out'
file_name = 'energy1.dat'
result_file = np.loadtxt("many_particle_data/"+file_name)

fig=plt.figure(figsize=(16, 6))
plt.title(file_name)
plt.plot(result_file[:,0],result_file[:,1], label='ke')
plt.plot(result_file[:,0],result_file[:,2], label='pe')
plt.plot(result_file[:,0],result_file[:,3], label='total')

plt.legend()_____no_output_____

# Play the movie_____no_output_____

sys.path.append('../../scripts/')
from viz_src.dump import dump
from viz_src.xyz import xyz
from viz_src.svg import svg
import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation
import pandas as pd
from IPython.display import HTML
%matplotlib notebook
%matplotlib notebook

class Particle_viz:
    def __init__(self, movie_file_name = None):

        self.data = dump(movie_file_name)
        self.data.map(1,"id")
        self.time_frames = self.data.time()
        #selecting last index from sorted time frames
        time_value = self.time_frames[-1]
        time_index = self.data.findtime(time_value)
        last_snap_shot = self.data.viz(time_index)

        self.box_x_lim = last_snap_shot[1][0::3]
        self.box_y_lim = last_snap_shot[1][1::3]
        self.box_z_lim = last_snap_shot[1][2::3]

        self.all_data_frames = {}
        for t_ in self.time_frames:
            time_index = self.data.findtime(t_)
            snap_shot = np.array(self.data.viz(time_index)[2])
            self.all_data_frames[t_] = snap_shot[:,2:5]

        #%matplotlib inline
        #figsize=(8, 8)
        self.fig = plt.figure(num='MD simulation output', figsize=(9, 8))
        self.ax = self.fig.add_subplot(111, projection='3d')
        #self.fig.suptitle('3D Test', fontsize=16)
        self.title = self.ax.set_title('Frame=-1')

        # Setting the axes properties
        self.ax.set_xlim3d(self.box_x_lim)
        self.ax.set_xlabel('X')

        self.ax.set_ylim3d(self.box_y_lim)
        self.ax.set_ylabel('Y')

        self.ax.set_zlim3d(self.box_z_lim)
        self.ax.set_zlabel('Z')

        #self.ax.set_xticklabels([])
        #self.ax.set_xticks([])
        #self.ax.set_yticklabels([])
        #self.ax.set_yticks([])
        #self.ax.set_zticklabels([])
        #self.ax.set_zticks([])

    def update_graph(self, num):
        data_frame_=self.all_data_frames[self.time_frames[num]]
        self.graph.set_data (data_frame_[:,0],data_frame_[:,1])
        self.graph.set_3d_properties(data_frame_[:,2])
        self.title.set_text('Frame={}'.format(self.time_frames[num]))
        return self.title, self.graph,

    def start(self):
        self.graph, = self.ax.plot(self.all_data_frames[self.time_frames[0]][:,1], self.all_data_frames[self.time_frames[0]][:,2], self.all_data_frames[self.time_frames[0]][:,0],
                 c='r', marker='o',  linestyle="", markersize=10)

        # Creating the Animation object
        self.animation = matplotlib.animation.FuncAnimation(fig=self.fig, func=self.update_graph, frames=len(self.time_frames),
                              interval=50, blit=True)
        #self.ax.plt.show()
        #animation_.save('matplot003.gif', writer='imagemagick')
        #animation_.save('myvideo.mp4', codec='h264')
        #HTML(animation_.to_html5_video())
        _____no_output_____viz_sim = Particle_viz(movie_file_name = "many_particle_data/propagation1.lammpstrj")
viz_sim.start()
#viz_simanimation.save('matplot003.gif', writer='imagemagick')
#viz_simanimation.save('myvideo.mp4', codec='h264')
#HTML(viz_sim.animation.to_html5_video())_____no_output_____


# Notebook from albeckr/Thesis_CoRoDro
Path: Exp_BASE.ipynb

# MultiAgentEnvironment: simple Map_____no_output_____

#%autosave 30
import glob
import sys
from operator import itemgetter
import numpy as np
import random
import time
import math

import networkx as nx
from networkx.algorithms.shortest_paths.generic import shortest_path_length

import ray
from ray import tune
#from ray.tune.logger import pretty_print
from ray.rllib.policy.policy import Policy, PolicySpec
from ray.rllib.models.tf.tf_modelv2 import TFModelV2
from ray.rllib.models.tf.fcnet import FullyConnectedNetwork

#from ray.rllib.agents.callbacks import DefaultCallbacks
from ray.rllib.env import MultiAgentEnv
from ray.rllib.utils.framework import try_import_tf

from gym.spaces import Discrete, Box, Tuple, MultiDiscrete, Dict, MultiBinary
from ray.rllib.utils.spaces.repeated import Repeated

import matplotlib.pyplot as plt

tf1, tf, tfv = try_import_tf()  # prefered TF import for Ray.io

from threading import Thread, Event

print("Imports successful")Imports successful
######## Utility Classes ########
class BoolTimer(Thread):
    """A boolean value that toggles after a specified number of seconds:

    Example:
        bt = BoolTimer(30.0, False)
        bt.start()

    Used in the Centrality Baseline to limit the computation time.
    """

    def __init__(self, interval, initial_state=True):
        Thread.__init__(self)
        self.interval = interval
        self.state = initial_state
        self.finished = Event()

    def __bool__(self):
        return bool(self.state)

    def run(self):
        self.finished.wait(self.interval)
        if not self.finished.is_set():
            self.state = not self.state
        self.finished.set()

######## Static helper functions ########
def shuffle_actions(action_dict, check_space = False):
    """
        Used to shuffle the action dict to ensure that agents with a lower id are not always preferred
        when picking up parcels over other agents that chose the same action.
        For debugging: Can also be used to check if all actions are in the action_space.
    """
    keys = list(action_dict)
    random.shuffle(keys)
    shuffled = {}

    for agent in keys:
        if check_space:  #assert actions are in action space -> disable for later trainig, extra check while development
            assert self.action_space.contains(action_dict[agent]),f"Action {action_dict[agent]} taken by agent {agent} not in action space"
        shuffled[agent] = action_dict[agent]

    return shuffled

def load_graph(data):
    """Loads topology (map) from json file into a networkX Graph and returns the graph"""

    nodes = data["nodes"]
    edges = data["edges"]

    g = nx.DiGraph()  # directed graph
    g.add_nodes_from(nodes)
    g.edges(data=True)

    for node in nodes:  # add attribute values
        g.nodes[node]["id"] = nodes[node]["id"]
        g.nodes[node]["type"] = nodes[node]["type"]

    for edge in edges:  # add edges with attributes
        f = edges[edge]["from"]
        t = edges[edge]["to"]

        weight_road, weight_air, _type = sys.maxsize, sys.maxsize, None

        if edges[edge]["road"] >= 0:
            weight_road = edges[edge]["road"]
            _type = 'road'
        if edges[edge]["air"] >= 0:
            weight_air = edges[edge]["air"]
            _type = 'both' if _type == 'road' else 'air'

        weight = min(weight_road, weight_air)  # needed for optimality baseline

        g.add_edge(f, t, type=_type, road= weight_road, air=weight_air, weight=weight)

    return g_____no_output_____

## Environment Code_____no_output_____

# Map definition
topology = {
    'nodes': {
            0: {'id': 1, 'type': 'parking'},
            1: {'id': 1, 'type': 'parking'},
            2: {'id': 2, 'type': 'parking'},
            3: {'id': 3, 'type': 'parking'},
            4: {'id': 4, 'type': 'parking'},
            5: {'id': 5, 'type': 'air'},
            6: {'id': 6, 'type': 'parking'},
            7: {'id': 7, 'type': 'air'},
            8: {'id': 8, 'type': 'parking'},
            9: {'id': 9, 'type': 'parking'},
            10: {'id': 10, 'type': 'parking'},
            11: {'id': 11, 'type': 'air'}
    },
    'edges': {
    ## Outer Ring    --> Road much (!) faster than drones
    "e01":{"from": 0,"to": 1,"road": 10, "air": 40},
    "e02":{"from": 1,"to": 0,"road": 10, "air": 40},
    "e03":{"from": 1,"to": 4,"road": 8, "air": 6},
    "e04":{"from": 4,"to": 1,"road": 8, "air": 6},
    "e03":{"from": 4,"to": 2,"road": 8, "air": 6},
    "e04":{"from": 2,"to": 4,"road": 8, "air": 6},
    "e05":{"from": 2,"to": 3,"road": 10, "air": 30},
    "e06":{"from": 3,"to": 2,"road": 10, "air": 30},
    "e07":{"from": 3,"to": 0,"road": 10, "air": 40},
    "e08":{"from": 0,"to": 3,"road": 8, "air": 35},
    ## Inner Nodes  --> Reinfahren langsamer als rausfahren
    ##              -->
    "e11":{"from": 4,"to": 6,"road": 2, "air": 5},
    "e12":{"from": 6,"to": 4,"road": 2, "air": 5},
    "e13":{"from": 0,"to": 6,"road": 6, "air": -1},
    "e14":{"from": 6,"to": 0,"road": 6, "air": -1},
    "e15":{"from": 3,"to": 6,"road": 6, "air": -1},
    "e16":{"from": 6,"to": 3,"road": 6, "air": -1},
    ## Outliers   --> Distance about equal if both exist!
    "e17":{"from": 4,"to": 5,"road": 2, "air": 2},
    "e18":{"from": 5,"to": 4,"road": 2, "air": 2},
    "e19":{"from": 6,"to": 7,"road": -1, "air": 2},
    "e20":{"from": 7,"to": 6,"road": -1, "air": 2},
    "e21":{"from": 3,"to": 11,"road": -1, "air": 3},
    "e22":{"from": 11,"to": 3,"road": -1, "air": 3},
    ## Upper left square  --> Hier besser mit den Drohnen agieren, Falls Sie da sind!!
    "e23":{"from": 0,"to": 8,"road": -1, "air": 4},
    "e24":{"from": 8,"to": 0,"road": -1, "air": 4},
    "e25":{"from": 8,"to": 9,"road": 3, "air": 3},
    "e26":{"from": 9,"to": 8,"road": 3, "air": 3},
    "e27":{"from": 9,"to": 10,"road": 3, "air": 3},
    "e28":{"from": 10,"to": 9,"road": 3, "air": 3},
    "e29":{"from": 0,"to": 10,"road": 3, "air": 3},
    "e30":{"from": 10,"to": 0,"road": 3, "air": 3}
    }
}_____no_output_____class Map_Environment(MultiAgentEnv):

    def __init__(self, env_config: dict = {}):
        # ensure config file includes all necessary settings
        assert 'NUMBER_STEPS_PER_EPISODE' in env_config
        assert 'NUMBER_OF_DRONES' in env_config
        assert 'NUMBER_OF_CARS' in env_config
        assert 'INIT_NUMBER_OF_PARCELS' in env_config
        assert 'TOPOLOGY' in env_config
        assert 'MAX_NUMBER_OF_PARCELS' in env_config
        assert 'THRESHOLD_ADD_NEW_PARCEL' in env_config
        assert 'BASELINE_FLAG' in env_config
        assert 'BASELINE_TIME_CONSTRAINT' in env_config
        assert 'BASELINE_OPT_CONSTANT' in env_config
        assert 'CHARGING_STATION_NODES' in env_config
        assert 'REWARDS' in env_config

        self.graph = load_graph(topology)

        # Map config
        self.NUMBER_OF_DRONES = env_config['NUMBER_OF_DRONES']
        self.NUMBER_OF_CARS = env_config['NUMBER_OF_CARS']
        self.NUMBER_OF_EDGES = self.graph.number_of_edges()
        self.NUMBER_OF_NODES = self.graph.number_of_nodes()
        self.CHARGING_STATION_NODES = env_config['CHARGING_STATION_NODES']
        self.MAX_BATTERY_POWER = env_config['MAX_BATTERY_POWER']

        # Simulation config
        self.NUMBER_STEPS_PER_EPISODE = env_config['NUMBER_STEPS_PER_EPISODE']
        self.INIT_NUMBER_OF_PARCELS = env_config['INIT_NUMBER_OF_PARCELS']
        self.RANDOM_SEED = env_config.get('RANDOM_SEED', None)
        self.MAX_NUMBER_OF_PARCELS = env_config['MAX_NUMBER_OF_PARCELS']
        self.THRESHOLD_ADD_NEW_PARCEL = env_config['THRESHOLD_ADD_NEW_PARCEL']
        self.BASELINE_FLAG = env_config['BASELINE_FLAG']
        self.BASELINE_TIME_CONSTRAINT = env_config['BASELINE_TIME_CONSTRAINT']
        self.BASELINE_OPT_CONSTANT = env_config['BASELINE_OPT_CONSTANT']

        self.DEBUG_LOG = env_config.get('DEBUG_LOGS', False)
        # Some Sanity Checks on the settings
        if self.DEBUG_LOG: assert self.MAX_NUMBER_OF_PARCELS >= self.INIT_NUMBER_OF_PARCELS, "Number of initial parcels exceeds max parcel limit"

        #Reward constants
        self.STEP_PENALTY = env_config['REWARDS']['STEP_PENALTY']
        self.PARCEL_DELIVERED = env_config['REWARDS']['PARCEL_DELIVERED']
        # compute other rewards
        self.BATTERY_DIED = self.STEP_PENALTY * self.NUMBER_STEPS_PER_EPISODE
        self.BATTERY_DIED_WITH_PARCEL = self.BATTERY_DIED * 2
        #self.DELIVERY_CONTRIBUTION: depends on active agents --> computed when given in prepare_global_reward()
        #self.ALL_DELIVERED_CONTRIB: depends on active agents --> computed when given in prepare_global_reward(episode_success=True)

        # Computed constants
        self.NUMBER_OF_AGENTS = self.NUMBER_OF_DRONES + self.NUMBER_OF_CARS
        self.PARCEL_STATE_DELIVERED = self.NUMBER_OF_AGENTS + self.NUMBER_OF_NODES
        self.NUMBER_OF_ACTIONS = 1 + self.NUMBER_OF_NODES + 1 + self.MAX_NUMBER_OF_PARCELS
        self.ACTION_DROPOFF = 1 + self.NUMBER_OF_NODES     # First Action NOOP is 0

        # seed RNGs
        self.seed(self.RANDOM_SEED)

        self.state = None
        self.current_step = None
        self.blocked_agents = None
        self.parcels_delivered = None
        self.done_agents = None
        self.all_done = None
        self.allowed_actions = None

        # baseline related
        self.baseline_missions = None
        self.agents_base = None
        self.o_employed = None

        # metrics for the evaluation
        self.parcel_delivered_steps = None # --> {p1: 20, p2: 240, p:140}
        self.parcel_added_steps = None     # --> {p1: 0, p2: 0, p: 50}
        self.agents_crashed = None         # --> {c_2: 120, d_0: 242}
        self.metrics = None

        self.agents = [*["d_" + str(i) for i in range(self.NUMBER_OF_DRONES)],*["c_" + str(i) for i in range(self.NUMBER_OF_DRONES, self.NUMBER_OF_DRONES + self.NUMBER_OF_CARS)]]

        # Define observation and action spaces for individual agents
        self.action_space = Discrete(self.NUMBER_OF_ACTIONS)

        #---- Repeated Obs Space: Represents a parcel with (id, location, destination)  --> # parcel_id starts at 1
        parcel_space = Dict({'id': Discrete(self.MAX_NUMBER_OF_PARCELS+1),
                             'location': Discrete(self.NUMBER_OF_NODES + self.NUMBER_OF_AGENTS + 1),
                             'destination': Discrete(self.NUMBER_OF_NODES)
                            })
        self.observation_space = Dict({'obs': Dict({'state':
                                                 Dict({ 'position': Discrete(self.NUMBER_OF_NODES),
                                                        'battery': Discrete(self.MAX_BATTERY_POWER + 1), #[0-100]
                                                        'has_parcel': Discrete(self.MAX_NUMBER_OF_PARCELS + 1),
                                                        'current_step': Discrete(self.NUMBER_STEPS_PER_EPISODE + 1)}),
                                                 'parcels': Repeated(parcel_space, max_len=self.MAX_NUMBER_OF_PARCELS)
                                                }),
                                        'allowed_actions': MultiBinary(self.NUMBER_OF_ACTIONS)
                                      })
        #TODO: why is reset() not called by env?
        self.reset()

    def step(self, action_dict):
        """conduct the state transitions caused by actions in action_dict
        :returns:
            - observation_dict: observations for agents that need to act in the next round
            - rewards_dict: rewards for agents following their chosen actions
            - done_dict: indicates end of episode if max_steps reached or all parcels delivered
            - info_dict: pass data to custom logger
        """

        if self.DEBUG_LOG: print(f"Debug log flag set to {self.DEBUG_LOG}")

        # ensure no disadvantage for agents with higher IDs if action conflicts with that taken by other agent
        action_dict = shuffle_actions(action_dict)

        self.current_step += 1

        # grant step penalty reward
        agent_rewards = {agent: self.STEP_PENALTY for agent in self.agents}

        # setting an agent done twice might cause crash when used with tune... -> https://github.com/ray-project/ray/issues/10761
        dones = {}

        self.metrics['step'] = self.current_step

        # dynamically add parcel
        if random.random() <= self.THRESHOLD_ADD_NEW_PARCEL and len(self.state['parcels']) < self.MAX_NUMBER_OF_PARCELS:

            p_id, parcel = self.generate_parcel()
            assert p_id not in self.state['parcels'], "Duplicate parcel ID generated"
            self.state['parcels'][p_id] = parcel

            if self.BASELINE_FLAG:
                self.compute_central_delivery(p_id, debug_log=False)
                self.metrics["optimal"] = self.compute_optimality_baseline(p_id, extra_charge=self.BASELINE_OPT_CONSTANT)

            for agent in self.agents:
                self.allowed_actions[agent][self.ACTION_DROPOFF + p_id] = np.array([1]).astype(bool)

        if self.BASELINE_FLAG:
            old_actions = action_dict
            action_dict = {}
            # Replace actions with actions recommended by the central baseline
            for agent, action in old_actions.items():
                if len(self.baseline_missions[agent]) > 0:
                    new_action = self.baseline_missions[agent][0]

                    if type(new_action) is tuple: # dropoff action with minimal time

                        if self.current_step >= new_action[1]:
                            new_action = new_action[0]
                            self.baseline_missions[agent].pop(0)
                        else: #agent has to wait for previous subroute agent
                            new_action = 0
                    else: # move or pickup or charge
                        self.baseline_missions[agent].pop(0)

                    action_dict[agent] = new_action

                else: # agent has no baseline mission -> Noop
                    action_dict[agent] = 0

        # carry out State Transition

        # handel NOP actions: -> action == 0
        noop_agents =  {agent: action for agent, action in action_dict.items() if action == 0}
        effectual_agents_items = {agent: action for agent, action in action_dict.items() if action > 0}.items()

        # now: transaction between agents modelled as pickup of just offloaded (=dropped) parcel --> handle dropoff first
        moving_agents = {agent: action for agent, action in effectual_agents_items if 0 < action and action <= self.NUMBER_OF_NODES}
        dropoff_agents = {agent: action for agent, action in effectual_agents_items if action == self.ACTION_DROPOFF}
        pickup_agents = {agent: action for agent, action in effectual_agents_items if action > self.ACTION_DROPOFF}

        # handle noop / charge decisions:
        for agent, action in noop_agents.items():
            # check if recharge is possible
            current_pos = self.state['agents'][agent]['position']
            if current_pos in self.CHARGING_STATION_NODES:
                self.state['agents'][agent]['battery'] = self.MAX_BATTERY_POWER

        # handle Movement actions:
        for agent, action in moving_agents.items():
            # get Current agent position from state
            self.state['agents'][agent]['battery'] += -1
            current_pos = self.state['agents'][agent]['position']

            # networkX: use node instead of edge:
            destination = action - 1

            if self.graph.has_edge(current_pos, destination):
                # Agent chose existing edge! -> check if type is suitable

                agent_type = 'road' if agent[:1] == 'c' else 'air'
                if self.graph[current_pos][destination]["type"] in [agent_type, 'both']:
                    # Edge has correct type
                    self.state['agents'][agent]['position'] = destination
                    self.state['agents'][agent]['battery'] += -(self.graph[current_pos][destination][agent_type] +1)
                    if self.state['agents'][agent]['battery'] < 0:        # ensure negative battery value does not break obs_space
                        #Battery below 0 --> reset to 0 (stay in obs space)
                        self.state['agents'][agent]['battery'] = 0

                    self.blocked_agents[agent] = self.graph[current_pos][destination][agent_type]
                    self.update_allowed_actions_nodes(agent)

        # handle Dropoff Decision: -> action == self.NUMBER_OF_NODES + 2
        for agent, action in dropoff_agents.items():
            self.state['agents'][agent]['battery'] += -1

            if self.state['agents'][agent]['has_parcel'] > 0:   # agent has parcel
                parcel_id = self.state['agents'][agent]['has_parcel']
                self.state['agents'][agent]['has_parcel'] = 0
                self.state['parcels'][parcel_id][0] = self.state['agents'][agent]['position']
                if self.state['parcels'][parcel_id][0] ==  self.state['parcels'][parcel_id][1]:
                    # Delivery successful

                    agent_rewards[agent] += self.PARCEL_DELIVERED  # local reward
                    # global contribution rewards
                    active_agents, reward = self.prepare_global_reward()
                    for a_id in active_agents: agent_rewards[a_id] += reward

                    self.state['parcels'][parcel_id][0] = self.PARCEL_STATE_DELIVERED
                    self.parcels_delivered[int(parcel_id) -1] = True                    # Parcel_ids start at 1

                    self.metrics['delivered'].update({"p_" + str(parcel_id): self.current_step})

                self.update_allowed_actions_parcels(agent)

        # handle Pickup Decision:
        for agent, action in pickup_agents.items():
            self.state['agents'][agent]['battery'] += -1

            # agent has free capacity
            if self.state['agents'][agent]['has_parcel'] == 0:  #free parcel capacity
                # convert action_id to parcel_id
                parcel_id = action - self.ACTION_DROPOFF

                if self.DEBUG_LOG: assert parcel_id in self.state['parcels'] # parcel {parcel_id} already in ENV ??

                elif self.state['parcels'][parcel_id][0] == self.state['agents'][agent]['position']:
                    # Successful pickup operation
                    self.state['parcels'][parcel_id][0] = self.NUMBER_OF_NODES + int(agent[2:])
                    self.state['agents'][agent]['has_parcel'] = int(parcel_id)

                    self.update_allowed_actions_parcels(agent)

        # unblock agents for next round
        self.blocked_agents = {agent: remaining_steps -1 for agent, remaining_steps in self.blocked_agents.items() if remaining_steps > 1}

        # handle dones - out of battery or max_steps or goal reached
        for agent in action_dict.keys():
            if agent not in self.done_agents and self.state['agents'][agent]['battery'] <= 0:
                agent_rewards[agent] = self.BATTERY_DIED_WITH_PARCEL if self.state['agents'][agent]['has_parcel'] != 0 else self.BATTERY_DIED
                dones[agent] = True
                self.done_agents.append(agent)

                self.metrics['crashed'].update({agent: self.current_step})

                if len(self.done_agents) == self.NUMBER_OF_AGENTS:
                    # all agents dead
                    self.all_done = True

        # check if episode terminated because of goal reached or all agents crashed -> avoid setting done twice
        if self.current_step >= self.NUMBER_STEPS_PER_EPISODE or (all(self.parcels_delivered) and len(self.parcels_delivered) == self.MAX_NUMBER_OF_PARCELS):
            # check if episode success:
            if self.current_step < self.NUMBER_STEPS_PER_EPISODE:
                    # grant global reward for all parcels delivered
                    active_agents, reward = self.prepare_global_reward(_episode_success=True)
                    for a_id in active_agents: agent_rewards[a_id] += reward

            self.all_done = True

        dones['__all__'] = self.all_done

        parcel_obs = self.get_parcel_obs()

        # obs \ rewards \ dones\ info
        return  {agent: { 'obs': {'state': {'position': self.state['agents'][agent]['position'], 'battery': self.state['agents'][agent]['battery'],
                                    'has_parcel': self.state['agents'][agent]['has_parcel'],'current_step': self.current_step},
                                    'parcels': parcel_obs},
                         'allowed_actions': self.allowed_actions[agent]} for agent in self.agents if agent not in self.blocked_agents and agent not in self.done_agents}, \
                { agent: agent_rewards[agent] for agent in self.agents}, \
                dones, \
                {}

    def seed(self, seed=None):
        tf.random.set_seed(seed)
        np.random.seed(seed)
        random.seed(seed)

    def reset(self):
        """resets variables; returns dict with observations, keys are agent_ids"""
        self.current_step = 0
        self.blocked_agents = {}
        self.parcels_delivered = [False for _ in range(self.MAX_NUMBER_OF_PARCELS)]
        self.done_agents = []
        self.all_done = False

        self.metrics = { "step": self.current_step,
            "delivered": {},
            "crashed": {},
            "added": {},
            "optimal": None
        }

        #Baseline
        self.baseline_missions = {agent: [] for agent in self.agents}
        self.o_employed = [0 for _ in range(self.NUMBER_OF_AGENTS)]

        self.agents_base = None  # env.agents in the pseudocode
        self.allowed_actions = { agent: np.array([1 for act in range(self.NUMBER_OF_ACTIONS)]) for agent in self.agents}

        #Reset State
        self.state = {'agents': {},
                      'parcels': {}}
        #generate initial parcels
        for _ in range(self.INIT_NUMBER_OF_PARCELS):
            p_id, parcel = self.generate_parcel()
            self.state['parcels'][p_id] = parcel

        parcel_obs = self.get_parcel_obs()

        # init agents
        self.state['agents'] = {agent: {'position': self._random_feasible_agent_position(agent),
                                        'battery': self.MAX_BATTERY_POWER, 'has_parcel': 0} for agent in self.agents}

        if self.BASELINE_FLAG:
            for parcel in self.state['parcels']:
                self.compute_central_delivery(parcel, debug_log=False)

                # TODO really return something here??
                self.metrics['optimal'] = self.compute_optimality_baseline(parcel, extra_charge=self.BASELINE_OPT_CONSTANT, debug_log=False)

        # compute allowed actions per agent --> move to function
        for agent in self.agents:
            self.update_allowed_actions_nodes(agent)
            self.update_allowed_actions_parcels(agent)

        agent_obs = {agent: {'obs': {'state':  {'position': state['position'], 'battery': state['battery'],
                                        'has_parcel': state['has_parcel'],'current_step': self.current_step},
                            'parcels': parcel_obs
                            },
                            'allowed_actions': self.allowed_actions[agent]
                            } for agent, state in self.state['agents'].items()
                    }

        return {**agent_obs}

    def _random_feasible_agent_position(self, agent_id):
        """Needed to avoid car agents being initialized at nodes only reachable by drones
        and thus trapped from the beginning. Ensures that car agents start at a node of type 'road' or 'parking'.
        """
        position = random.randrange(self.NUMBER_OF_NODES)
        if agent_id[0] == 'c':    # agent is car
            while self.graph.nodes[position]['type'] == 'air':      # position not reachable by car
                position = random.randrange(self.NUMBER_OF_NODES)
        return position

    def update_allowed_actions_nodes(self, agent):
        new_pos = self.state['agents'][agent]['position']
        next_steps = list(self.graph.neighbors(new_pos))
        agent_type = 'air' if agent[0]=='d' else 'road'

        allowed_nodes = np.zeros(self.NUMBER_OF_NODES)
        for neighbor in next_steps:
            if self.graph[new_pos][neighbor]['type'] in [agent_type, 'both']:
                allowed_nodes[neighbor] = 1
        self.allowed_actions[agent][1:self.NUMBER_OF_NODES+1] = np.array(allowed_nodes).astype(bool)

    def update_allowed_actions_parcels(self, agent):
        """ Allow only the Dropoff or Pickup actions, depending on the has_parcel value of the agent.
        Pickup is not concerned with the parcel actually being at the agents current location, only with free capacity
        and the parcel already being added to the ENV"""
        num_parcels = len(self.state['parcels'])

        allowed_parcels = np.zeros(self.MAX_NUMBER_OF_PARCELS)
        dropoff = 1
        if self.state['agents'][agent]['has_parcel'] == 0:
            dropoff = 0
            allowed_parcels = np.concatenate([np.ones(num_parcels), np.zeros(self.MAX_NUMBER_OF_PARCELS - num_parcels)])

        self.allowed_actions[agent][self.NUMBER_OF_NODES+1:] = np.array([dropoff, *allowed_parcels]).astype(bool)

    def get_parcel_obs(self):
        parcel_obs = [{'id':pid, 'location': parcel[0], 'destination':parcel[1]} for (pid, parcel) in self.state['parcels'].items()]
        return parcel_obs

    def generate_parcel(self):
        """generate new parcel id and new parcel with random nodes for location and destination.
            p_ids (int) start at 1, later nodes for parcel will be sampled to avoid parcels that already spawn at their destination"""
        p_id = len(self.state['parcels']) + 1

        parcel = random.sample(range(self.NUMBER_OF_NODES), 2)  # => initial location != destination
        self.metrics['added'].update({p_id: self.current_step})

        return p_id, parcel

    def prepare_global_reward(self, _episode_success=False):
        """ computes a global reward for all active agents still in the environment.
            If _episode_success is set to True, all parcels have been delivered an the ALL_DELIVERED reward is granted.
            :Returns: list of active agents and the reward value
        """
        agents_alive = list(set(self.agents).difference(set(self.done_agents)))
        if self.DEBUG_LOG: assert len(agents_alive) > 0
        reward = self.PARCEL_DELIVERED * (self.NUMBER_STEPS_PER_EPISODE - self.current_step) / self.NUMBER_STEPS_PER_EPISODE if _episode_success else self.PARCEL_DELIVERED / len(agents_alive)

        return agents_alive, reward

    #------ BASELINE related methods ------###
    def compute_optimality_baseline(self, parcel_id, extra_charge=2.5, debug_log=False):
        """Used in the optimality baseline
            Input: parcel_id, (extra_charge)
            Output: new total delivery rounds needed for all parcels
        """
        parcel = self.state['parcels'][parcel_id]
        path_time = 2 + shortest_path_length(self.graph, parcel[0], parcel[1], 'weight')

        _time = math.ceil(path_time * extra_charge) # round to next higher integer

        min_index = self.o_employed.index(min(self.o_employed))
        self.o_employed[min_index] += _time

        return max(self.o_employed)

    def compute_central_delivery(self, p_id, debug_log = False):
        """Used in the central baseline, iteratively tries to find a good delivery route
        with the available agents in the time specified in BASELINE_TIME_CONSTRAINT
            Input: parcel_id
            Output: Dict: {agent_id: [actions], ...}  --> update that dict! (merge in this function with prev actions!)
        """

        if self.agents_base is None:
            self.agents_base = {a_id: (a['position'], 0) for (a_id, a) in self.state["agents"].items()}  # last instructed pos + its step count

        min_time = None
        new_missions = {}   # key: agent, value: [actions]

        source = self.state["parcels"][p_id][0]
        target = self.state["parcels"][p_id][1]
        shortest_paths_generator = nx.shortest_simple_paths(self.graph, source, target, weight='weight')

        running = BoolTimer(self.BASELINE_TIME_CONSTRAINT)
        running.start()

        while running:
            # Default --> assign full route to nearest drone
            if min_time is None:

                air_route = nx.shortest_path(self.graph, source= source, target= target, weight="air")
                air_route_time = nx.shortest_path_length(self.graph, source= source, target= target, weight="air")
                air_route.pop(0)  # remove source node from path

                # Assign most suitable Drone
                best_drone = None
                for (a_id, a_tp) in self.agents_base.items():
                    if a_id[0] != 'd':                      # filter for correct agent type
                        continue

                    journey_time = nx.shortest_path_length(self.graph, source=a_tp[0], target=source, weight="air") + a_tp[1]
                    if min_time is None or journey_time < min_time:
                        min_time = journey_time
                        best_drone = (a_id, a_tp)

                # construct path for agent
                drone_route = nx.shortest_path(self.graph, source= best_drone[1][0], target= source, weight="air")

                drone_route_actions = [x+1 for x in drone_route[1:]] + [(self.ACTION_DROPOFF + p_id, 0)] + [x+1 for x in air_route[1:]] + [self.ACTION_DROPOFF]   # increment node_ids by one to get corresponding action
                min_time += air_route_time + 2 # add 2 steps for pick & drop

                self._add_charging_stops_to_route(drone_route_actions, debug_log=debug_log)

                new_missions[best_drone[0]] = (drone_route_actions, min_time)

            else:   # try to improve the existing base mission
                try:
                    shortest_route = next(shortest_paths_generator)
                except StopIteration:
                    break               # all existing shortest paths already tried

                subroutes = self._path_to_subroutes(shortest_route, debug_log=debug_log)
                duration, min_agents = self._find_best_agents(subroutes, min_time, debug_log=debug_log)

                if duration < min_time:
                    # faster delivery route found!
                    assert duration < min_time, "Central Baseline prefered a longer route..."

                    # update min_time
                    min_time = duration
                    new_missions = self._build_missions(min_agents, subroutes, p_id, debug_log=debug_log)

        #---- end while = Timer
        # now save the best mission found in the ENV
        for agent in new_missions.keys():
            # retrieve target node from mission, depends on charging stops and case no move necessary
            target = None
            if isinstance(new_missions[agent][0][-2], int):
                target = new_missions[agent][0][-2]
                if target == 0: target = new_missions[agent][0][-3]    # charging action was added before dropoff
                target -= 1                                            # action-1 = node_id
            else:  # handle case no move necessary -> pick action before dropoff
                target = self.agents_base[agent][0]

            self.agents_base[agent] = (target, self.agents_base[agent][1] + new_missions[agent][1])
            self.baseline_missions[agent].extend(new_missions[agent][0])

        return new_missions

    def _find_best_agents(self, subroutes, min_time, debug_log=False):
        """ For use in centrality_baseline. Finds best available agents for traversing a set of subroutes
            and returns these with the total duration for doing so.
            Input: subroutes = [ edge_type, [actions]]
            """

        min_agents = {}
        temp_agents_base = {k:v for k,v in self.agents_base.items()}  # deep copy for temporary planning transfer or stick with agent??

        for i,r in enumerate(subroutes):

            # init some helper vars
            a_type = "d" if r[0] == 'air' else 'c'
            min_time_sub = None    # best time (min) for this subroute (closest agent)
            best_agent_sub = None

            # iterate over agents of correct type!  --> later: Busy / unbusy
            for (a_id, a_tp) in temp_agents_base.items():
                #reminder: a_tp is tuple of latest future position (node, timestep)

                # filter for correct agent type - even if type is 'both' one agent can still take only its edge type
                weight_agent = r[0]
                if r[0] == 'both':
                    weight_agent = 'road' if a_id[0] == 'c' else 'air'
                else:                     # wrong agent type
                    if a_id[0] != a_type:  # todo replace with parameter variable in function!
                        continue

                journey_time = nx.shortest_path_length(self.graph, source=a_tp[0], target=r[1][0], weight=weight_agent) + a_tp[1] # earliest time agent can be there
                if min_time_sub is None or journey_time < min_time_sub:
                    min_time_sub = journey_time
                    best_agent_sub = (a_id, a_tp)    # a_id, a_tp = (latest location, timestep)

            # closest available agent found
            best_agent_weight = weight_agent = 'road' if best_agent_sub[0] == 'c' else 'air'
            duration_sub = min_time_sub + nx.shortest_path_length(self.graph, source=r[1][0], target=r[1][-1], weight=best_agent_weight) + 1

            # update agent state in temporary planning
            temp_agents_base[best_agent_sub[0]] = (r[1][-1], duration_sub)

            # agent_tuple, duration_subroutes_until_then
            min_agents[i] = (best_agent_sub, duration_sub)

            if debug_log: assert duration_sub < sys.maxsize, "Non existent edge taken somewhere..."

            # check if current subroute already longer than the min one
            if duration_sub > min_time:
                break # already worse, check next simple path

        return duration_sub, min_agents

    def _build_missions(self, min_agents, subroutes, parcel_id, debug_log = False):
        """For use in centrality_baseline. Computes list of actions for delivery of parcel parcel_id
        and necessary duration for execution from list of agents, subroutes"""
        new_mission = {}

        for i, s in enumerate(subroutes):
            best_agent_pos = min_agents[i][0][1][0]
            #earliest time to start that action
            time_pickup = min_agents[i-1][1] if i > 0 else 0 # First subroute pickup as soon as possible

            # construct the actual delivery path to pickup
            delivery_route = nx.shortest_path(self.graph, source=best_agent_pos, target=s[1][0], weight=s[0])
            delivery_route_actions = [x+1 for x in delivery_route[1:]] + [(self.ACTION_DROPOFF + parcel_id, time_pickup)] + [x+1 for x in s[1][1:]] + [self.ACTION_DROPOFF]

            if debug_log: assert min_agents[i][0] not in new_mission  #I see no preferable case where agent picks-up 1 parcel 2 times --> holding always better than following

            self._add_charging_stops_to_route(delivery_route_actions, debug_log=debug_log)
            new_mission[min_agents[i][0][0]] = (delivery_route_actions, min_agents[i][1])

        return new_mission

    def _add_charging_stops_to_route(self, route_actions, debug_log=False):
        """For use in centrality_baseline. Iterates through a list of actions and inserts a charging action
        after every move action to a node with a charging station.
        Tuples representing Dropoff actions with minimal executions time are updated to account for eventual delays. """

        delay = 0
        for i, n in enumerate(route_actions):
            if type(n) is tuple:
                if delay > 0: route_actions[i] = (n[0], n[1] + delay)
            else:
                if n-1 in self.CHARGING_STATION_NODES:
                    delay += 1
                    route_actions.insert(i+1, 0)

    def _path_to_subroutes(self, path, debug_log= False):
        """For use in centrality_baseline. Takes path in the graph as input and returns list of subroutes
        split at changes of the edge types with the type"""

        # get subroutes by their edge_types
        e_type_prev = None
        e_type_changes = [] # save indices of source nodes before new edge type
        subroutes = []
        _subroute = [path[0]]

        if len(path) > 1:
            e_type_prev = self.graph.edges[path[0], path[1]]['type']

            for i, node in enumerate(path[1:-1], start=1):

                e_type_next = self.graph.edges[node, path[i+1]]['type']

                _subroute.append(node)
                if e_type_next != e_type_prev:
                    subroutes.append((e_type_prev, _subroute))
                    _subroute = [node]
                    e_type_prev = e_type_next

            _subroute.append(path[-1])
        subroutes.append((e_type_prev, _subroute))   # don't forget last subroute

        return subroutes
_____no_output_____

## Agent Model and Experiment Evaluation Code_____no_output_____

from gym.spaces import Discrete, Box, Tuple, MultiDiscrete, Dict, MultiBinary
#from ray.rllib.utils.spaces.space_utils import flatten_space
#from ray.rllib.models.preprocessors import DictFlatteningPreprocessor

# Parametric-action agent model  --> apply Action Masking!
class ParametricAgentModel(TFModelV2):
    def __init__(self, obs_space, action_space, num_outputs, model_config, name, *args, **kwargs):
        super(ParametricAgentModel, self).__init__(obs_space, action_space, num_outputs, model_config, name, *args, **kwargs)

        assert isinstance(action_space, Discrete), f'action_space is a {type(action_space)}, but should be Discrete!'

        # Adjust for number of agents/parcels/Nodes!! -> Simply copy found shape from the thrown exception
        true_obs_shape = (1750, )
        action_embed_size = action_space.n

        self.action_embed_model = FullyConnectedNetwork(
            Box(0, 1, shape=true_obs_shape),  # TODO hier nochmal die 1 anpassen?? --> muss das hier der obs entsprechen??
            action_space,
            action_embed_size,
            model_config,
            name + '_action_embedding')

    def forward(self, input_dict, state, seq_lens):

        action_mask = input_dict['obs']['allowed_actions']
        action_embedding, _ = self.action_embed_model.forward({'obs_flat': input_dict["obs_flat"]}, state, seq_lens)
        intent_vector = tf.expand_dims(action_embedding, 1)
        action_logits = tf.reduce_sum(intent_vector, axis=1)
        inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min)

        return action_logits + inf_mask, state

    def value_function(self):
        return self.action_embed_model.value_function()
    _____no_output_____## Proposed way to train / evaluate MARL policy from github Issues: --> https://github.com/ray-project/ray/issues/9123 and  https://github.com/ray-project/ray/issues/9208

def train(config, name, save_dir, stop_criteria, num_samples, verbosity=1):
    """
    Train an RLlib PPO agent using tune until any of the configured stopping criteria is met.
    :param stop_criteria: Dict with stopping criteria.
        See https://docs.ray.io/en/latest/tune/api_docs/execution.html#tune-run
    :return: Return the path to the saved agent (checkpoint) and tune's ExperimentAnalysis object
        See https://docs.ray.io/en/latest/tune/api_docs/analysis.html#experimentanalysis-tune-experimentanalysis
    """

    print("Start training")
    analysis = ray.tune.run(PPOTrainer, verbose=verbosity, config=config, local_dir=save_dir,
                            stop=stop_criteria, name=name, num_samples=num_samples,
                            checkpoint_at_end=True)

    # list of lists: one list per checkpoint; each checkpoint list contains 1st the path, 2nd the metric value
    checkpoints = analysis.get_trial_checkpoints_paths(trial=analysis.get_best_trial('episode_reward_mean', mode='max'),
                                                       metric='episode_reward_mean')
    # retriev the checkpoint path; we only have a single checkpoint, so take the first one
    checkpoint_path = checkpoints[0][0]
    print(f"Saved trained model in checkpoint {checkpoint_path} - achieved episode_reward_mean: {checkpoints[0][1]}")
    return checkpoint_path, analysis

def load(config, path):
    """
    Load a trained RLlib agent from the specified path. Call this before testing the trained agent.
    """
    agent = PPOTrainer(config=config) #, env=env_class)
    agent.restore(path)
    return agent

def test(env_class, env_config, policy_mapping_fcn, agent):
    """Test trained agent for a single episode. Return the retrieved env metrics for this episode and the episode reward"""
    # instantiate env class
    env = env_class(env_config)
    episode_reward = 0
    done = False
    obs = env.reset()

    while not done: # run until episode ends
        actions = {}
        for agent_id, agent_obs in obs.items():
            # Here: policy_id == agent_id - added this to avoid confusion for other policy mappings
            policy_id = policy_mapping_fcn(agent_id, episode=None, worker=None)
            actions[agent_id] = agent.compute_action(agent_obs, policy_id=policy_id)
        obs, reward, done, info = env.step(actions)
        done = done['__all__']

        # sum up reward for all agents
        episode_reward += sum(reward.values())

    # Retrieve custom metrics from ENV
    return env.metrics, episode_reward
_____no_output_____
def train_and_test_scenarios(config, seeds=None):
    """ Trains for a single scenario indicated by a seed """
    # TODO how to distinguish between the different algos ?
    print("Starte: run_function_trainer!")

    # prepare the config dicts
    NAME = config['NAME']
    SAVE_DIR = config['SAVE_DIR']
    ENVIRONMENT = config['ENV']

    # Simulations
    NUMBER_STEPS_PER_EPISODE = config['NUMBER_STEPS_PER_EPISODE']
    STOP_CRITERIA = config['STOP_CRITERIA']
    NUMBER_OF_SAMPLES = config['NUMBER_OF_SAMPLES']

    #MAP / ENV
    NUMBER_OF_DRONES = config['NUMBER_OF_DRONES']
    NUMBER_OF_CARS = config['NUMBER_OF_CARS']
    NUMBER_OF_AGENTS = NUMBER_OF_DRONES + NUMBER_OF_CARS
    MAX_NUMBER_OF_PARCELS = config['MAX_NUMBER_OF_PARCELS']

    # TESTING
    SEEDS = config['SEEDS']

    env_config = {
            'DEBUG_LOGS':False,
            'TOPOLOGY': config['TOPOLOGY'],
            # Simulation config
            'NUMBER_STEPS_PER_EPISODE': NUMBER_STEPS_PER_EPISODE,
            #'NUMBER_OF_TIMESTEPS': NUMBER_OF_TIMESTEPS,
            'RANDOM_SEED': None, # 42
            # Map
            'CHARGING_STATION_NODES': config['CHARGING_STATION_NODES'],
            # Entities
            'NUMBER_OF_DRONES': NUMBER_OF_DRONES,
            'NUMBER_OF_CARS': NUMBER_OF_CARS,
            'MAX_BATTERY_POWER': config['MAX_BATTERY_POWER'],  # TODO split this for drone and car??
            'INIT_NUMBER_OF_PARCELS': config['INIT_NUMBER_OF_PARCELS'],
            'MAX_NUMBER_OF_PARCELS': config['MAX_NUMBER_OF_PARCELS'],
            'THRESHOLD_ADD_NEW_PARCEL': config['THRESHOLD_ADD_NEW_PARCEL'],
            # Baseline settings
            'BASELINE_FLAG': False,  # is set True in the test function when needed
            'BASELINE_OPT_CONSTANT': config['BASELINE_OPT_CONSTANT'],
            'BASELINE_TIME_CONSTRAINT': config['BASELINE_TIME_CONSTRAINT'],
            # TODO
            #Rewards
            'REWARDS': config['REWARDS']
        }

    run_config = {
        'num_gpus': config['NUM_GPUS'],
        'num_workers': config['NUM_WORKERS'],
        'env': ENVIRONMENT,
        'env_config': env_config,
        'multiagent': {
            'policies': {
                # tuple values: policy, obs_space, action_space, config
                **{a: (None, None, None, { 'model': {'custom_model': ParametricAgentModel }, 'framework': 'tf'}) for a in ['d_'+ str(j) for j in range(NUMBER_OF_DRONES)] + ['c_'+ str(i) for i in range(NUMBER_OF_DRONES, NUMBER_OF_CARS + NUMBER_OF_DRONES)]}
            },
            'policy_mapping_fn': policy_mapping_fn,
            'policies_to_train': ['d_'+ str(i) for i in range(NUMBER_OF_DRONES)] + ['c_'+ str(i) for i in range(NUMBER_OF_DRONES, NUMBER_OF_CARS + NUMBER_OF_DRONES)]
        },
        #'log_level': "INFO",
        #"hiddens": [],     # For DQN
        #"dueling": False,  # For DQN
    }

    # Train and Evaluate the agents !
    checkpoint, analysis = train(run_config, NAME, SAVE_DIR, STOP_CRITERIA, NUMBER_OF_SAMPLES)
    print("Training finished - Checkpoint: ", checkpoint)

    env_class = ENVIRONMENT
    # Restore trained policies for evaluation
    agent = load(run_config, checkpoint)
    print("Agent loaded - Agent: ", agent)

    # Run the test cases for the specified seeds
    runs = {'max_steps': NUMBER_STEPS_PER_EPISODE, 'max_parcels': MAX_NUMBER_OF_PARCELS, 'max_agents': NUMBER_OF_AGENTS}

    for seed in SEEDS:
        print(seed)
        env_config['RANDOM_SEED'] = seed
        # TODO check if
        assert run_config['env_config']['RANDOM_SEED'] == seed
        result = test_scenario(run_config, agent)
        runs.update(result)

    return runs

def test_scenario(config, agent):
    """
    Loads a pretrained agent, initializes an environment from the seed
    and then evaluates it over one episode with the Marl agents and the central baseline.

    Returns: dict with results for graph creation for both evaluation / inference runs
    """
    #TODO unterscheidung run_config > env_config
    # Wie die beiden Runs abspeichern --> {Seed + [marl / base] : result }

    env_class = config['env']
    env_config = config['env_config']
    seed = env_config['RANDOM_SEED']
    policy_mapping_fn = config['multiagent']['policy_mapping_fn']

    # Test with MARL

    metrics_marl, reward_marl = test(env_class, env_config, policy_mapping_fn, agent)

    # Test with CentralBase
    env_config['BASELINE_FLAG'] = True
    metrics_central, reward_central = test(env_class, env_config, policy_mapping_fn, agent)
    env_config['BASELINE_FLAG'] = False

    # ASSERT that both optimal values are equal
    #assert metrics_marl['optimal'] == metrics_central['optimal']

    return {"M_" + str(seed): metrics_marl, "C_" + str(seed): metrics_central}

_____no_output_____from ray.rllib.agents.ppo import PPOTrainer

def policy_mapping_fn(agent_id, episode, worker, **kwargs):
    return agent_id

basic_config = {
        # experiment
        'NAME': 'small_no',
        'SAVE_DIR': 'Exp_environment',
        'ALGO': PPOTrainer,
        'ENV': Map_Environment,
        'DEBUG_LOGS':False,
        'NUM_GPUS': 2,
        'NUM_WORKERS': 32,
        'NUMBER_OF_SAMPLES': 2,
        # Simulation config
        'NUMBER_STEPS_PER_EPISODE': 1200,
        'RANDOM_SEED': None, # 42
        # Map
        'TOPOLOGY': topology,
        'CHARGING_STATION_NODES': [0,1,2,3,4,6,9],
        # Entities
        'NUMBER_OF_DRONES': 2,
        'NUMBER_OF_CARS': 2,
        'INIT_NUMBER_OF_PARCELS': 10,
        'MAX_NUMBER_OF_PARCELS': 10,
        'THRESHOLD_ADD_NEW_PARCEL': 0.1, # 10% chance
        'MAX_BATTERY_POWER': 100,
        #Baseline
        'BASELINE_TIME_CONSTRAINT': 10,
        'BASELINE_OPT_CONSTANT': 2.5,
        #TESTING
        'SEEDS': [72, 21, 44, 66, 86, 14],
        #Rewards
        'REWARDS': {
            'PARCEL_DELIVERED': 200,
            'STEP_PENALTY': -0.1,
        },
        'STOP_CRITERIA': {
            'timesteps_total': 24_000_000,
        }
}

print("Test evaluation functions")
experiment_results = train_and_test_scenarios(basic_config)
print("Test evaluation finished ;)")_____no_output_____ experiment_results_____no_output_____#  cached a result to test the graph generation
_cached_4_graph_____no_output_____
def create_chart_bars(results_dict):
    """ Function that plots a bar graph with the duration of one episode
        run with MARL agents/ Centrality Baseline/ Optimality Baseline, recorded in the :param results_dict:.
    """

    # Design choices
    colors = {
        "marl": 'blue',
        "central": 'green',
        "optimal": 'red'
    }

    # Retrieve settings values from results dict
    max_steps = results_dict['max_steps']
    max_parcels = results_dict['max_parcels']
    max_agents = results_dict['max_agents']

    # Deep copy for further computations
    scenario_results = {k:v for k,v in results_dict.items() if k[0:3] != 'max'}

    # filter
    runs = scenario_results.keys()
    values = scenario_results.values()

    # TODO statt all dead => einfach not all delivered marker
    merged = {}  # key is seed as str, value is a dict with [marl, central, optimal, all_dead_marl, all_dead_cent]

    # Retrieve the data
    for run_id, res in scenario_results.items():

        split_id = run_id.split('_')  # --> type, seed
        key_type, key_seed = split_id[0], split_id[1]
        # merge data from the runs with same seed (marl + baselines)

        # add new dict if seed not encountered yet
        if key_seed not in merged:
            merged[key_seed] = {}

        _key_delivered = 'marl'
        _key_crashed = 'all_dead_marl'
        if key_type == 'C':
            # Baseline run
            merged[key_seed]['optimal'] = res['optimal']
            _key_delivered = 'central'
            _key_crashed = 'all_dead_central'

        # Retrieve number of steps in run
        last_step = res['step']

        merged[key_seed][_key_delivered] = last_step
        all_parcels_delivered = len(res['delivered']) == max_parcels     # were all parcels delivered
        merged[key_seed][_key_delivered + '_all'] = all_parcels_delivered
        if not all_parcels_delivered:
            merged[key_seed][_key_delivered] = max_steps

    #print("Merged: ", merged)

    #  example data = [[30, 25, 50, 20],
#    [40, 23, 51, 17],
#    [35, 22, 45, 19]]

    data = [[],[],[],[], []]
    labels = []

    # split data into type of run
    for seed,values in merged.items():
        labels.append('S_'+seed)
        data[0].append(values['marl'])
        data[1].append(values['central'])
        data[2].append(values['optimal'])
        data[3].append(values['marl_all'])
        data[4].append(values['central_all'])

    print("Data: ", data)

    X = np.arange(len(labels))
    #print(X)

    fig = plt.figure()
    ax = fig.add_axes([0,0,1,1])

    ax.bar(X + 0.00, data[2], color = colors['optimal'], label="Optimality Baseline", width = 0.25)
    ax.bar(X + 0.25, data[0], color = colors['marl'], label="MARL System", width = 0.25, alpha=0.8)
    ax.bar(X + 0.50, data[1], color = colors['central'], label="Centrality Baseline", width = 0.25)

    plt.xlabel("Experiments")
    plt.ylabel("Steps taken")

    # Add experiment identifiers x-Axis
    plt.xticks(X + 0.25, labels)
    # Add legend
    ax.legend()
_____no_output_____# Plot duration bar graphs from the results
create_chart_bars(experiment_results)_____no_output_____import matplotlib.ticker as mtick

def create_chart_episode_events(results_dict, draw_crashes = False):
    """ Function that plots a graph with either the deliveries of parcels or the crashes of agents
        over the course of an episode, recorded in the :param results_dict:.
        Set the :param draw_crashes: flag for plotting crashes, default are deliveries.
    """
    # TODO Graphen auch abspeichern oder nur hier anzeigen ??
    # TODO - Parcel additions ??
    # TODO - overthink fill with max_value for mean computation

    # Design choices
    colors = {
        "marl": 'blue',
        "central": 'green',
        "optimal": 'red'
    }
    alpha = 0.6  # Opacity of individual value lines
    alpha_opt = 0.2
    opt_marker_size = 15
    line_width_mean = 5
    line_width_optimal = 2

    fig = plt.figure()
    ax = fig.add_subplot(111)

    # Retrieve settings values from results dict
    max_steps = results_dict['max_steps']
    max_parcels = results_dict['max_parcels']
    max_agents = results_dict['max_agents']

    _len = max_agents if draw_crashes else max_parcels
    _len += 1     # start plot at origin

    _key = 'crashed' if draw_crashes else 'delivered'

    # Deep copy for further computations
    scenario_results = {k:v for k,v in results_dict.items() if k[0:3] != 'max'}

    # for computation of mean
    m_values, c_values, o_values = [], [], []
    Y = [str(i) for i in range(0, _len)]

    # iterate over configs
    for scenario, results in scenario_results.items():

        # Default settings -> MARL run
        color = colors["marl"]
        _type = m_values
        label= "marl_system"

        if scenario[0] == 'C':
            # Baseline run
            color = colors["central"]
            _type = c_values
            label = "centrality_baseline"

            # Retrieve and plot optimality baseline
            optimal_time = results['optimal']
            assert optimal_time is not None
            if not draw_crashes: ax.plot(optimal_time, 0, "*", color = colors["optimal"], label="optimality_baseline", markersize=opt_marker_size, alpha= alpha_opt, clip_on=False)
            o_values.append(optimal_time)

        _num_steps = results['step']

        X = [0] + list(results[_key].values())

        X = X + [max_steps]*(_len - len(X))           # Fill X up with max_step values for not delivered parcels / not crashed agents

        _type.append(X)  # add X to the respective mean list
        #Y = [str(i) for i in range(0, len(results[_key].values())+1)]

        #print("Data: ", results[_key].values())
        #print("new X: ", X)
        #print("new Y: ", Y)

        ax.step(X, Y, label=label, where='post', color=color, alpha=alpha)

        # Attempt to improve the filling mess in the plot...
        #X = X + [max_steps]*(_len - len(X))           # Fill X up with max_step values for not delivered parcels / not crashed agents

        #_type.append(X)  # add X to the respective mean list

    # compute mean values
    m_mean = np.mean(np.array(m_values), axis=0)
    c_mean = np.mean(np.array(c_values), axis=0)
    o_mean = np.mean(np.array(o_values), axis=0)

    ax.step(m_mean, Y, label="marl_system", where='post', color=colors["marl"], linewidth=line_width_mean)
    ax.step(c_mean, Y, label="centrality_baseline", where='post', color=colors["central"], linewidth=line_width_mean)

    # star for opt:    if not draw_crashes: plt.plot(o_mean, 0, "*", label="optimality_baseline", color="r", markersize=opt_marker_size, clip_on=False)
    # better?: vertical line for opt
    if not draw_crashes: ax.axvline(o_mean, label="optimality_baseline", color=colors["optimal"], linewidth=2, alpha=alpha_opt+0.3)

    # y axis as percentage
    max_percent = max_agents if draw_crashes else max_parcels
    xticks = mtick.PercentFormatter(max_percent)
    ax.yaxis.set_major_formatter(xticks)

    # Lables and Legend
    plt.xlabel("steps")
    ylabel = "% of crashed agents" if draw_crashes else "% of delivered parcels"
    plt.ylabel(ylabel)
    handles, labels = plt.gca().get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    plt.legend(by_label.values(), by_label.keys())

    # Margins
    plt.ylim(bottom=0)
    plt.xlim()
    plt.margins(x=0, y=0)

    plt.show()
    _____no_output_____# Plot delivery graphs from the results
create_chart_episode_events(experiment_results, draw_crashes=False)
_____no_output_____# Plot crash graphs from the results
create_chart_episode_events(experiment_results, draw_crashes=True)_____no_output_____

### Manual Actions for debugging_____no_output_____

##################
env_config = {
            'DEBUG_LOGS':False,
            'TOPOLOGY': topology,
            # Simulation config
            'NUMBER_STEPS_PER_EPISODE': 1000,
            #'NUMBER_OF_TIMESTEPS': NUMBER_OF_TIMESTEPS,
            'RANDOM_SEED': None, # 42
            # Map
            'CHARGING_STATION_NODES': [0,1,2,3,4],
            # Entities
            'NUMBER_OF_DRONES': 2,
            'NUMBER_OF_CARS': 2,
            'MAX_BATTERY_POWER': 100,  # TODO split this for drone and car??
            'INIT_NUMBER_OF_PARCELS': 3,
            'MAX_NUMBER_OF_PARCELS': 3,
            'THRESHOLD_ADD_NEW_PARCEL': 0.01,
            # Baseline settings
            'BASELINE_FLAG': False,  # is set True in the test function when needed
            'BASELINE_OPT_CONSTANT': 2.5,
            'BASELINE_TIME_CONSTRAINT': 5,
            # TODO
            #Rewards
            'REWARDS': {
                'PARCEL_DELIVERED': 200,
                'STEP_PENALTY': -0.1,
            },
        }

env = Map_Environment(env_config)
env.state
#env.ACTION_DROPOFF_____no_output_____# TODO select actions to give agent 0 reward!!
#print(env.action_space)

actions_1 = {'d_0': 6, 'd_1': 2, 'c_2': 3, 'c_3': 4}
#actions_1 = {'d_0': 0, 'd_1': 0, 'c_2': 5, 'c_3': 5}
actions_2 = {'d_0': 2, 'd_1': 8, 'c_2': 7, 'c_3':0}
actions_3 = {'d_0': 5, 'd_1': 5, 'c_2':2 }

new_obs, rewards, dones, infos = env.step(actions_1)
print(infos)
print(f"New Obs are: {new_obs}")
print(rewards)
print("------------------")
new_obs2, rewards2, dones2, infos2 = env.step(actions_2)
print(f"New Obs are: {new_obs2}")
print(rewards2)
print("------------------")
new_obs3, rewards3, dones3, infos3 = env.step(actions_3)
print(f"New Obs are: {new_obs3}")
print(rewards3)

actions_4 = {'d_0': 0, 'd_1': 0, 'c_2': 1, 'c_3': 0}
actions_5 = {'d_0': 0, 'd_1': 0, 'c_2': 0, 'c_3': 0}
actions_6 = {'d_0': 0, 'd_1': 0, 'c_2': 5, 'c_3': 0}

new_obs, rewards, dones, infos = env.step(actions_4)
print(f"New Obs are: {new_obs}")
print(rewards)
print("------------------")
new_obs2, rewards2, dones2, infos2 = env.step(actions_5)
print(f"New Obs are: {new_obs2}")
print(rewards2)
print("------------------")
new_obs3, rewards3, dones3, infos3 = env.step(actions_6)
print(f"New Obs are: {new_obs3}")
print(rewards3)

print("------------------")
print(dones3)_____no_output_____# TENSORBOARD
# Load the TensorBoard notebook extension
%load_ext tensorboard
#Start tensorboard below the notebook
#%tensorboard --logdir logs_____no_output_____


# Notebook from sfarrens/The-Anatomy-of-a-Python-Class
Path: Classes_II.ipynb

_____no_output_____> Author: Samuel Farrens
> Email: samuel.farrens@cea.fr

# The Anatomy of a Python Class - Part II

This notebook introduces some more advanced concepts in the manipulation of Python classes. Before starting you should make sure you have completed the [first notebook](./Classes_I.ipynb) or are at least familiar with the topics covered therein. As with the first part, this tutorial is in no way exhaustive and you are encouraged to suppliment your understanding with further reading.

If you are new to Jupyter notebooks note that cells are executed by pressing SHIFT+ENTER (⇧+ ⏎). See the Jupyter documentation for more details.

## Contents

1. [Set-Up](#1.-Set-Up)
1. [Inheritance](#2.-Inheritance)
  1. [Inherting Attributes](#Inherting-Attributes)
  1. [Overriding](#Overriding)
  1. [Handling Instantiation](#Handling-Instantiation)
  1. [Multiple Parents](#Multiple-Parents)
  1. [Method Resolution Order](#Method-Resolution-Order)
1. [Composition](#3.-Composition)
1. [Abstract Classes](#4.-Abstract-Classes)
  1. [Abstract Methods](#Abstract-Methods)
  1. [Abstract Properties](#Abstract-Properties)_____no_output_____## 1. Set-Up

The following cell contains some set-up commands. Be sure to execute this cell before continuing._____no_output_____

# Notebook Set-Up Commands

def print_error(error):
    """ Print Error

    Function to print exceptions in red.

    Parameters
    ----------
    error : string
        Error message

    """
    print('\033[1;31m{}\033[1;m'.format(error))_____no_output_____

## 2. Inheritance

A powerful property of classes is the ability to inherit attributes and methods from other classes. This helps avoid repetition of code, which makes it easier to develop and maintain.

### Inherting Attributes

We will begin by defining a simple class with a single static method that will act as a *parent* class._____no_output_____

# Define a parent class
class Parent:

    @staticmethod
    def add(a, b):
        return a + b

# Print the parent dictionary
print('Parent.__dict__ =', Parent.__dict__)_____no_output_____

Note that this class has no special properties, it the same as classes we saw in the previous notebook.

Now we will define a second class, also with a single static function. This class will serve as a *child* class that will inherit from the parent class. This is done by simply passing the parent class name in `()` when defining the child class. _____no_output_____

# Define a child class
class Child(Parent):

    @staticmethod
    def subtract(a, b):
        return a - b

print('Child.__dict__ =', Child.__dict__)_____no_output_____

We can see that the contents of the child class dictionary are only what is expected from a normal class definition. Using the special `__bases__` attribute, however, we can see a tuple of classes that the child inherits from._____no_output_____

print('Child.__bases__ =', Child.__bases__)_____no_output_____

This also allows use to see the parent dictionary._____no_output_____

print(Child.__bases__[0].__dict__)_____no_output_____

Finally, we can demonstrate that the child has indeed inherited new attributes from the parent as follows._____no_output_____

print('1 + 2 =', Child.add(1, 2))
print('3 - 2 =', Child.subtract(3, 2))_____no_output_____

We can clearly see that `Child` has inhertited the static `add` method from `Parent`.

In the following example, we can see that this works for any type of class attribute._____no_output_____

class Parent:

    x = 1

class Child(Parent):

    @classmethod
    def show(cls):
        return print('x =', cls.x)

Child.show()_____no_output_____

### Overriding

We saw in the previous notebook that class instance attributes will override class attributes of the same name. The same happens with child attributes with respect to parent attributes. _____no_output_____

class Parent:

    x = 1

class Child(Parent):

    x = 2

    @classmethod
    def show(cls):
        return print('x =', cls.x)

Child.show()_____no_output_____

### Hierarchy

It is possible to define a hierarchy of parent classes, the attributes of which will be inheried by a given child class._____no_output_____

class GrandParent:

    x = 1
    name = 'Grandparent'

class Parent(GrandParent):

    y = 2
    name = 'Parent'

class Child(Parent):

    z = 3
    name = 'Child'

print(Child.name, Child.x, Child.y, Child.z)_____no_output_____

Note that overriding will also act hierarchically, meaning that the last class in the chain will be given precedence.
_____no_output_____### **Exercise**

How can we see which grand parent clases `Child` inherits from? 🤔_____no_output_____

############
# EXERCISE #
############
# Add your solution here _____no_output_____

### Handling Instantiation

If the parent class contains an `__init__` method, this too can be inherited by the child._____no_output_____

class Parent:

    def __init__(self, value):
        self.myattr = value

class Child(Parent):

    def show(self):
        return print('myattr =', self.myattr)

Child('A string').show()_____no_output_____

### **Exercise**

Try defining the `Parent` class with the `show` method and the `Child` class with the `__init__` method. Does it still work? 🤔_____no_output_____

############
# EXERCISE #
############
# Add your solution here _____no_output_____

This works because the child will lookup any attributes not listed in its own dictionary in all the parent (or base) classes.

If both classes have an `__init__` method, then only the child class will be initialised (see [Overriding](#Overriding))._____no_output_____

class Parent:

    def __init__(self):
        self.pval = 'parent value'

class Child(Parent):

    def __init__(self):
        self.cval = 'child value'

inst = Child()
print('cval =', inst.cval)_____no_output_____try:
    print('pval =', inst.pval)
except Exception as error:
    print_error(error)_____no_output_____

We can fix this by manually instantiating the parent class._____no_output_____

class Parent:

    def __init__(self):
        self.pval = 'parent value'

class Child(Parent):

    def __init__(self):
        self.cval = 'child value'
        Parent.__init__(self)

inst = Child()
print('cval =', inst.cval)
print('pval =', inst.pval)_____no_output_____

Python has a useful shortcut to make this process easier using the [`super`](https://docs.python.org/3/library/functions.html?highlight=super#super)function. This allows us to instantiate the parent without naming it, this particularly useful if the name of the parent class were to change._____no_output_____

class Parent:

    def __init__(self):
        self.pval = 'parent value'

class Child(Parent):

    def __init__(self):
        self.cval = 'child value'
        super().__init__()

inst = Child()
print('cval =', inst.cval)
print('pval =', inst.pval)_____no_output_____

### Multiple Parents

It is possible for a child to inherit attributes from multiple parents._____no_output_____

class Mother:

    def __init__(self, value):
        self.mval = value

    def show_mother(self):
        print('mother value =', self.mval)

class Father:

    def __init__(self, value):
        self.fval = value

    def show_father(self):
        print('father value =', self.fval)

class Child(Mother, Father):

    def __init__(self, value1, value2, value3):
        self.cval = value3
        Mother.__init__(self, value1)
        Father.__init__(self, value2)

    def show_child(self):
        print('child value =', self.cval)

inst = Child(1, 2, 3)
inst.show_mother()
inst.show_father()
inst.show_child()_____no_output_____

Note that the parent classes were explicitly instantiated in the previous example. Attempting the same thing using the `super` function will raise an error._____no_output_____

class Mother:

    def __init__(self, value):
        self.mval = value

    def show_mother(self):
        print('mother value =', self.mval)

class Father:

    def __init__(self, value):
        self.fval = value

    def show_father(self):
        print('father value =', self.fval)

class Child(Mother, Father):

    def __init__(self, value1, value2):
        self.cval = value2
        super().__init__(value1)

    def show_child(self):
        print('child value =', self.cval)

inst = Child(1, 2)_____no_output_____try:
    inst.show_mother()
    inst.show_father()
    inst.show_child()
except Exception as error:
    print_error(error)_____no_output_____

As we can see the `Mother` parent class was properly instantiated, but the `Father` parent class was not. The reasons for this are explained in the following subsection.

> In general, caution should be used when building a class archtecture that requires inheritance for multiple parent classes.

### Method Resolution Order

The Method Resolution Order (MRO) dictates the order in which the child class will search the base classes for a given attribute. We can see the MRO of the child class from the previous example using the `__mro__` attribute as follows._____no_output_____

print('Child.__mro__ =', Child.__mro__)_____no_output_____

We can see that the `Mother` class appears before the `Father` class in the MRO, hence the `super` method in `Child` instantiates `Mother`.

In order to make the previous example work, we would need to add another `super` in the `__init__` of the `Mother` class as follows.

> Note: Don't do this in your code!_____no_output_____

class Mother:

    def __init__(self, value):
        self.mval = value
        super().__init__(value)

    def show_mother(self):
        print('mother value =', self.mval)

class Father:

    def __init__(self, value):
        self.fval = value

    def show_father(self):
        print('father value =', self.fval)

class Child(Mother, Father):

    def __init__(self, value1, value2):
        self.cval = value2
        super().__init__(value1)

    def show_child(self):
        print('child value =', self.cval)

inst = Child(1, 2)
inst.show_mother()
inst.show_father()
inst.show_child()_____no_output_____

This, however, is not recommended as it adds unncessary ambiguity to the code and makes debugging extremely difficult._____no_output_____### **Exercise**

Create a parent class and a child class that can identify its progenitor.
1. Your parent class should have the class attribute `parent_name` with the value of your choice.
1. Your child class should have the instance attribute `name`, the value of which can be provided when creating an instance.
1. Printing an instance of your child class should contain its name and its parent's name. *e.g.*

```python
print(Child('Thor'))
Thor Odinson
```_____no_output_____

############
# EXERCISE #
############
# Add your solution here_____no_output_____

## 3. Composition

Inhertitance is not the only way in which a class can access attributes from other classes. This comes back one of main take-away messages of this tutorial, namely that everything in Python is an object and any object can be assigned as an attribute.

This means that a class attribute can actually be an instance of another class._____no_output_____

class Composer:

    def __init__(self):
        self.myattr = 'composer value'

class myClass:

    def __init__(self):
        self.comp = Composer()

inst = myClass()
print("myattr =", inst.comp.myattr)_____no_output_____

### **Exercise**

How would you go about viewing the contents of the `Composer` dictionary from an instance of `myClass`? 🤔_____no_output_____

############
# EXERCISE #
############
# Add your solution here_____no_output_____

This can particularly useful as we can pass classes or class intances to our new class without knowing anything about the other classes. _____no_output_____

class Star:

    def __init__(self):
        self.whoami = 'I am a star!'

class Galaxy:

    def __init__(self):
        self.whoami = 'I am a Galaxy!'

class myClass:

    def __init__(self, composer):
        self.comp = composer()

inst1 = myClass(Star)
inst2 = myClass(Galaxy)
print("whoami =", inst1.comp.whoami)
print("whoami =", inst2.comp.whoami)_____no_output_____

We can even achieve inheritance-like properties if we know the name of a given attribute that the composer class should have._____no_output_____

class Star:

    def __init__(self):
        self.whoami = 'I am a star!'

class Galaxy:

    def __init__(self):
        self.whoami = 'I am a Galaxy!'

class myClass:

    def __init__(self, composer):
        self.whoami = composer().whoami

inst1 = myClass(Star)
inst2 = myClass(Galaxy)
print("whoami =", inst1.whoami)
print("whoami =", inst2.whoami)_____no_output_____

> Note: This will break if the composer does not have an attribute called `whoami`.

Composition also allows pre-initialised class instances to be passed._____no_output_____

class Star:

    def __init__(self, mag):
        self.mag = mag

class Galaxy:

    def __init__(self, mag):
        self.mag = mag

class myClass:

    def __init__(self, composer):
        self.comp = composer
        self.mag = self.comp.mag

star = Star(11.05)
inst1 = myClass(star)
inst2 = myClass(Galaxy(8.2))
print("mag =", inst1.mag)
print("mag =", inst2.mag)_____no_output_____

This can be useful, especially if we want to be able to change the original instance attributes._____no_output_____

inst1.mag = 12.3
inst1.comp.mag = 13.1
print("inst1 mag =", inst1.mag)
print("star mag =", star.mag)_____no_output_____

Notice how the `mag` attribute of the instance `star` changed!_____no_output_____## 4. Abstract Classes

Abstract classes are classes that contain abstract methods. So it makes sense to move right along and explain what an abstract method is.

### Abstract Methods

An abstract method is a method that is defined but not implemented. This is useful for definining parent classes that impose some conditions on the child._____no_output_____

from abc import ABC, abstractmethod

class Parent(ABC):

    @abstractmethod
    def get_id(self):
        pass

print ('Parent.__dict__', Parent.__dict__)_____no_output_____

We can see that the method `get_id` has been defined but does nothing.

We now define a child class that inherits this parent._____no_output_____

class Child(Parent):

    def __init__(self):
        self.whoami = 'I am a child!'_____no_output_____try:
    inst = Child()
except Exception as error:
    print_error(error)_____no_output_____

This raises an error when trying to create an instance. This is because the child class needs to override the abstract method `get_id`. _____no_output_____

class Child(Parent):

    def __init__(self):
        self.get_id()

    def get_id(self):
        self.whoami = 'I am a child!'

inst = Child()
print(inst.whoami)_____no_output_____

This allows the parent class, which may be written with no knowledge of the properties of the child class, to impose some conditions on what a given child class should do.

### Abstract Properties

In the previous notebook we learned how class properties can be defined to impose conditions on the values of certain class attributes. Using Abstract methods we can combine these tools to ensure that a given class has the properties we want.

We start by defining an abstract class with the abstract property `whoami` and its corresponding setter._____no_output_____

class AstroObject(ABC):

    @property
    @abstractmethod
    def whoami(self):
        pass

    @whoami.setter
    @abstractmethod
    def whoami(self, value):
        pass_____no_output_____

We can then define a child class that inherits this abstract class._____no_output_____

class Star(AstroObject):

    def __init__(self):
        self.whoami = 'I am a star!'_____no_output_____try:
    star = Star()
except Exception as error:
    print_error(error)_____no_output_____

We can see that while a value for the attribute is provided the property is not explicitly defined. This way we can impose that the `Star` class must define a getter and setter for the attribute `whoami`._____no_output_____

class Star(AstroObject):

    def __init__(self):
        self.whoami = 'I am a star!'

    @property
    def whoami(self):
        return self._whoami

    @whoami.setter
    def whoami(self, value):
        if not isinstance(value, str):
            raise ValueError('whoami must be a string')
        self._whoami = value

star = Star()
print("whoami =", star.whoami)
_____no_output_____

This in turn could be used in conjunction with what we learned about class composition to ensure that composers have the properties needed to work with the child class._____no_output_____