import pygame, sys
import matplotlib.pyplot as plt
import numpy as np
import math
from scipy.integrate import ode

# set up the colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

DEFAULT_ANGLE = 45.
DEFAULT_SPEED = 50.

# clock object that ensure that animation has the same
# on all machines, regardless of the actual machine speed.
clock = pygame.time.Clock()

def load_image(name):
    image = pygame.image.load(name)
    return image

class MyCircle(pygame.sprite.Sprite):
    def __init__(self, color, width, height):
        pygame.sprite.Sprite.__init__(self)

        self.image = pygame.Surface([width, height])
        self.rect = self.image.get_rect()
        self.image.fill(WHITE)
        cx = self.rect.centerx
        cy = self.rect.centery
        pygame.draw.circle(self.image, color, (int(width/2), int(height/2)), cx, cy)
        self.rect = self.image.get_rect()

    def update(self):
        pass

class Simulation:
    def __init__(self):
        self.pos = [0, 0]
        self.vel = [0, 0]
        self.fric = 0.0001 # friction coefficient
        self.g = -9.8 # gravity acts downwards
        self.dt = 0.033 # 33 millisecond, which corresponds to 30 fps
        self.cur_time = 0

        self.paused = True

        self.solver = ode(self.f)
        self.solver.set_integrator('dop853')
        self.solver.set_f_params(self.fric, self.g)

    def f(self, t, state, arg1, arg2):
        dx  = state[2]
        dy  = state[3]
        dvx = - state[2] * arg1
        dvy = arg2 - state[3] * arg1
        return [dx, dy, dvx, dvy]

    def setup(self, speed, angle_degrees):
        self.vel[0] = speed*math.cos(angle_degrees)
        self.vel[1] = speed*math.sin(angle_degrees)
        self.cur_time = 0
        self.solver.set_initial_value([self.pos[0],self.pos[1], self.vel[0], self.vel[1]], self.cur_time)
        self.trace_x = [self.pos[0]]
        self.trace_y = [self.pos[1]]

    def step(self):
        self.cur_time += self.dt

        if self.solver.successful():
            self.solver.integrate(self.cur_time)
            self.pos = self.solver.y[0:2]
            self.vel = self.solver.y[2:4]
            self.trace_x.append(self.pos[0])
            self.trace_y.append(self.pos[1])

    def pause(self):
        self.paused = True

    def resume(self):
        self.paused = False

def sim_to_screen(win_height, x, y):
    '''flipping y, since we want our y to increase as we move up'''
    x += 10
    y += 10

    return x, win_height - y

def deg_to_rad(deg):
    return math.pi*deg/180

def usage():
    print ('--------------------------------')
    print ('Usage:')
    print ('--------------------------------')
    print ('python lab2.py <angle(degrees)> <speed>')
    print ('or')
    print ('python lab2.py')
    print ('(Angle default: 45°)')
    print ('(Speed default: 50)')
    print ('--------------------------------')

def main():

    # Error checking command line args and setting angle
    num_args = len(sys.argv) - 1
    if num_args == 2:
        angle = float(sys.argv[1])
        speed = float(sys.argv[2])
    elif num_args == 1 or num_args > 2:
        usage()
        sys.exit(0)
    else:
        angle = DEFAULT_ANGLE
        speed = DEFAULT_SPEED

    # initializing pygame
    pygame.init()

    # top left corner is (0,0)
    win_width = 640
    win_height = 480
    screen = pygame.display.set_mode((win_width, win_height))
    pygame.display.set_caption('2D projectile motion')

    # setting up a sprite group, which will be drawn on the
    # screen
    my_sprite = MyCircle(RED, 10, 10)
    my_group = pygame.sprite.Group(my_sprite)

    # setting up simulation
    sim = Simulation()

    sim.setup(speed, deg_to_rad(angle))

    print ('Usage during simulation:')
    print ('--------------------------------')
    print ('Press (r) to start/resume simulation')
    print ('Press (p) to pause simulation')
    print ('Press (q) to quit simulation')
    print ('Press (space) to step forward simulation when paused')
    print ('--------------------------------')

    while True:
        # 30 fps
        clock.tick(30)

        # update sprite x, y position using values
        # returned from the simulation
        my_sprite.rect.x, my_sprite.rect.y = sim_to_screen(win_height, sim.pos[0], sim.pos[1])

        event = pygame.event.poll()
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit(0)

        if event.type == pygame.KEYDOWN and event.key == pygame.K_p:
            sim.pause()
            continue
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_r:
            sim.resume()
            continue
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_q:
            pygame.quit()
            break
        else:
            pass

        # clear the background, and draw the sprites
        screen.fill(WHITE)
        my_group.update()
        my_group.draw(screen)
        pygame.display.flip()

        if sim.pos[1] <= -1.:
            pygame.quit()
            break

        # update simulation
        if not sim.paused:
            sim.step()
        else:
            if event.type == pygame.KEYDOWN and event.key == pygame.K_SPACE:
                sim.step()

    plt.figure(1)
    plt.plot(sim.trace_x, sim.trace_y)
    plt.xlabel('x')
    plt.ylabel('y')
    plt.axis('equal')
    plt.title('2D projectile trajectory')
    #plt.show()

    print("Angle:       " + str(angle) + "°")
    print("Speed:       " + str(speed))
    print("Distance(x): " + str(sim.trace_x[-1]))

    #with open("out.txt", "a") as myfile:
    #    myfile.write(str(angle) + ", " + str(speed) + ", " + str(sim.trace_x[-1]) + "\n")

if __name__ == '__main__':
    main()


#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
from randomwalk import RandomWalk

n_vals,mean_distance_vals = [],[]
max_n = 2000
for n in range(100,max_n,50):
	rw = RandomWalk(p=0.5,n=n,stats=False)
	rw.plot_monte_carlo(number_of_trials=3000,show=False)

	mean_distance_vals.append(rw.mc_mean_distance)
	n_vals.append(np.sqrt(rw.n))

slope, intercept = np.polyfit(n_vals, mean_distance_vals, 1)
print("Constant of proportionality k =",round(slope,3))
print("Intercept c =",round(intercept,3))

fig = plt.figure("Mean distance as a dependence on \u221An, Monte Carlo simulated")
plt.scatter(n_vals,mean_distance_vals,s=4)
plt.suptitle("s vs \u221An")
plt.xlabel('Square root of n')
plt.ylabel('mean distance, s')
plt.show()


import os
import torch

from ldm.util import instantiate_from_config

def get_resume_checkpoint_path(opt):
    if not os.path.exists(opt.resume):
        raise ValueError("Cannot find {}".format(opt.resume))
    if os.path.isfile(opt.resume):
        paths = opt.resume.split("/")
        # idx = len(paths)-paths[::-1].index("logs")+1
        # logdir = "/".join(paths[:idx])
        logdir = "/".join(paths[:-2])
        ckpt = opt.resume
    else:
        assert os.path.isdir(opt.resume), opt.resume
        logdir = opt.resume.rstrip("/")
        ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")
    _tmp = logdir.split("/")
    nowname = _tmp[-1]
    return nowname, logdir, ckpt

def create_new_logdir_and_nowname(opt, now):
    if opt.name:
        name = "_" + opt.name
    elif opt.base:
        cfg_fname = os.path.split(opt.base[0])[-1]
        cfg_name = os.path.splitext(cfg_fname)[0]
        name = "_" + cfg_name
    else:
        name = ""
    nowname = now + name + opt.postfix
    logdir = os.path.join(opt.logdir, nowname)
    return nowname, logdir

def load_model_from_config(config, ckpt, verbose=False):
    print(f"Loading model from {ckpt}")
    pl_sd = torch.load(ckpt, map_location="cpu")
    if "global_step" in pl_sd:
        print(f"Global Step: {pl_sd['global_step']}")
    sd = pl_sd["state_dict"]
    model = instantiate_from_config(config.model)
    m, u = model.load_state_dict(sd, strict=False)
    if len(m) > 0 and verbose:
        print("missing keys:")
        print(m)
    if len(u) > 0 and verbose:
        print("unexpected keys:")
        print(u)

    model.cuda()
    model.eval()
    return model


#!/usr/bin/env python
# -*- coding:utf-8 -*-
# __author__ = 'chenquanbin'
# __date__ = '2018/5/13'

def insert_sort(datalist):
    for i in range(1,len(datalist)):
        temp = datalist[i]
        for j in range(i,0,-1):
            if datalist[j-1] > temp:
                # 往右移动
                datalist[j] = datalist[j-1]
            else:
                break
            datalist[j-1] = temp

if __name__ == '__main__':
    datalist = [5, 3, 8, 1, -4, 0, 100, 66, 32, 75]
    insert_sort(datalist)
    print(datalist)


#!/usr/bin/python3
"""
Class FileStorage
"""

import json
from models.amenity import Amenity
from models.base_model import BaseModel
from models.city import City
from models.place import Place
from models.review import Review
from models.state import State
from models.user import User

classes = {"Amenity": Amenity, "BaseModel": BaseModel, "City": City,
           "Place": Place, "Review": Review, "State": State, "User": User}

class FileStorage:
    """objects with Strings/JSON formart to transfer"""

    __file_path = "file.json"
    __objects = {}

    def all(self, cls=None):
        """get all objectss"""
        if cls is not None:
            new_dict = {}
            for key, value in self.__objects.items():
                if cls == value.__class__ or cls == value.__class__.__name__:
                    new_dict[key] = value
            return new_dict
        return self.__objects

    def new(self, obj):
        """add object with HASH"""
        if obj is not None:
            key = obj.__class__.__name__ + "." + obj.id
            self.__objects[key] = obj

    def save(self):
        """save JSON file"""
        json_objects = {}
        for key in self.__objects:
            json_objects[key] = self.__objects[key].to_dict()
        with open(self.__file_path, 'w') as f:
            json.dump(json_objects, f)

    def reload(self):
        """load JSON file"""
        try:
            with open(self.__file_path, 'r') as f:
                jo = json.load(f)
            for key in jo:
                self.__objects[key] = classes[jo[key]["__class__"]](**jo[key])
        except:
            pass

    def delete(self, obj=None):
        """remove item/s Object"""
        if obj is not None:
            key = obj.__class__.__name__ + '.' + obj.id
            if key in self.__objects:
                del self.__objects[key]

    def close(self):
        """call again fn reload"""
        self.reload()


import Converter

def present_options(converter_one, converter_two):
    print(f'Input 1 for {converter_one.__name__.replace("_", " ")}')
    print(f'Input 2 for {converter_two.__name__.replace("_", " ")}')
    chosen_option = int(input())

    if chosen_option == 1:
        user_input = input('input the amount(s)')
        if len(user_input.split(" ")) > 1:
            inputs = list(map(lambda num: int(num), user_input.split(" ")))
            print(converter_one(inputs[0], inputs[1]))
        else:
            print(converter_one(int(user_input)))
    elif chosen_option == 2:
        user_input = int(input('input the amount'))
        print(converter_two(user_input))
    else:
        raise Exception("Invalid menu option was chosen")

def program():
    while True:
        print('Select a conversion option:')
        print('1. Length')
        print('2. Mass')
        print('3. Temperature')
        print('4. Time')
        print('5. Quit')
        chosen_option = int(input())

        try:
            if chosen_option == 1:
                present_options(Converter.feet_and_inches_to_meters, Converter.meters_to_feet_and_inches)
            elif chosen_option == 2:
                present_options(Converter.pounds_to_kg, Converter.kg_to_pounds)
            elif chosen_option == 3:
                present_options(Converter.kelvin_to_celsius, Converter.celsius_to_kelvin)
            elif chosen_option == 4:
                present_options(Converter.hours_and_minutes_to_seconds, Converter.seconds_to_hours_and_minutes)
            elif chosen_option == 5:
                return
        except Exception as e:
            print(str(e))

program()
# didn't implement user input validaiton - because this is not graded


import string
import sys

alphabet = list(string.ascii_lowercase)

with open("input.txt") as f:
    input = f.read().splitlines()
    input = [[n for n in line] for line in input]

def dijkstra(graph, start, target):
    distances = {}
    prev = {}
    Q = set()

    for n in graph:
        distances[n] = sys.maxsize
        prev[n] = None
        Q.add(n)

    distances[start] = 0

    while Q:
        min_dist = sys.maxsize
        u = None

        for n in Q:
            if distances[n] < min_dist:
                min_dist = distances[n]
                u = n

        if u is None:
            u = Q.pop()
        else:
            Q.remove(u)

        if u == target:
            return distances, prev

        xu, yu = u

        candidates = []
        for move in [(1, 0), (0, 1), (-1, 0), (0, -1)]:
            candidates.append((xu + move[0], yu + move[1]))

        candidates = [c for c in candidates if c in Q]

        for c in candidates:
            alt = distances[u]

            xc, yc = c

            if input[yc][xc] == "E":
                current_elevation = alphabet.index("z") + 1
            elif input[yc][xc] == "S":
                current_elevation = 0
            else:
                current_elevation = alphabet.index(input[yc][xc])

            if input[yu][xu] not in ["S", "E"]:
                if current_elevation > alphabet.index(input[yu][xu]) + 1:
                    alt = sys.maxsize
                else:
                    alt += 1

            if alt < distances[c]:
                distances[c] = alt
                prev[c] = u

    return distances, prev

def get_graph_start_target(data):
    graph = set()
    start = None
    target = None

    for y in range(len(data)):
        for x in range(len(data[0])):
            graph.add((x, y))
            if data[y][x] == "S":
                start = (x, y)
            if data[y][x] == "E":
                target = (x, y)

    return graph, start, target

graph, start, target = get_graph_start_target(input)
_, prev = dijkstra(graph, start, target)

path = []
current = target
while True:
    if current == target:
        break
    path.append(current)
    current = prev[current]

print(len(path))

def get_candidates(data):
    candidates = []

    for y in range(len(data)):
        for x in range(len(data[0])):
            if data[y][x] in ["a", "S"]:
                candidates.append((x, y))

    return candidates

start_positions = get_candidates(input)
start_positions = [(x, y) for (x, y) in start_positions if x == 0]

lengths = []
for s in start_positions:
    d, p = dijkstra(graph, s, target)

    path = []
    current = target
    last = None
    while True:
        if current == s:
            break

        path.append(current)
        current = p[current]

    lengths.append(len(path))

print(min(lengths))


from typing import Optional, Type, List

import pytorch_lightning
import torch

from agz.MCTS_recurse import MCTS
from games.Game import Game

class Simulator:
    def __init__(self, game: Type[Game], mcts_iterations=100, device="cpu"):
        self.do_mcts = True
        self.game = game
        self.mcts_object = MCTS(mcts_iterations, device=device)
        self.temperature_threshold = 0.3

    def nn_move(self, network, game: Game, mcts=True):
        """
          Tile format (8 bits);
            [ NN C TTTTT ]

        Selects a move based on the policy vector and updates the board state.
        Returns the policy vector with which the move was chosen.

        :param network: the neural network for which to compute the move
        :param game: the gamestate for which to compute the move
        :param mcts: use mcts to compute an improvement policy?
        :return: the taken move index in the policy vector
        """

        # Select only best moves with low enough temperature
        # temperature = (float(game.turn_limit) - float(game.node.turn)) / float(game.turn_limit)
        temperature = game.node.turn < 15

        if mcts:
            policy = self.mcts_object.process(game, network)
        else:
            game.node.expand()
            policy, _ = network.predict(game.node.to_np(game.node.to_play))

            invalid_mask = torch.ones(policy.shape, dtype=torch.bool)
            for child in game.node.children:
                invalid_mask[child.encode()] = False

            policy[invalid_mask] = 0
            p_sum = torch.sum(policy)
            if p_sum == 0:
                policy = invalid_mask / torch.sum(invalid_mask)
            else:
                policy /= p_sum

        # Select a move based on updated_policy's weights out of the valid moves (encodings)
        best_idx = torch.multinomial(policy, 1).item()

        selected = None
        for child in game.node.children:
            if child.encode() == best_idx:
                selected = child
                break

        assert selected is not None, f"{best_idx} not " \
                                     f"in {[child.encode() for child in game.node.children]} " \
                                     f"(policy: {policy})" \
                                     f"mcts: {mcts}"
        game.select_child(selected)

        return policy

    def play(self,
             game: Game,
             networks: List[Optional[pytorch_lightning.LightningModule]],
             mcts=True):
        """
        Does a playout of the game with two neural networks, returns all board states and the final result.
        The returned zip contains all boards which were played by player 1 including the policy vector.
        This can later be used to train the neural network.

        :param mcts: Use MCTS for policy enhancement
        :param game: the board state from which to choose
        :param networks: the neural networks for all players
        :return: zip(boards, policies), result
        """
        policy_vectors = []
        # Play game until a terminal state has occurred.
        result = -1
        for i in range(game.turn_limit):

            nn = networks[game.to_play()]

            policy_vector = None
            if nn is None:
                game.ai_move(algorithm="random")
            else:
                policy_vector = self.nn_move(nn, game, mcts=mcts)
                policy_vector = policy_vector.to("cpu")

            policy_vectors.append(policy_vector)

            result = game.winner()

            if result is not None:
                break

        # Return all board states and the final result.
        return zip(game.history[:-1], policy_vectors), result

    def parallel_play(self, nn_perspective: int, net1: pytorch_lightning.LightningModule,
                      net2: Optional[pytorch_lightning.LightningModule] = None):
        """
        Wrapper for the play function preparing the data and initializing the Hive object.
        It will convert the arrays to tensors to be used by pytorch.
        This function is threadsafe.
        """
        tensors = []
        policy_vectors = []

        game = self.game()

        node_result, result = self.play(game, [net1, net2])
        node_result = list(node_result)
        for board_idx, (node, policy_vector) in enumerate(node_result):
            # Convert all data to be usable
            # if board_idx > len(node_result) - 5:
            #     print(node.to_np(node.to_play))
            #     print(policy_vector)
            if node.to_play == nn_perspective:
                tensors.append(torch.Tensor(node.to_np(node.to_play)))
                policy_vectors.append(policy_vector)

        # Convert game outcome to numerical value vector
        result_value = (
            0 if (result == -1)
            else 1 if (result == nn_perspective)
            else -1
        )
        result_values = [torch.Tensor([result_value])] * len(tensors)

        # Invert tensors
        # new_tensors = game.get_inverted(tensors)
        # if len(new_tensors) > 0:
        #     tensors += new_tensors
        #     policy_vectors += policy_vectors
        #     result_values += ([torch.Tensor([result_value * -1])] * len(new_tensors))

        return tensors, policy_vectors, result_values, nn_perspective


import sys, os
import numpy as np
import jax
import jax.numpy as jnp
from jax import random, grad, jit, vmap, jacfwd, jacrev
from functools import partial

class NuclearPotential(object):
    def __init__(self, nwalk, pot_name, pot_3b_name):
        self.nwalk=nwalk
        self.hc = 197.327053
        self.alpha = 1 / 137.03599
        self.b = 4.27
        if (pot_name == 'pionless_2'):
            self.vkr = 2.0
            self.v0r = -133.3431
            self.v0s = -9.0212
            self.ar3b = np.sqrt(68.48830)
            self.v_2b = self.pionless_2b
            self.v_3b = self.pionless_3b
        elif (pot_name == 'pionless_4'):
            self.vkr = 4.0
            self.v0r = -487.6128
            self.v0s = -17.5515
            self.ar3b = np.sqrt(677.79890)
            self.v_2b = self.pionless_2b
            self.v_3b = self.pionless_3b
        elif (pot_name == 'pionless_6'):
            self.vkr = 6.0
            self.v0r = -1064.5010
            self.v0s = -26.0830
            self.ar3b = np.sqrt(2652.65100)
            self.v_2b = self.pionless_2b
            self.v_3b = self.pionless_3b
        elif (pot_name == 'pionless_lo_a'):
            self.R0 = 1.7
            self.R1 = 1.5
            self.C01 = -4.38524414
            self.C10 = -8.00783936
            lamchi = 1000./self.hc
            fpi = 92.40/self.hc
            if (pot_3b_name == 'R3_1.0'):
               self.R3 = 1.0
               self.ce3b = 1.8354
            elif (pot_3b_name == 'R3_1.5'):
               self.R3 = 1.5
               self.ce3b = 4.6301
            elif (pot_3b_name == 'R3_2.0'):
               self.R3 = 2.0
               self.ce3b =  11.6871
            elif (pot_3b_name == 'R3_2.5'):
               self.R3 = 2.5
               self.ce3b =  27.4702
            else:
               self.R3 = 1.
               self.ce3b =  0
            self.ce3b = jnp.sqrt( self.ce3b / lamchi / fpi**4 * self.hc / jnp.pi**3 / self.R3**6 )
            self.v_2b = self.pionless_2b_lo
            self.v_3b = self.pionless_3b_lo
        elif (pot_name == 'pionless_lo_o'):
            self.R0 = 1.54592984
            self.R1 = 1.83039397
            self.C01 = -5.27518671
            self.C10 = -7.04040080
            lamchi = 1000./self.hc
            fpi = 92.40/self.hc
            if (pot_3b_name == 'R3_1.0'):
               self.R3 = 1.0
               self.ce3b = 1.0786
            if (pot_3b_name == 'R3_1.1'):
               self.R3 = 1.1
               self.ce3b = 1.2945
            elif (pot_3b_name == 'R3_1.5'):
               self.R3 = 1.5
               self.ce3b = 2.7676
            elif (pot_3b_name == 'R3_2.0'):
               self.R3 = 2.0
               self.ce3b =  6.95356
            elif (pot_3b_name == 'R3_2.5'):
               self.R3 = 2.5
               self.ce3b =  16.21993
            else:
               self.R3 = 1.
               self.ce3b =  0
            self.ce3b = jnp.sqrt( self.ce3b / lamchi / fpi**4 * self.hc / jnp.pi**3 / self.R3**6 )
            self.v_2b = self.pionless_2b_lo
            self.v_3b = self.pionless_3b_lo

        nr_test = 100
        r_test = jnp.linspace(0,5,nr_test)
        file_pot = open("pot_nn.dat", "w")
        pot_test = np.zeros(shape=(8,nr_test))
        pot_test[0,:] = r_test
        for i in range(r_test.shape[0]):
            pot_test[1:7,i] =  self.v_2b(r_test[i])
            pot_test[7,i] =  self.v_em(r_test[i])
        np.savetxt(file_pot, np.transpose(pot_test), delimiter=' ', newline=os.linesep)
        file_pot.close()

    @partial(jit, static_argnums=(0,))
    def pionless_2b(self, rr):
        pot_2b=jnp.zeros(6)
        x = self.vkr * rr
        vr = jnp.exp( -x**2 / 4.0 )
        pot_2b = pot_2b.at[0].set(self.v0r * vr)
        pot_2b = pot_2b.at[2].set(self.v0s * vr)
        return pot_2b

    @partial(jit, static_argnums=(0,))
    def pionless_3b(self, rr):
        x = self.vkr * rr
        vr = jnp.exp( -x**2 / 4.0 )
        pot_3b = self.ar3b * vr
        return pot_3b

    @partial(jit, static_argnums=(0,))
    def pionless_2b_lo(self, rr):
        pot_2b = jnp.zeros(6)
        C0_r = 1. / (jnp.sqrt(jnp.pi)*self.R0)**3*jnp.exp( -( rr / self.R0 )**2 )
        C1_r = 1. / (jnp.sqrt(jnp.pi)*self.R1)**3*jnp.exp( -( rr / self.R1 )**2 )
        pot_2b = pot_2b.at[0].set( 3. * ( self.C01 * C1_r + self.C10 * C0_r ) )
        pot_2b = pot_2b.at[1].set( self.C01 * C1_r - 3. * self.C10 * C0_r )
        pot_2b = pot_2b.at[2].set( -3. * self.C01 * C1_r + self.C10 * C0_r )
        pot_2b = pot_2b.at[3].set( -1. * ( self.C01 * C1_r + self.C10 * C0_r ) )
        pot_2b = pot_2b / 16. * self.hc
        return pot_2b

    @partial(jit, static_argnums=(0,))
    def pionless_3b_lo(self, rr):
        pot_3b = self.ce3b * jnp.exp( - ( rr / self.R3 )**2 )
        return pot_3b

    @partial(jit, static_argnums=(0,))
    def v_em(self, rr):
        rr = jnp.maximum(rr,0.0001)
        br = self.b * rr
        fcoul = 1 - (1 + 11 * br / 16 + 3 * br**2 / 16 + br**3 / 48) * jnp.exp(-br)
        pot_em = self.alpha * self.hc * fcoul / rr
        return pot_em


"""
This file uses modified code from nvidia apex (https://github.com/NVIDIA/apex)
"""

import torch
import torch.nn as nn
from torch.autograd import Variable
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
from torch.optim import Optimizer
import copy

class MultiSequential(nn.Sequential):
    def __init__(self, *args):
        super(MultiSequential, self).__init__(*args)

    def forward(self, *input):
        for module in self._modules.values():
            input = module(*input)
        return input

class LossScaler:
    """
    Class that manages a static loss scale.  This class is intended to interact with
    :class:`FP16_Optimizer`, and should not be directly manipulated by the user.

    Use of :class:`LossScaler` is enabled via the ``static_loss_scale`` argument to
    :class:`FP16_Optimizer`'s constructor.

    Args:
        scale (float, optional, default=1.0):  The loss scale.
    """

    def __init__(self, scale=1):
        self.cur_scale = scale

    # `params` is a list / generator of torch.Variable
    def has_overflow(self, params):
        return False

    # `x` is a torch.Tensor
    def _has_inf_or_nan(x):
        return False

    def update_scale(self, overflow):
        pass

    @property
    def loss_scale(self):
        return self.cur_scale

    def scale_gradient(self, module, grad_in, grad_out):
        return tuple(self.loss_scale * g for g in grad_in)

    def backward(self, loss):
        scaled_loss = loss*self.loss_scale
        scaled_loss.backward()

class DynamicLossScaler:
    """
    Class that manages dynamic loss scaling.  It is recommended to use :class:`DynamicLossScaler`
    indirectly, by supplying ``dynamic_loss_scale=True`` to the constructor of
    :class:`FP16_Optimizer`.  However, it's important to understand how :class:`DynamicLossScaler`
    operates, because the default options can be changed using the
    the ``dynamic_loss_args`` argument to :class:`FP16_Optimizer`'s constructor.

    Loss scaling is designed to combat the problem of underflowing gradients encountered at long
    times when training fp16 networks.  Dynamic loss scaling begins by attempting a very high loss
    scale.  Ironically, this may result in OVERflowing gradients.  If overflowing gradients are
    encountered, :class:`DynamicLossScaler` informs :class:`FP16_Optimizer` that an overflow has
    occurred.
    :class:`FP16_Optimizer` then skips the update step for this particular iteration/minibatch,
    and :class:`DynamicLossScaler` adjusts the loss scale to a lower value.
    If a certain number of iterations occur without overflowing gradients detected,
    :class:`DynamicLossScaler` increases the loss scale once more.
    In this way :class:`DynamicLossScaler` attempts to "ride the edge" of
    always using the highest loss scale possible without incurring overflow.

    Args:
        init_scale (float, optional, default=2**32):  Initial loss scale attempted by :class:`DynamicLossScaler.`
        scale_factor (float, optional, default=2.0):  Factor used when adjusting the loss scale. If an overflow is encountered, the loss scale is readjusted to loss scale/``scale_factor``.  If ``scale_window`` consecutive iterations take place without an overflow, the loss scale is readjusted to loss_scale*``scale_factor``.
        scale_window (int, optional, default=1000):  Number of consecutive iterations without an overflow to wait before increasing the loss scale.
    """

    def __init__(self,
                 init_scale=2**32,
                 scale_factor=2.,
                 scale_window=1000):
        self.cur_scale = init_scale
        self.cur_iter = 0
        self.last_overflow_iter = -1
        self.scale_factor = scale_factor
        self.scale_window = scale_window

    # `params` is a list / generator of torch.Variable
    def has_overflow(self, params):
        for p in params:
            if p.grad is not None and DynamicLossScaler._has_inf_or_nan(p.grad.data):
                return True

        return False

    # `x` is a torch.Tensor
    def _has_inf_or_nan(x):
        try:
            # if x is half, the .float() incurs an additional deep copy, but it's necessary if
            # Pytorch's .sum() creates a one-element tensor of the same type as x
            # (which is true for some recent version of pytorch).
            cpu_sum = float(x.float().sum())
            # More efficient version that can be used if .sum() returns a Python scalar
            # cpu_sum = float(x.sum())
        except RuntimeError as instance:
            # We want to check if inst is actually an overflow exception.
            # RuntimeError could come from a different error.
            # If so, we still want the exception to propagate.
            if "value cannot be converted" not in instance.args[0]:
                raise
            return True
        else:
            if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
                return True
            return False

    # `overflow` is boolean indicating whether the gradient overflowed
    def update_scale(self, overflow):
        if overflow:
            # self.cur_scale /= self.scale_factor
            self.cur_scale = max(self.cur_scale/self.scale_factor, 1)
            self.last_overflow_iter = self.cur_iter
        else:
            if (self.cur_iter - self.last_overflow_iter) % self.scale_window == 0:
                self.cur_scale *= self.scale_factor
        self.cur_iter += 1

    @property
    def loss_scale(self):
        return self.cur_scale

    def scale_gradient(self, module, grad_in, grad_out):
        return tuple(self.loss_scale * g for g in grad_in)

    def backward(self, loss):
        scaled_loss = loss*self.loss_scale
        scaled_loss.backward()

class Convert(nn.Module):

    def __init__(self):
        super(Convert, self).__init__()

    def forward(self, *inputs):

        return (input.half() for input in inputs)

def BN_convert_float(module):
    '''
    BatchNorm layers need parameters in single precision.
    Find all layers and convert them back to float. This can't
    be done with built in .apply as that function will apply
    fn to all modules, parameters, and buffers. Thus we wouldn't
    be able to guard the float conversion based on the module type.
    '''
    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
        module.float()

    for child in module.children():
        BN_convert_float(child)

    return module

def half_model(model):
    """
    Convert model to half precision in a batchnorm-safe way.
    """

    return MultiSequential(Convert(), BN_convert_float(model.half()))

def convertToFP16(model,optimizer,loss_scale=None,dynamic_scale_args=None):

    dynamic_loss_scale = True
    static_loss_scale = 128
    if loss_scale != None:
        dynamic_loss_scale = False
        static_loss_scale = loss_scale

    return half_model(model), FP16_Optimizer(optimizer,dynamic_loss_scale=dynamic_loss_scale,static_loss_scale=static_loss_scale,dynamic_loss_args=dynamic_scale_args)

def convertToFP32(model):

    if type(model) == nn.DataParallel:
        model = model.module

    if isinstance(model,MultiSequential):

        for child in model.children():

            if not isinstance(child,Convert):
                model = child
                break

    model = copy.deepcopy(model).float()
    return model

def backwards_debug_hook(grad):
    raise RuntimeError("master_params recieved a gradient in the backward pass!")

def prep_param_lists(model, flat_master=False):
    """
    Creates a list of FP32 master parameters for a given model, as in
    `Training Neural Networks with Mixed Precision:  Real Examples`_.

    Args:
        model (torch.nn.Module): Existing Pytorch model
        flat_master (bool, optional, default=False):  Flatten the master parameters into a single tensor, as a performance optimization.
    Returns:
        A tuple (``model_params``, ``master_params``). ``model_params`` is a list of the model's parameters for later use with :func:`model_grads_to_master_grads` and :func:`master_params_to_model_params`.  ``master_params`` is a list of FP32 master gradients.  If ``flat_master=True``, ``master_params`` will be a list with one element.

    Example::

        model_params, master_params = prep_param_lists(model)

    .. warning::
        Currently, if ``flat_master=True``, all the model's parameters must be the same type.  If the model has parameters of different types, use ``flat_master=False``, or use :class:`FP16_Optimizer`.

    .. _`Training Neural Networks with Mixed Precision:  Real Examples`:
        http://on-demand.gputechconf.com/gtc/2018/video/S81012/
    """
    model_params = [param for param in model.parameters() if param.requires_grad]

    if flat_master:
        # Give the user some more useful error messages
        try:
            # flatten_dense_tensors returns a contiguous flat array.
            # http://pytorch.org/docs/master/_modules/torch/_utils.html
            master_params = _flatten_dense_tensors([param.data for param in model_params]).float()
        except:
            print("Error in prep_param_lists:  model may contain a mixture of parameters "
                      "of different types.  Use flat_master=False, or use F16_Optimizer.")
            raise
        master_params = torch.nn.Parameter(master_params)
        master_params.requires_grad = True
        # master_params.register_hook(backwards_debug_hook)
        if master_params.grad is None:
            master_params.grad = master_params.new(*master_params.size())
        return model_params, [master_params]
    else:
        master_params = [param.clone().float().detach() for param in model_params]
        for param in master_params:
            param.requires_grad = True
        return model_params, master_params

def model_grads_to_master_grads(model_params, master_params, flat_master=False):
    """
    Copy model gradients to master gradients.

    Args:
        model_params:  List of model parameters created by :func:`prep_param_lists`.
        master_params:  List of FP32 master parameters created by :func:`prep_param_lists`.  If ``master_params`` was created with ``flat_master=True``, ``flat_master=True`` should also be supplied to :func:`model_grads_to_master_grads`.
    """
    if flat_master:
        # The flattening may incur one more deep copy than is necessary.
        master_params[0].grad.data.copy_(
            _flatten_dense_tensors([p.grad.data for p in model_params]))
    else:
        for model, master in zip(model_params, master_params):
            if model.grad is not None:
                if master.grad is None:
                    master.grad = Variable(master.data.new(*master.data.size()))
                master.grad.data.copy_(model.grad.data)
            else:
                master.grad = None

def master_params_to_model_params(model_params, master_params, flat_master=False):
    """
    Copy master parameters to model parameters.

    Args:
        model_params:  List of model parameters created by :func:`prep_param_lists`.
        master_params:  List of FP32 master parameters created by :func:`prep_param_lists`.  If ``master_params`` was created with ``flat_master=True``, ``flat_master=True`` should also be supplied to :func:`master_params_to_model_params`.
    """
    if flat_master:
        for model, master in zip(model_params,
                                 _unflatten_dense_tensors(master_params[0].data, model_params)):
            model.data.copy_(master)
    else:
        for model, master in zip(model_params, master_params):
            model.data.copy_(master.data)

# Backward compatibility fixes

def to_python_float(t):
    if hasattr(t, 'item'):
        return t.item()
    else:
        return t[0]

TORCH_MAJOR = int(torch.__version__.split('.')[0])
TORCH_MINOR = int(torch.__version__.split('.')[1])
if TORCH_MAJOR == 0 and TORCH_MINOR <= 4:
    clip_grad_norm = torch.nn.utils.clip_grad_norm
else:
    clip_grad_norm = torch.nn.utils.clip_grad_norm_

# TODO:  Update overflow check + downscale to use Carl's fused kernel.
class FP16_Optimizer(Optimizer):
    """
    :class:`FP16_Optimizer` is designed to wrap an existing PyTorch optimizer,
    and manage static or dynamic loss scaling and master weights in a manner transparent to the user.
    For standard use, only two lines must be changed:  creating the :class:`FP16_Optimizer` instance,
    and changing the call to ``backward``.

    Example::

        model = torch.nn.Linear(D_in, D_out).cuda().half()
        optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
        # Name the FP16_Optimizer instance to replace the existing optimizer
        # (recommended but not required):
        optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
        ...
        # loss.backward() becomes:
        optimizer.backward(loss)
        ...

    Example with dynamic loss scaling::

        ...
        optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
                                   # optional arg to control dynamic loss scaling behavior
                                   # dynamic_loss_args={'scale_window' : 500})
                                   # Usually, dynamic_loss_args is not necessary.

    Args:
        init_optimizer (torch.optim.optimizer):  Existing optimizer created with the parameters to optimize.  Internally, :class:`FP16_Optimizer` replaces the passed optimizer's fp16 parameters, if any, with fp32 master parameters copied from the original ones.  :class:`FP16_Optimizer` also stores references to the original fp16 parameters, and updates these fp16 parameters from the master fp32 copy at the end of each :attr:`step`.
        static_loss_scale (float, optional, default=1.0):  Loss scale used internally to scale gradients computed by the model.  Any fp16 gradients will be copied to fp32, then downscaled before being applied to the fp32 master params, so ``static_loss_scale`` should not affect learning rate.
        dynamic_loss_scale (bool, optional, default=False):  Use dynamic loss scaling.  If True, this will override any ``static_loss_scale`` option.
        dynamic_loss_args (dict, optional, default=None):  Dict of kwargs that will be forwarded to the internal :class:`DynamicLossScaler` instance's constructor.  Keys of this dict must match kwargs accepted by :class:`DynamicLossScaler`'s constructor.  If ``dynamic_loss_args`` is unspecified, :class:`DynamicLossScaler`'s defaults will be used.
        verbose (bool, optional, default=True):  By default, FP16_Optimizer's constructor prints out the parameters and parameter groups it is ingesting, as a sanity check.  If this becomes annoying (e.g. for large models), it can be disabled by passing ``verbose=False``.  ``verbose=False`` will not disable printing when the loss scale is readjusted during dynamic loss scaling.

    ``init_optimizer`` is expected to have been constructed in the ordinary way.
    It is recommended (although not required) that the newly constructed :class:`FP16_Optimizer` instance be
    named to replace ``init_optimizer``, for two reasons:
    First, it means that references to the same name
    later in the file will not have to change.
    Second, :class:`FP16_Optimizer` reserves the right (as an implementation detail) to
    modify ``init_optimizer``.  If you do choose a unique name for the new
    :class:`FP16_Optimizer` instance, you should only work with this new instance,
    because the preexisting optimizer might no longer behave as expected.

    ``init_optimizer`` may be any Pytorch optimizer.
    It may contain a mixture of fp16 and fp32 parameters organized into any number of
    ``param_groups`` with different hyperparameters.  The :class:`FP16_Optimizer` constructor will
    ingest these ``param_groups`` and remember them.

    Calls to ::

        loss.backward()

    must be replaced with ::

        optimizer.backward(loss)

    because :class:`FP16_Optimizer` requires ownership of the backward pass to implement
    loss scaling and copies to master gradients.

    .. note::
        Loss scaling, either static or dynamic, is orthogonal to learning rate, because gradients
        are downscaled before being applied.  This means that adjusting the loss scale, or using
        dynamic loss scaling, should not require retuning the learning rate or any other
        hyperparameters.

    **Advanced options**

    **Closures**:  :class:`FP16_Optimizer` can wrap a Pytorch optimizer that receives a closure.
    See docstring for :attr:`step`.

    **Gradient clipping**:  Use :attr:`clip_master_grads`.

    **Multiple losses**:  If your model accumulates gradients from multiple losses,
    this can be made more efficient by supplying ``update_master_grads=False``
    to :attr:`backward`.  See docstring for :attr:`backward`.

    **Manually adjusting loss scale**:  The current loss scale can be retrieved or set via ::

        print(optimizer.loss_scale)
        optimizer.loss_scale = new_loss_scale

    For static loss scaling, manually adjusting the loss scale over time is a reasonable
    thing to do.  During later epochs, gradients may become smaller, and a
    higher loss scale may be required, analogous to scheduling the learning rate.  Dynamic loss
    scaling is more subtle (see :class:`DynamicLossScaler`) and in this case, manually adjusting
    the loss scale is not recommended.

    **Multi_GPU training**:  If the wrapped ``init_optimizer`` was created from a model wrapped in
    Pytorch DistributedDataParallel or Apex DistributedDataParallel, :class:`FP16_Optimizer`
    should still work as intended.
    """

    def __init__(self,
                 init_optimizer,
                 static_loss_scale=1.0,
                 dynamic_loss_scale=True,
                 dynamic_loss_args=None,
                 verbose=False):

        if not torch.cuda.is_available:
            raise SystemError("Cannot use fp16 without CUDA.")

        self.verbose = verbose

        self.optimizer = init_optimizer

        self.param_groups = init_optimizer.param_groups
        self.defaults = init_optimizer.defaults
        # init_state_dict sets up an alternative way to cast per-param state tensors.
        # Stashing here in case https://github.com/pytorch/pytorch/issues/7733 makes it necessary.
        # init_state_dict = init_optimizer.state_dict()

        self.fp16_groups = []
        self.fp32_from_fp16_groups = []
        self.fp32_from_fp32_groups = []
        for i, param_group in enumerate(self.optimizer.param_groups):
            self.maybe_print("FP16_Optimizer processing param group {}:".format(i))
            fp16_params_this_group = []
            fp32_params_this_group = []
            fp32_from_fp16_params_this_group = []
            for i, param in enumerate(param_group['params']):
                if param.requires_grad:
                    if param.type() == 'torch.cuda.HalfTensor':
                        self.maybe_print("FP16_Optimizer received torch.cuda.HalfTensor with {}"
                                         .format(param.size()))
                        fp16_params_this_group.append(param)
                        master_param = param.detach().clone().float()
                        master_param.requires_grad = True
                        param_group['params'][i] = master_param
                        fp32_from_fp16_params_this_group.append(master_param)
                        # Reset existing state dict key to the new master param.
                        # We still need to recast per-param state tensors, if any, to FP32.
                        if param in self.optimizer.state:
                           self.optimizer.state[master_param] = self.optimizer.state.pop(param)
                    elif param.type() == 'torch.cuda.FloatTensor':
                        self.maybe_print("FP16_Optimizer received torch.cuda.FloatTensor with {}"
                                         .format(param.size()))
                        fp32_params_this_group.append(param)
                        param_group['params'][i] = param
                    else:
                        raise TypeError("Wrapped parameters must be either "
                                        "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
                                        "Received {}".format(param.type()))

            self.fp16_groups.append(fp16_params_this_group)
            self.fp32_from_fp16_groups.append(fp32_from_fp16_params_this_group)
            self.fp32_from_fp32_groups.append(fp32_params_this_group)

        # Leverage state_dict() and load_state_dict() to recast preexisting per-param state tensors
        self.optimizer.load_state_dict(self.optimizer.state_dict())
        # alternative way to cast per-param state tensors:
        # self.optimizer.load_state_dict(init_state_dict)

        if dynamic_loss_scale:
            self.dynamic_loss_scale = True
            if dynamic_loss_args is not None:
                self.loss_scaler = DynamicLossScaler(**dynamic_loss_args)
            else:
                self.loss_scaler = DynamicLossScaler()
        else:
            self.dynamic_loss_scale = False
            self.loss_scaler = LossScaler(static_loss_scale)

        self.overflow = False
        self.first_closure_call_this_step = True

        self.clip_grad_norm = clip_grad_norm

    def maybe_print(self, msg):
        if self.verbose:
            print(msg)

    def __getstate__(self):
        return self.optimizer.__getstate__()

    def __setstate__(self, state):
        self.optimizer.__setstate__(state)

    def zero_grad(self, set_grads_to_None=False):
        """
        Zero fp32 and fp16 parameter grads.
        """
        # In principle, only the .grad attributes of the model params need to be zeroed,
        # because gradients are copied into the FP32 master params.  However, we zero
        # all gradients owned by the optimizer, just to be safe:
        for group in self.optimizer.param_groups:
             for p in group['params']:
                 if set_grads_to_None:
                     p.grad = None
                 else:
                     if p.grad is not None:
                         p.grad.detach_()
                         p.grad.zero_()

        # Zero fp16 gradients owned by the model:
        for fp16_group in self.fp16_groups:
            for param in fp16_group:
                if set_grads_to_None:
                    param.grad = None
                else:
                    if param.grad is not None:
                        param.grad.detach_() # as in torch.optim.optimizer.zero_grad()
                        param.grad.zero_()

    def _check_overflow(self):
        params = []
        for group in self.fp16_groups:
            for param in group:
                params.append(param)
        for group in self.fp32_from_fp32_groups:
            for param in group:
                params.append(param)
        self.overflow = self.loss_scaler.has_overflow(params)

    def _update_scale(self, has_overflow=False):
        self.loss_scaler.update_scale(has_overflow)

    def _master_params_to_model_params(self):
        for fp16_group, fp32_from_fp16_group in zip(self.fp16_groups, self.fp32_from_fp16_groups):
            master_params_to_model_params(fp16_group, fp32_from_fp16_group)

    # To consider:  Integrate distributed with this wrapper by registering a hook on each variable
    # that does the overflow check, gradient copy + downscale, and fp32 allreduce in a different stream.
    def _model_grads_to_master_grads(self):
        for fp16_group, fp32_from_fp16_group in zip(self.fp16_groups, self.fp32_from_fp16_groups):
            model_grads_to_master_grads(fp16_group, fp32_from_fp16_group)

    def _downscale_master(self):
        if self.loss_scale != 1.0:
            for group in self.optimizer.param_groups:
                for param in group['params']:
                    if param.grad is not None:
                        param.grad.data.mul_(1./self.loss_scale)

    def clip_master_grads(self, max_norm, norm_type=2):
        """
        Clips fp32 master gradients via ``torch.nn.utils.clip_grad_norm``.

        Args:
            max_norm (float or int): max norm of the gradients
            norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
                infinity norm.

        Returns:
            Total norm of the current fp32 gradients (viewed as a single vector).

        .. warning::
            Returns -1 if the most recently computed fp16 gradients overflowed (that is, if ``self.overflow`` is ``True``).
        """
        if not self.overflow:
            fp32_params = []
            for param_group in self.optimizer.param_groups:
                for param in param_group['params']:
                    fp32_params.append(param)
            return self.clip_grad_norm(fp32_params, max_norm, norm_type)
        else:
            return -1

    def state_dict(self):
        """
        Returns a dict containing the current state of this :class:`FP16_Optimizer` instance.
        This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict
        of the contained Pytorch optimizer.
        Example::

            checkpoint = {}
            checkpoint['model'] = model.state_dict()
            checkpoint['optimizer'] = optimizer.state_dict()
            torch.save(checkpoint, "saved.pth")
        """
        state_dict = {}
        state_dict['loss_scaler'] = self.loss_scaler
        state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale
        state_dict['overflow'] = self.overflow
        state_dict['first_closure_call_this_step'] = self.first_closure_call_this_step
        state_dict['optimizer_state_dict'] = self.optimizer.state_dict()
        state_dict['fp32_from_fp16'] = self.fp32_from_fp16_groups
        return state_dict

    def load_state_dict(self, state_dict):
        """
        Loads a state_dict created by an earlier call to state_dict().
        If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``,
        whose parameters in turn came from ``model``, it is expected that the user
        will call ``model.load_state_dict()`` before
        ``fp16_optimizer_instance.load_state_dict()`` is called.

        Example::

            model = torch.nn.Linear(D_in, D_out).cuda().half()
            optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
            optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
            ...
            checkpoint = torch.load("saved.pth")
            model.load_state_dict(checkpoint['model'])
            optimizer.load_state_dict(checkpoint['optimizer'])
        """
        # I think it should actually be ok to reload the optimizer before the model.
        self.loss_scaler = state_dict['loss_scaler']
        self.dynamic_loss_scale = state_dict['dynamic_loss_scale']
        self.overflow = state_dict['overflow']
        self.first_closure_call_this_step = state_dict['first_closure_call_this_step']
        self.optimizer.load_state_dict(state_dict['optimizer_state_dict'])
        # At this point, the optimizer's references to the model's fp32 parameters are up to date.
        # The optimizer's hyperparameters and internal buffers are also up to date.
        # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still
        # out of date.  There are two options.
        # 1:  Refresh the master params from the model's fp16 params.
        # This requires less storage but incurs precision loss.
        # 2:  Save and restore the fp32 master copies separately.
        # We choose option 2.
        #
        # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device
        # of their associated parameters, because it's possible those buffers might not exist yet in
        # the current optimizer instance.  In our case, as long as the current FP16_Optimizer has been
        # constructed in the same way as the one whose state_dict we are loading, the same master params
        # are guaranteed to exist, so we can just copy_() from the saved master params.
        for current_group, saved_group in zip(self.fp32_from_fp16_groups, state_dict['fp32_from_fp16']):
            for current, saved in zip(current_group, saved_group):
                current.data.copy_(saved.data)

    def step(self, closure=None): # could add clip option.
        """
        If no closure is supplied, :attr:`step` should be called after
        ``fp16_optimizer_obj.backward(loss)``.
        :attr:`step` updates the fp32 master copy of parameters using the optimizer supplied to
        :class:`FP16_Optimizer`'s constructor, then copies the updated fp32 params into the fp16 params
        originally referenced by :class:`FP16_Optimizer`'s constructor, so the user may immediately run
        another forward pass using their model.

        If a closure is supplied, :attr:`step` may be called without a prior call to
        :attr:`backward(loss)`.
        This control flow is identical to `ordinary Pytorch optimizer use`_ with closures.
        However, the user should take care that any ``loss.backward()`` call within the closure
        has been replaced by ``fp16_optimizer_obj.backward(loss)``.

        Args:
           closure (optional):  Closure that will be supplied to the underlying optimizer originally passed to :class:`FP16_Optimizer`'s constructor.  closure should call :attr:`zero_grad()` on the :class:`FP16_Optimizer` object, compute the loss, call :attr:`backward(loss)`, and return the loss.

        Example with closure::

            # optimizer is assumed to be an FP16_Optimizer object, previously constructed from an
            # existing pytorch optimizer.
            for input, target in dataset:
                def closure():
                    optimizer.zero_grad()
                    output = model(input)
                    loss = loss_fn(output, target)
                    # loss.backward() becomes:
                    optimizer.backward(loss)
                    return loss
                optimizer.step(closure)

        .. warning::
            Currently, calling :attr:`step` with a closure is not compatible with dynamic loss scaling.

        .. _`ordinary Pytorch optimizer use`:
            http://pytorch.org/docs/master/optim.html#optimizer-step-closure
        """

        scale = self.loss_scaler.loss_scale
        self._update_scale(self.overflow)

        if self.overflow:
            """
            print("OVERFLOW! Skipping step. Attempted loss scale: {}, reducing to {}"
                .format(scale, self.loss_scale))
            """
            return

        if closure is not None:
            retval = self._step_with_closure(closure)
        else:
            retval = self.optimizer.step()

        self._master_params_to_model_params()

        return retval

    def _step_with_closure(self, closure):
        def wrapped_closure():
            # helpful for debugging
            # print("Calling wrapped_closure, first_closure_call_this_step = {}"
            #       .format(self.first_closure_call_this_step))
            if self.first_closure_call_this_step:
                # We expect that the fp16 params are initially fresh on entering self.step(),
                # so _master_params_to_model_params() is unnecessary the first time wrapped_closure()
                # is called within self.optimizer.step().
                self.first_closure_call_this_step = False
            else:
                # If self.optimizer.step() internally calls wrapped_closure more than once,
                # it may update the fp32 params after each call.  However, self.optimizer
                # doesn't know about the fp16 params at all.  If the fp32 params get updated,
                # we can't rely on self.optimizer to refresh the fp16 params.  We need
                # to handle that manually:
                self._master_params_to_model_params()
            # Our API expects the user to give us ownership of the backward() call by
            # replacing all calls to loss.backward() with optimizer.backward(loss).
            # This requirement holds whether or not the call to backward() is made within a closure.
            # If the user is properly calling optimizer.backward(loss) within "closure,"
            # calling closure() here will give the fp32 master params fresh gradients
            # for the optimizer to play with, so all wrapped_closure needs to do is call
            # closure() and return the loss.
            temp_loss = closure()
            while(self.overflow):
                scale = self.loss_scaler.loss_scale
                self._update_scale(self.overflow)
                """
                print("OVERFLOW within closure! Skipping step. Attempted loss scale: {}, "
                      "reducing to {}".format(scale, self.loss_scale))
                """
                temp_loss = closure()
            return temp_loss

        retval = self.optimizer.step(wrapped_closure)

        self.first_closure_call_this_step = True

        return retval

    def backward(self, loss, update_master_grads=True):
        """
        :attr:`backward` performs the following conceptual steps:

        1. fp32_loss = loss.float() (see first Note below)
        2. scaled_loss = fp32_loss*loss_scale
        3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's leaves (which may be fp16, fp32, or a mixture, depending how your model was defined).
        4. fp16 grads are then copied to the master params' ``.grad`` attributes (see second Note), which are guaranteed to be fp32.
        5. Finally, master grads are divided by loss_scale.

        In this way, after :attr:`backward`, the master params have fresh gradients,
        and :attr:`step` may be called.

        .. note::
            :attr:`backward` internally converts the loss to fp32 before applying the loss scale.
            This provides some additional safety against overflow if the user has supplied an
            fp16 loss value.
            However, for maximum overflow safety, the user should
            compute the loss criterion (MSE, cross entropy, etc) in fp32 before supplying it to
            :attr:`backward`.

        .. warning::
            The gradients found in a model's leaves after the call to
            :attr:`backward` should not be regarded as valid in general,
            because it's possible
            they have been scaled (and in the case of dynamic loss scaling,
            the scale factor may change over time).
            If the user wants to inspect gradients after a call to :attr:`backward`,
            only the master gradients should be regarded as valid.  These can be retrieved via
            :attr:`inspect_master_grad_data()`.

        Args:
            loss:  The loss output by the user's model.  loss may be either float or half (but see first Note above).
            update_master_grads (bool, optional, default=True):  Option to copy fp16 grads to fp32 grads on this call.  By setting this to False, the user can delay the copy, which is useful to eliminate redundant fp16->fp32 grad copies if :attr:`backward` is being called on multiple losses in one iteration.  If set to False, the user becomes responsible for calling :attr:`update_master_grads` before calling :attr:`step`.

        Example::

            # Ordinary operation:
            optimizer.backward(loss)

            # Naive operation with multiple losses (technically valid, but less efficient):
            # fp32 grads will be correct after the second call,  but
            # the first call incurs an unnecessary fp16->fp32 grad copy.
            optimizer.backward(loss1)
            optimizer.backward(loss2)

            # More efficient way to handle multiple losses:
            # The fp16->fp32 grad copy is delayed until fp16 grads from all
            # losses have been accumulated.
            optimizer.backward(loss1, update_master_grads=False)
            optimizer.backward(loss2, update_master_grads=False)
            optimizer.update_master_grads()
        """
        # To consider:  try multiple backward passes using retain_grad=True to find
        # a loss scale that works.  After you find a loss scale that works, do a final dummy
        # backward pass with retain_graph=False to tear down the graph.  Doing this would avoid
        # discarding the iteration,  but probably wouldn't improve overall efficiency.
        self.loss_scaler.backward(loss.float())
        if update_master_grads:
            self.update_master_grads()

    def update_master_grads(self):
        """
        Copy the ``.grad`` attribute from stored references to fp16 parameters to
        the ``.grad`` attribute of the fp32 master parameters that are directly
        updated by the optimizer.  :attr:`update_master_grads` only needs to be called if
        ``fp16_optimizer_obj.backward`` was called with ``update_master_grads=False``.
        """
        if self.dynamic_loss_scale:
            self._check_overflow()
            if self.overflow: return
        self._model_grads_to_master_grads()
        self._downscale_master()

    def inspect_master_grad_data(self):
        """
        When running with :class:`FP16_Optimizer`,
        ``.grad`` attributes of a model's fp16 leaves should not be
        regarded as truthful, because they might be scaled.
        After a call to :attr:`fp16_optimizer_obj.backward(loss)`, if no overflow was encountered,
        the fp32 master params' ``.grad``
        attributes will contain valid gradients properly divided by the loss scale.  However,
        because :class:`FP16_Optimizer` flattens some parameters, accessing them may be
        nonintuitive.  :attr:`inspect_master_grad_data`
        allows those gradients to be viewed with shapes corresponding to their associated model leaves.

        Returns:
            List of lists (one list for each parameter group).  The list for each parameter group
            is a list of the ``.grad.data`` attributes of the fp32 master params belonging to that group.
        """
        raise NotImplementedError("Currently not implemented, working on it...")
        fp32_grads_each_group = []
        if self.overflow:
            print("Warning:  calling FP16_Optimizer.inspect_master_grad_data while in an overflow state.  "
                  "Gradients are currently invalid (may be inf, nan, or stale).  Returning None.")
            return None
        else:
            return None

    # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale"
    def _get_loss_scale(self):
        return self.loss_scaler.loss_scale

    def _set_loss_scale(self, value):
        self.loss_scaler.cur_scale = value

    loss_scale = property(_get_loss_scale, _set_loss_scale)

    # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state"
    def _get_state(self):
        return self.optimizer.state

    def _set_state(self, value):
        self.optimizer.state = value

    state = property(_get_state, _set_state)

    # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups"
    # (for example, to adjust the learning rate)
    def _get_param_groups(self):
        return self.optimizer.param_groups

    def _set_param_groups(self, value):
        self.optimizer.param_groups = value

    param_groups = property(_get_param_groups, _set_param_groups)


#!/usr/bin/sudo python

from mininet.topo import Topo
from mininet.net import Mininet
from mininet.log import setLogLevel
from mininet.node import OVSKernelSwitch, RemoteController, CPULimitedHost
from mininet.cli import CLI
from mininet.link import TCLink
import os
import threading
import numpy as np
from random import randint, choice
from time import sleep

from intent_manager import IntentManager
from link_manager import LinkManager
from settings import *
import subprocess
import socket
import re
from switch_movement import SwitchMover
from media_server import MediaServer
from energy_management import BatteryCapacityLogger
from switch_operations import SwitchOperator
import traceback
import time
import json
import shutil
import pickle
from process_cleaner import stop_children_processes

switch_mover = None
media_server = None
media_server_info = {}
drone_battery_manager = None
link_manager = None
intent_manager = None
processes = []
link_options_host2switch = dict(bw=54 / DATA_SCALING, delay='100us', loss=0.01, max_queue_size=1000, use_htb=True)
link_options_switch2switch = dict(bw=100 / DATA_SCALING, delay='100us', loss=0.01, max_queue_size=1000, use_htb=True)
# linkOptionsMediaServer2Switch = dict(bw=100, penalty='0ms', loss=0.01, max_queue_size=100000, use_htb=True)
link_options_switch_dummy_host2switch = dict(bw=1000 / DATA_SCALING, delay='0', loss=0, max_queue_size=1000,
                                             use_htb=True)
listen_host_availability_thread = None
request_file_thread = None
switch_mobility_thread = None
busy_hosts = []
switch_name_id_table = {}
switch_id_name_table = {}
links_h2s = []
links_s2s = {}
current_links = []
current_time = 0
requests_from_depleted_drones = 0
simulation_id = str(int(time.time() // 10))

simulation_snapshot = {"settings": {}}

with open('settings.py', 'r') as f:
    exec (f.read(), simulation_snapshot["settings"])
if '__builtins__' in simulation_snapshot["settings"]:
    del simulation_snapshot["settings"]['__builtins__']

def debug(msg):
    if PRINT_DUMMY_MESSAGES:
        print(msg)

def print_link_changes(message):
    if PRINT_LINK_CHANGES:
        print(message)

class InbandController(RemoteController):

    def checkListening(self):
        if self.port is None:
            self.port = 6633
        return

class DroneTopology(Topo):
    def build(self, s=NUMBER_OF_SWITCHES, n=NUMBER_OF_HOSTS_PER_SWITCH):
        global switch_name_id_table
        switches = []
        host_number = 1
        for sw in range(s):
            new_switch = self.addSwitch('s%d' % (sw + 1))
            cache_host_in_switch = self.addHost('hs%d' % (sw + 1),
                                                mac='00:00:00:00:%s:00' % hex(sw + 1).split('x')[-1].zfill(2),
                                                ip='10.0.%d.1/16' % (sw + 1),
                                                cpu=.5 / (s * n)
                                                )
            self.addLink(cache_host_in_switch, new_switch, **link_options_switch_dummy_host2switch)
            switch_name_id_table['s%d' % (sw + 1)] = sw
            switch_id_name_table[sw] = 's%d' % (sw + 1)
            for h in range(n):
                host = self.addHost('h%d' % host_number,
                                    cpu=.30 / (s * n),
                                    mac='00:00:00:00:00:' + hex(host_number).split('x')[-1].zfill(2))
                host_number += 1
                self.addLink(host, new_switch, bw=switch_mover.switch2host_link_capacities[sw] / DATA_SCALING,
                             delay='0',
                             loss=0,
                             max_queue_size=10000, use_htb=True)
            switches.append(new_switch)

        for link in switch_mover.active_link_pair_capacity_table:
            # pair = "{}-{}".format(switch_id_name_table[link[0]], switch_id_name_table[link[1]])
            self.addLink(switch_id_name_table[link[0]], switch_id_name_table[link[1]], bw=100 / DATA_SCALING,
                         delay='0ms',
                         loss=0,
                         max_queue_size=10000, use_htb=True)
        # serverNode = self.addHost('msrv', mac='00:00:00:00:00:FF', ip='10.0.200.200/8')
        # self.addLink(serverNode, switch_id_name_table[switch_mover.content_server_switch_id],
        # **linkOptionsMediaServer2Switch)

def wait_exponentially(scale_of_exponential):
    time_until_next_request = np.random.exponential(scale_of_exponential)
    sleep(time_until_next_request)

def request_file(scale_of_exponential=REQUEST_INTERVAL, host_number=NUMBER_OF_HOSTS, alpha_zipf_exponent=1.1,
                 number_of_files=NUMBER_OF_FILES):
    global requests_from_depleted_drones
    s = np.random.zipf(alpha_zipf_exponent, number_of_files * 10)
    random_file_selection_list = s[s < number_of_files + 1] - 1
    t = threading.currentThread()
    flow_id = 0
    while getattr(t, "is_running", True):
        # start = time.time()
        requesting_host_id = randint(1, host_number)
        if requesting_host_id - 1 in drone_battery_manager.depleted_battery_ids:
            requests_from_depleted_drones += 1
            wait_exponentially(scale_of_exponential)
            continue
        requested_file_id = str(choice(random_file_selection_list))
        # File request protocol message: requesting_host_ip;requested_file_id

        try:
            host = net.get("h%d" % requesting_host_id)
            host_ip = host.IP()

            if DISABLE_MULTIPLE_REQUESTS_FROM_SAME_HOST and host_ip in busy_hosts:
                debug("h%d already requested a file" % requesting_host_id)
                continue
            else:
                debug("%d: h%s requests file %s" % (flow_id, requesting_host_id, requested_file_id))
            if DISABLE_MULTIPLE_REQUESTS_FROM_SAME_HOST and GENERATE_TRAFFIC:
                busy_hosts.append(host_ip)
            request_message = "{};{}".format(host_ip, requested_file_id)
            media_source_node_ip, redirect_message = media_server.process_file_request(request_message,
                                                                                media_server_info["ip"])
            if media_source_node_ip in [SOURCE_OWN_CACHE, SOURCE_NOT_AVAILABLE]:
                if DISABLE_MULTIPLE_REQUESTS_FROM_SAME_HOST:
                    busy_hosts.remove(host_ip)
                wait_exponentially(scale_of_exponential)
                continue
            # debug(request_message)

            # host.cmd("echo '{};{}' | nc {} {} -w 0.5 &> /dev/null &".format(redirect_message, flow_id,
            # media_source_node_ip, str(FILE_REQUEST_LISTEN_PORT)))
            if ONOS_INTENT_ADD_MODE:
                intent_manager.create_host_2_host_intent(host_ip, media_source_node_ip)
            host.cmd("echo '{};{}' >/dev/tcp/{}/{} &".format(redirect_message, flow_id, media_source_node_ip,
                                                             str(FILE_REQUEST_LISTEN_PORT)))

            flow_id += 1
            # debug(result)
            # I assume request rate is Poisson Distribution so the time intervals are exponential distributions
            # end = time.time()
            # print(end - start)
            wait_exponentially(scale_of_exponential)
        except Exception as e:
            print(traceback.format_exc())
            print(e)
            print(media_server_info)

# def moveHosts(host_number=NUMBER_OF_HOSTS, switch_number=NUMBER_OF_SWITCHES, interval=MOVE_EVENT_INTERVALS):
#     while True:
#         try:
#             movingHostID = randint(1, host_number)
#             newSwitchID = randint(1, switch_number)
#             host = net.get("h%d" % movingHostID)
#             switch = net.get("s%d" % newSwitchID)
#             host.deleteIntfs()
#             print ("Host %d hand over to %d " % (movingHostID, newSwitchID))
#             sleep(interval / 2)
#             net.addLink(host, switch, **link_options_host2switch)
#             sleep(interval / 2)
#         except Exception as e:
#             print(traceback.format_exc())
#             print (e)

def update_links_h2s():
    for switchID, new_link_capacity in enumerate(switch_mover.switch2host_link_capacities):
        previous_link_capacity = links_h2s[switchID][1]
        if previous_link_capacity != new_link_capacity:
            interface = links_h2s[switchID][0]
            interface.config(bw=new_link_capacity / DATA_SCALING, smooth_change=True)
            links_h2s[switchID][1] = new_link_capacity
            print_link_changes(
                "H2S {} capacity has changed from {} to {}".format(switchID, previous_link_capacity, new_link_capacity))

def update_links_s2s():
    new_link_capacities_s2s = {}
    new_links = []
    for link in switch_mover.active_link_pair_capacity_table:
        sw1 = switch_id_name_table[link[0]]
        sw2 = switch_id_name_table[link[1]]
        new_capacity = link[2]
        pair = "{}-{}".format(sw1, sw2)
        if sw1 == sw2:
            print("Link to itself: %s" % sw1)
            print("Operating switches:")
            print(switch_mover.operating_switches)
            print("Neighbors:")
            print(switch_mover.neighbors)
            print("Active link table:")
            print(switch_mover.active_link_pair_capacity_table)
            continue
        new_link_capacities_s2s[pair] = new_capacity
        new_links.append(pair)
        # Check link exists
        if pair in links_s2s:
            if not links_s2s[pair][3]:
                # Link was down, make link up
                interface1 = links_s2s[pair][1][0]
                interface2 = links_s2s[pair][1][1]
                if links_s2s[pair][2] != new_capacity:
                    # Apply the new bw
                    interface1.config(bw=new_capacity / DATA_SCALING, smooth_change=True)
                    print_link_changes(
                        "{} capacity has changed from {} to {}".format(pair, links_s2s[pair][2], new_capacity))
                net.get(sw1).attach(interface1)
                print_link_changes("{} is up".format(pair))
                if ONOS_LINK_ADD_MODE:
                    link_manager.send_two_way_link_to_controller(interface1.sumo_id, interface2.sumo_id)
        else:
            # Create the link
            try:
                net.addLink(sw1, sw2, bw=100 / DATA_SCALING)
            except Exception as e:
                print(traceback.format_exc())
                print (e)

            interfaces = net.get(sw1).connectionsTo(net.get(sw2))[0]
            interface1 = interfaces[0]
            interface2 = interfaces[1]
            net.get(sw1).attach(interface1)
            net.get(sw2).attach(interface2)
            links_s2s[pair] = [sw1, interfaces, new_capacity, True]
            print_link_changes("{} is created.".format(pair))
            if ONOS_LINK_ADD_MODE:
                link_manager.send_two_way_link_to_controller(interface1.sumo_id, interface2.sumo_id)

    for previously_active_link in current_links:
        # From the previous time slot, check the link changes one by one
        if previously_active_link in new_link_capacities_s2s:
            # Link continues to operate
            previous_capacity = links_s2s[previously_active_link][2]
            new_capacity = new_link_capacities_s2s[previously_active_link]
            if previous_capacity != new_capacity:
                # Link capacity has changed.
                interface1 = links_s2s[previously_active_link][1][0]
                interface1.config(bw=new_capacity / DATA_SCALING, smooth_change=True)
                print_link_changes(
                    "{} capacity has changed from {} to {}".format(previously_active_link, previous_capacity,
                                                                   new_capacity))
            else:
                print_link_changes("{}-> Alles gut.".format(previously_active_link))
        else:
            # Link is down
            links_s2s[previously_active_link][3] = False
            interfaces = links_s2s[previously_active_link][1]
            interface1 = interfaces[0]
            interface2 = interfaces[1]
            if ONOS_LINK_ADD_MODE:
                link_manager.delete_two_way_link_from_controller(interface1.sumo_id, interface2.sumo_id)
            net.get(links_s2s[previously_active_link][0]).detach(interface1)
            print_link_changes("{} is down".format(previously_active_link))

    # Save the new link pairs
    current_links[:] = new_links

def update_links():
    update_links_h2s()
    update_links_s2s()

def remove_depleted_switches(depleted_batteries):
    switch_mover.number_of_switches_alive -= len(depleted_batteries)
    switch_mover.operating_switches = np.setdiff1d(switch_mover.operating_switches, depleted_batteries)
    for switch in depleted_batteries:
        net.get(switch_id_name_table[switch]).stop()
    if switch_mover.number_of_switches_alive <= 0:
        terminate_simulation()
        return
    if switch_mover.content_server_switch_id in depleted_batteries:
        print("Content server switch is depleted")
        global media_server_info
        switch_mover.get_central_node()
        media_server_dummy_host_name = "h" + switch_id_name_table[switch_mover.content_server_switch_id]
        media_server_info = {"node": net.get(media_server_dummy_host_name),
                             "ip": net.get(media_server_dummy_host_name).IP(),
                             "id": switch_mover.content_server_switch_id}

def save_simulation_snapshot_to_file(data_time, file_name=None):
    if file_name is None:
        file_name = "resultsV2/situation-%s.pkl" % simulation_id
    simulation_snapshot[data_time] = {"links": switch_mover.active_link_pair_capacity_table[:],
                                      "content_server": switch_mover.content_server_switch_id,
                                      "battery_levels": np.copy(drone_battery_manager.battery_levels),
                                      "media_server_logs": media_server.get_caching_statistics(),
                                      "requests_from_depleted_drones": requests_from_depleted_drones,
                                      "caching_switches": switch_operator.caching_switches[:],
                                      "cached_switches": switch_operator.cached_switches[:],
                                      "depleted_battery_ids": np.copy(drone_battery_manager.depleted_battery_ids),
                                      "energy_consumption_statistics": np.copy(
                                          drone_battery_manager.energy_consumption_statistics),
                                      "energy_cost_matrix": media_server.energy_cost_matrix[:],
                                      "flow_hop_matrix": media_server.flow_hop_matrix[:],
                                      }
    with open(file_name, 'wb+') as handle:
        pickle.dump(simulation_snapshot, handle, protocol=pickle.HIGHEST_PROTOCOL)

def move_switches(interval=MOVE_EVENT_INTERVALS):
    global current_time
    t = threading.currentThread()
    print("Time is %d." % (current_time / 60))
    # try:
    #     os.remove("batteryLevelsSpecific.csv")
    # except:
    #     traceback
    # drone_battery_manager.printProportionalBatteryLevelsToFile("batteryLevelsSpecific.csv", (current_time / 60))
    # switch_mover.printLinksToFile("links/links_%d" % (current_time / 60))
    save_simulation_snapshot_to_file(current_time)
    while switch_mover.number_of_switches_alive and getattr(t, "is_running", True):
        sleep(interval)
        current_time += interval
        drone_battery_manager.decrease_battery_fly()
        switch_operator.get_switch_overheads()
        depleted_batteries = drone_battery_manager.check_for_depleted_batteries()
        if depleted_batteries:
            remove_depleted_switches(depleted_batteries)
        switch_mover.move_switches_for_one_time_interval(current_time)
        switch_operator.update_device_coordinates_on_controller()
        update_links()
        if current_time % 60 == 0:
            "Caching host placement time"
            switch_operator.select_caching_nodes()
            # switch_mover.printLinksToFile("links/links_%d" % (current_time / 60))
            # drone_battery_manager.printProportionalBatteryLevelsToFile("batteryLevelsSpecific.csv",(current_time/60))
            # drone_battery_manager.printEnergyConsumptionStatisticsToFile(current_time / 60)
            save_simulation_snapshot_to_file(current_time)
            print(drone_battery_manager.battery_levels)
            print("Time is %d." % (current_time / 60))

def get_host_availability(accepted_socket):
    message = accepted_socket.recv(1024).decode()
    accepted_socket.close()
    print(message)
    if ONOS_INTENT_ADD_MODE:
        message_parsed = message.split(";", 3)
        src_ip = message_parsed[0]
        dst_ip = message_parsed[1]
        intent_manager.delete_host_2_host_intent(src_ip, dst_ip)
    if DISABLE_MULTIPLE_REQUESTS_FROM_SAME_HOST:

        try:
            busy_hosts.remove(dst_ip)
        except Exception as e:
            print(traceback.format_exc())
            print (e)

def listen_host_availability():
    t = threading.currentThread()
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((MININET_IP, HOST_AVAILABILITY_LISTEN_PORT))
    s.listen(5)
    while getattr(t, "is_running", True):
        # print("...ABC...")
        accepted_socket, address = s.accept()
        try:
            if not getattr(t, "is_running", True):
                return
            get_host_availability(accepted_socket)
        except Exception as e:
            print(traceback.format_exc())
            print(e)

def get_initial_link_objects():
    global media_server_info
    print(switch_mover.active_link_pair_capacity_table)
    media_server_dummy_host_name = "h" + switch_id_name_table[switch_mover.content_server_switch_id]
    media_server_info = {"node": net.get(media_server_dummy_host_name),
                         "ip": net.get(media_server_dummy_host_name).IP(),
                         "id": switch_mover.content_server_switch_id}
    drone_battery_manager.add_extra_battery(switch_mover.content_server_switch_id)
    for switchID, linkCapacity in enumerate(switch_mover.switch2host_link_capacities):
        switch_name = switch_id_name_table[switchID]
        interface = net.get(switch_name).intf()
        links_h2s.append([interface, linkCapacity])
    for link in switch_mover.active_link_pair_capacity_table:
        sw1 = switch_id_name_table[link[0]]
        sw2 = switch_id_name_table[link[1]]
        capacity = link[2]
        interfaces = net.get(sw1).connectionsTo(net.get(sw2))[0]
        interface1 = interfaces[0]
        interface2 = interfaces[1]
        # Initialize link with high value then set one side real value which is bottleneck for the whole link
        interface1.config(bw=capacity / DATA_SCALING, smooth_change=True)
        pair = "{}-{}".format(sw1, sw2)
        links_s2s[pair] = [sw1, interfaces, capacity, True]
        current_links.append(pair)
        if ONOS_LINK_ADD_MODE:
            link_manager.send_two_way_link_to_controller(interface1.sumo_id, interface2.sumo_id)

def initialize_network():
    # Tell mininet to print information
    setLogLevel(MININET_LOG_LEVEL)
    "Create drone network"
    topology = DroneTopology()
    mininet_network = Mininet(topo=topology, controller=RemoteController('main_controller', ip='127.0.0.1', port=6653),
                              switch=OVSKernelSwitch, host=CPULimitedHost, link=TCLink)
    return mininet_network

def start_network():
    global request_file_thread, listen_host_availability_thread, switch_mobility_thread
    # c1 = network.addController('c0', controller=RemoteController, ip=CONTROLLER_IP)
    # Add internet connection
    net.addNAT().configDefault()
    net.start()

    listen_host_availability_thread = threading.Thread(target=listen_host_availability)
    listen_host_availability_thread.setDaemon(True)
    listen_host_availability_thread.start()

    # CLI(net)
    get_initial_link_objects()
    CLI(net)
    # mediaServerIP = network.get("msrv").IP()
    # debug("mediaServerIP %s" % mediaServerIP)
    for h in net.hosts:
        h.popen("ping -c 1 10.0.0.1")
        if re.match(r"h[0-9]+", h.name):
            # itg_rec_signal_port = int(h.name[1:]) + 9000
            # processes.append(h.popen('ITGRecv -Sp %d -a %s' % (itg_rec_signal_port, h.IP())))
            if GENERATE_TRAFFIC:
                processes.append(h.popen('python streamer_host.py %s %s' % (h.IP(), simulation_id)))
        elif re.match(r"hs[0-9]+", h.name):
            if GENERATE_TRAFFIC:
                processes.append(h.popen('python streamer_host.py %s %s' % (h.IP(), simulation_id)))
    CLI(net)
    # hostMobilityThread = threading.Thread(target=moveHosts, kwargs={"host_number": NUMBER_OF_HOSTS,
    #                                                                 "switch_number": NUMBER_OF_SWITCHES,
    #                                                                 "interval": MOVE_EVENT_INTERVALS})
    # hostMobilityThread.setDaemon(True)
    # hostMobilityThread.start()
    switch_mobility_thread = threading.Thread(target=move_switches)
    switch_mobility_thread.setDaemon(True)
    switch_mobility_thread.start()

    request_file_thread = threading.Thread(target=request_file,
                                           kwargs={"scale_of_exponential": REQUEST_INTERVAL,
                                                   "host_number": NUMBER_OF_HOSTS,
                                                   "alpha_zipf_exponent": ALPHA_ZIPF,
                                                   "number_of_files": NUMBER_OF_FILES})
    request_file_thread.setDaemon(True)
    request_file_thread.start()

    debug(links_s2s)
    debug(links_h2s)

def copytree(src, dst, symlinks=False, ignore=None):
    for item in os.listdir(src):
        s = os.path.join(src, item)
        d = os.path.join(dst, item)
        if os.path.isdir(s):
            shutil.copytree(s, d, symlinks, ignore)
        else:
            shutil.copy2(s, d)

def terminate_simulation():
    print("We are done!!")
    save_simulation_snapshot_to_file(current_time)
    report = media_server.get_caching_statistics()
    report["requests_from_depleted_drones"] = requests_from_depleted_drones
    report["current_time"] = current_time
    print(report)
    results_path = "results/%s" % simulation_id
    result_file = os.path.join(results_path, "result.json")
    if not os.path.isfile(result_file):
        os.makedirs(results_path)
        with open(result_file, "w+") as outfile:
            json.dump(report, outfile)
        subprocess.call(['chmod', '-R', '777', 'results'])
    try:
        listen_host_availability_thread.is_running = False
    except Exception as e:
        print(traceback.format_exc())
        print(e)
    try:
        request_file_thread.is_running = False
    except Exception as e:
        print(traceback.format_exc())
        print(e)
    try:
        switch_mobility_thread.is_running = False
    except Exception as e:
        print(traceback.format_exc())
        print(e)

if __name__ == '__main__':
    stop_children_processes(processes)

    link_manager = LinkManager()
    intent_manager = IntentManager()
    switch_mover = SwitchMover()
    drone_battery_manager = BatteryCapacityLogger()
    switch_operator = SwitchOperator(switch_mover, drone_battery_manager)
    media_server = MediaServer(switch_operator)

    net = initialize_network()
    start_network()

    CLI(net)
    terminate_simulation()

    stop_children_processes(processes)
    net.stop()


numbers = [0, 0, 0]
numbers[0] = float(input("1º número: "))
numbers[1] = float(input("2º número: "))
numbers[2] = float(input("3º número: "))

sortedlist = sorted(numbers)

print(f"menor: {sortedlist[0]}, maior: {sortedlist[2]}")

numbers = [0, 0, 0]
numbers[0] = float(input("1º número: "))
numbers[1] = float(input("2º número: "))
numbers[2] = float(input("3º número: "))

acc_min = numbers[0]
acc_max = numbers[0]

for n in numbers:
    if n < acc_min:
        acc_min = n

    if n > acc_max:
        acc_max = n

print(f"Maior: {acc_max}\nMenor: {acc_min}")

numbers = [0, 0, 0]
numbers[0] = float(input("1º número: "))
numbers[1] = float(input("2º número: "))
numbers[2] = float(input("3º número: "))

min = numbers[0]
max = numbers[0]

for n in numbers:
    if n < min:
        min = n

    if n > max:
        max = n

print(f"Maior: {max}\nMenor: {min}")

numbers = [0, 0, 0]
numbers[0] = float(input("1º número: "))
numbers[1] = float(input("2º número: "))
numbers[2] = float(input("3º número: "))

print(f"Maior: {max(numbers)}\nMenor: {min(numbers)}")