#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Oct  2 11:08:09 2020

@author: alvarezguido
GITHUB: https://github.com/alvarezguido
"""

"""
SYNOPSIS
----
----
-----

"""

import simpy
import random
import numpy as np
import math
#import sys
#import re
#import matplotlib.pyplot as plt
#import os
#import operator
#from mpl_toolkits import mplot3d
#from mpl_toolkits.mplot3d import Axes3D
#import PIL
import random
import re
import os
import datetime
import sys

name = "ER"
mode_debbug = 0

if not mode_debbug:
    null = open(os.devnull, 'w')
    old_stdout = sys.stdout
    sys.stdout = null
####WE START BY USING SF=12 ADN BW=125 AND CR=1, FOR ALL NODES AND ALL TRANSMISIONS######

if mode_debbug:
    RANDOM_SEED = 5
    chan = 1
    packetlen = 20
    total_data = 60
    beacon_time = 120
    maxBSReceives = 500
    multi_nodes = [1500]
else:
    RANDOM_SEED = int(sys.argv[1])
    chan = int(sys.argv[2])
    packetlen = int(sys.argv[3])   ##NODES SEND PACKETS OF JUST 20 Bytes
    total_data = int(sys.argv[4]) ##TOTAL DATA ON BUFFER, FOR EACH NODE (IT'S THE BUFFER O DATA BEFORE START SENDING)
    beacon_time = int(sys.argv[5]) ###SAT SENDS BEACON EVERY CERTAIN TIME
    maxBSReceives = int(sys.argv[6]) ##MAX NUMBER OF PACKETS THAT BS (ie SATELLITE) CAN RECEIVE AT SAME TIME

    #multi_nodes = [int(sys.argv[7]), int(sys.argv[8]) ,int(sys.argv[9])]
    multi_nodes = [int(sys.argv[7]), int(sys.argv[8]) ,int(sys.argv[9]), int(sys.argv[10]),int(sys.argv[11]),int(sys.argv[12]),int(sys.argv[13]),int(sys.argv[14]),int(sys.argv[15]),int(sys.argv[16]),int(sys.argv[17]),int(sys.argv[18]),int(sys.argv[19]),int(sys.argv[20])]
random.seed(RANDOM_SEED) #RANDOM SEED IS FOR GENERATE ALWAYS THE SAME RANDOM NUMBERS (ie SAME RESULTS OF SIMULATION)
nodesToSend = []
packetsToSend = math.ceil(total_data/packetlen)
###GLOBAL PARAMS ####
bsId = 1 ##ID OF BASE STATION (NOT USED)

avgSendTime = 3  ## NOT USED! --> A NODE SENDS A PACKET EVERY X SECS

back_off = beacon_time * 0.95 ###BACK OFF TIME FOR SEND A PACKET
packetsAtBS = [] ##USED FOR CHEK IF THERE ARE ALREADY PACKETS ON THE SATELLITE
c = 299792.458 ###SPEED LIGHT [km/s]
Ptx = 14
G_device = 0; ##ANTENNA GAIN FOR AN END-DEVICE
G_sat = 12;   ##ANTENNA GAIN FOR SATELLITE
nodes = [] ###EACH NODE WILL BE APPENDED TO THIS VARIABLE
freq =868e6 ##USED FOR PATH LOSS CALCULATION
frequency = [868100000, 868300000, 868500000] ##FROM LORAWAN REGIONAL PARAMETERS EU863-870 / EU868

nrLost = 0 ### TOTAL OF LOST PACKETS DUE Lpl
nrCollisions = 0 ##TOTAL OF COLLIDED PACKETS
nrProcessed = 0 ##TOTAL OF PROCESSED PACKETS
nrReceived = 0 ###TOTAL OF RECEIVED PACKETS
nrNoProcessed = 0 ##TOTAL OF INTRA-PACKETS NO PROCESSED
nrIntraTot = 0
nrLostMaxRec = 0
nrCollFullPacket = 0
nrSentIntra = 0 ##TOTAL OF SENT INTRA-PACKTES
nrReceivedIntra = 0 ##TOTAL OF RECEIVED INTRA-PACKETS

##ARRAY WITH MEASURED VALUES FOR SENSIBILITY, NEW VALUES
##THE FOLLOWING VALUES CORRESPOND TO:
#   - FIRST ELEMENT: IT'S THE SF (NOT USABLE)
#   - SECOND ELEMENT: SENSIBILITY FOR 125KHZ BW
#   - THIRD ELEMENT: SENSIBILITY FOR 250KHZ BW
#   - FOURTH ELEMENT: SENSIBILITY FOR 500KHZ BW
# NOTICE THAT SENSIBILITY DECREASE ALONG BW INCREASES, ALSO WITH LOWER SF
# THIS VALUES RESPONDS TO:
# wf = -174 + 10 log(BW) +NF +SNRf
sf7 = np.array([7,-123,-120,-117.0])
sf8 = np.array([8,-126,-123,-120.0])
sf9 = np.array([9,-129,-126,-123.0])
sf10 = np.array([10,-132,-129,-126.0])
sf11 = np.array([11,-134.53,-131.52,-128.51])
sf12 = np.array([12,-137,-134,-131.0])

sensi = np.array([sf7,sf8,sf9,sf10,sf11,sf12])

#DR = ["dr8","dr9","dr10","dr11"]
DR = [-137,-134.5,-134,-131.5]

## READ PARAMS FROM DIRECTORY ##
path = "./wider_scenario_2/"

### -137dB IS THE MINIMUN TOLERABLE SENSIBILITY, FOR SF=12 AND BW=125KHz ###

leo_pos=np.loadtxt( path + "LEO-XYZ-Pos.csv",skiprows=1,delimiter=',',usecols=(1,2,3))
## WHERE:
    ## leo_pos[i,j]:
        ## i --> the step time in sat pass
        ## j --> 0 for x-position, 1 for y-position, 2 for z-position

sites_pos = np.loadtxt( path + "SITES-XYZ-Pos.csv",skiprows=1,delimiter=',',usecols=(1,2,3))
## WHERE:
    ## sites_pos[i,j]:
        ## i --> the node i
        ## j --> 0 for x-position, 1 for y-position, 2 for z-position

dist_sat = np.zeros((sites_pos.shape[0],3,leo_pos.shape[0]))
t = 0
for i in range(leo_pos.shape[0]):
    t+=1
    dist_sat [:,:,i] = leo_pos[i,:] - sites_pos
## WHERE:
    ## dist_sat[i,j,k]:
        ## i --> the node i
        ## j --> 0 for x-position, 1 for y-position, 2 for z-position
        ## k --> the step time in sat pass

#### FOR COMPUTE DISTANCE MAGNITUDE (ABS) FROM END-DEVICE TO SAT PASSING BY ####
distance = np.zeros((sites_pos.shape[0],leo_pos.shape[0]))
distance[:,:] = (dist_sat[:,0,:]**2 + dist_sat[:,1,:]**2 + dist_sat[:,2,:]**2)**(1/2)
## WHERE:
    ## distance[i,j]:
        ## i --> the node i
        ## j --> the step time in sat pass

##MATRIX FOR LINK BUDGET Lpl ###
Lpl = np.zeros((sites_pos.shape[0],leo_pos.shape[0]))
Lpl = 20*np.log10(distance*1000) + 20*np.log10(freq) - 147.55 #DISTANCE MUST BE IN METERS
## WHERE:
    ## Lpl[i,j]:
        ## i --> the node i
        ## j --> the step time in sat pass

##MATRIX FOR LINK BUDGET, USING Prx ###
Prx = np.zeros((sites_pos.shape[0],leo_pos.shape[0]))
Prx = Ptx + G_sat + G_device -20*np.log10(distance*1000) - 20*np.log10(freq) + 147.55 #DISTANCE IS CONVERTED TO METERS
## WHERE:
    ## Prx[i,j]:
        ## i --> the node i
        ## j --> the step time in sat pass

distance = np.concatenate((distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance,\
                           distance,distance,distance,distance,distance,distance,distance,distance,distance,distance))

Lpl = np.concatenate((Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,\
                      Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl,Lpl))

Prx = np.concatenate((Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,\
                      Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx,Prx))

elev = np.degrees(np.arcsin(599/distance))

# =============================================================================
# IS7 = np.array([1,-8,-9,-9,-9,-9])
# IS8 = np.array([-11,1,-11,-12,-13,-13])
# IS9 = np.array([-15,-13,1,-13,-14,-15])
# IS10 = np.array([-19,-18,-17,1,-17,-18])
# IS11 = np.array([-22,-22,-21,-20,1,-20])
# IS12 = np.array([-25,-25,-25,-24,-23,1])
# IsoThresholds = np.array([IS7,IS8,IS9,IS10,IS11,IS12])
#
# =============================================================================
ISDR8 = np.array([1,-23,-24,-25])
ISDR9 = np.array([-20,1,-20,-21])
ISDR10 = np.array([-18,-17,1,-17])
ISDR11 = np.array([-15,-14,-13,1])

IsoThresholds = np.array([ISDR8,ISDR9,ISDR10,ISDR11])

#THIS IS THE MATRIX OF CROSS INTERFERENCE BETWEEN SF

Collmap = [[0 for i in range(0,6)] for j in range(0,6)]

def simulate_scenario (nrNodes):
    env = simpy.Environment()

    def powerCollision_2(p1, p2):
        #powerThreshold = 6
        global Collmap
        #print ("SF: node {} has {} ; node {} has {}".format(p1.nodeid,p1.sf, p2.nodeid, p2.sf))
        #print ("pwr: node {} with rssi {} dBm and node {} with rssi {} dBm; diff {:3.2f} dBm".format(p1,p1.rssi[math.ceil(env.now)], p2.nodeid,p2.rssi[math.ceil(env.now)], p1.rssi[math.ceil(env.now)] - p2.rssi[math.ceil(env.now)]))
        if True: #p1.sf == p2.sf:
           if abs(p1.rssi[math.ceil(env.now)] - p2.rssi[math.ceil(env.now)]) < IsoThresholds[p1.dr-8][p2.dr-8]:
               print ("collision pwr both node {} and node {}".format(p1.nodeid, p2.nodeid))
               #Collmap[p1.sf-7][p2.sf-7] += 1
               #Collmap[p2.sf-7][p1.sf-7] += 1
               # packets are too close to each other, both collide
               # return both packets as casualties
               return (p1, p2)
           elif p1.rssi[math.ceil(env.now)] - p2.rssi[math.ceil(env.now)] < IsoThresholds[p1.dr-8][p2.dr-8]:
               # p2 overpowered p1, return p1 as casualty
               print ("collision pwr node {} overpowered node {}".format(p2.nodeid, p1.nodeid))
               print ("capture - p2 wins, p1 lost")
               #Collmap[p1.sf-7][p2.sf-7] += 1
               return (p1,)
           print ("capture - p1 wins, p2 lost")
           # p2 was the weaker packet, return it as a casualty
           #Collmap[p2.sf-7][p1.sf-7] += 1
           return (p2,)

    def checkcollision(packet):
        col = 0 # flag needed since there might be several collisions for packet
        processing = 0
        #print ("MAX RECEIVE IS: ", maxBSReceives)
        for i in range(0,len(packetsAtBS)):
            if packetsAtBS[i].header.processed == 1 or packetsAtBS[i].intraPacket.processed == 1:
                processing = processing + 1
        if (processing > maxBSReceives):
            packet.processed = 0
        else:
            packet.processed = 1
        if packetsAtBS:
            #print ("{:3.5f} || >> FOUND overlap... node {} (sf:{} bw:{} freq:{}) others: {}".format(env.now,packet.nodeid, packet.sf, packet.bw,packet.freq,len(packetsAtBS)))
            for other in packetsAtBS:
                if other.nodeid != packet.nodeid:
                   #print ("{:3.5f} || >> node {} overlapped with node {} (sf:{} bw:{} freq:{}). Let's check Freq...".format(env.now,packet.nodeid, other.nodeid, other.packet.sf, other.packet.bw,other.packet.freq))
                   # simple collision
                   #if frequencyCollision(packet, other.packet) and sfCollision(packet, other.packet):
                   if frequencyCollision(packet, other.header):# and timingCollision(packet, other.packet):
                       c = powerCollision_2(packet, other.header)
                       for p in c:
                          p.col = 1
                          if p == packet:
                             col = 1
                   if frequencyCollision(packet, other.intraPacket):# and timingCollision(packet, other.packet):
                       c = powerCollision_2(packet, other.intraPacket)
                       for p in c:
                          p.col = 1
                          if p == packet:
                             col = 1
            return col
        return 0

    def frequencyCollision(p1,p2):
        if (p1.ch == p2.ch):
            #print ("{:3.5f} || >> same channel for header on node {} and node {}.. Let's check sub-channels...".format(env.now,p1.nodeid,p2.nodeid))
            #if (p1.freqHopHeader[replica] == p2.freqHopHeader[replica]):
            if (p1.subCh == p2.subCh):
                #print ("{:3.5f} || >> same sub-channel for header on node {} and node {}".format(env.now,p1.nodeid,p2.nodeid))
                #print ("{:3.5f} || >> Header {} from node {} collided!!!".format(env.now,replica,p1.nodeid))
                return True
            else:
                #print ("{:3.5f} || >> No sub-channel collision".format(env.now))
                return False
        else:
            #print ("{:3.5f} || >> No header channel collision..".format(env.now))
            return False

    class myNode():
        def __init__(self, nodeid, bs, avgSendTime, packetlen, total_data):
            global channel
            global DR
            self.dr = 8
            #carriers = list(range(280))
            #random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
            self.nodeid = nodeid
            self.avgSendTime = avgSendTime
            self.bs = bs
            self.dist = distance[nodeid,:]
            self.elev = elev[nodeid,:]
            self.mindist = np.amin(distance[nodeid,:])
            self.mindist_pos = int(np.where(distance[nodeid,:] == np.amin(distance[nodeid,:]))[0])
            #print('node %d' %nodeid, "dist: ", self.dist[0])
            self.buffer = total_data
            self.packetlen = packetlen
            self.ch = int(random.choice(channel))
            self.packet = myPacket(self.nodeid, packetlen, self.dist)
            #self.freqHop = carriers[0:35]
            self.sent = 0 #INITIAL SENT PACKETS
            self.totalLost = 0 #INITIAL TOTAL LOST FOR PARTICULAR NODE
            self.totalColl = 0
            self.totalRec = 0
            self.totalProc = 0
            if self.dr == 8:
                carriers = list(range(280))
                random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
                self.freqHop = carriers[0:35]
            elif self.dr == 9:
                carriers = list(range(280))
                random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
                self.freqHop = carriers[0:35]
            elif self.dr == 10:
                carriers = list(range(688))
                random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
                self.freqHop = carriers[0:86]
            elif self.dr == 11:
                carriers = list(range(688))
                random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
                self.freqHop = carriers[0:86]

            self.header = myHeader(self.nodeid,self.dist,self.ch,self.freqHop, self.dr)
            self.intraPacket = myIntraPacket(self.nodeid,self.dist,self.ch,self.freqHop,self.dr)

    class myHeader ():
        def __init__(self,nodeid,dist,ch,freqHop,dr):
            global Ptx
            global Prx
            global Lpl
            global c
            global distance
            global channel
            global frequency
            self.nodeid = nodeid
            self.txpow = Ptx
            self.transRange = 150
            self.arriveTime = 0
            self.rssi = Prx[nodeid,:]
            self.rectime = 0.233
            #self.rectime = 1.5
            self.proptime = distance[nodeid,:]*(1/c)
            self.collided = 0
            self.noCollided = 0
            self.processed = 0
            self.noProcessed = 0
            self.ch = ch
            self.lost = bool
            self.Nlost = 0
            self.subCh = 0
            self.sentIntra = 0
            self.dr = dr
            self.col =0
            if dr == "dr8":
                self.freqHopHeader = freqHop[0:3]
            elif dr == "dr9":
                self.freqHopHeader = freqHop[0:2]
            elif dr == "dr10":
                self.freqHopHeader = freqHop[0:3]
            elif dr == "dr11":
                self.freqHopHeader = freqHop[0:2]

    class myIntraPacket ():
        def __init__(self,nodeid,dist,ch,freqHop,dr):
            global Ptx
            global Prx
            global Lpl
            global c
            global distance
            global channel
            global frequency
            self.nodeid = nodeid
            self.txpow = Ptx
            self.transRange = 150
            self.arriveTime = 0
            self.rssi = Prx[nodeid,:]
            self.freqHopIntraPacket = freqHop[3:]
            self.rectime = 50e-3
            #self.rectime = 3
            self.proptime = distance[nodeid,:]*(1/c)
            self.collided = 0
            self.noCollided = 0
            self.nrColl = 0
            self.processed = 0
            self.noProcessed = 0
            self.ch = ch
            self.lost = bool
            self.Nlost = 0
            self.subCh = 0
            self.sentIntra = 0
            self.dr = dr
            self.col =0
            if dr == "dr8":
                self.freqHopIntraPacket = freqHop[3:]
            elif dr == "dr9":
                self.freqHopIntraPacket = freqHop[2:]
            elif dr == "dr10":
                self.freqHopIntraPacket = freqHop[3:]
            elif dr == "dr11":
                self.freqHopIntraPacket = freqHop[2:]

    class myPacket():
        def __init__(self, nodeid, packetlen, dist):
            #global experiment
            global Ptx
            global Prx
            #global gamma
            #global d0
            #global var
            global Lpl
            #global freq
            #global GL
            global c
            global distance
            global channel
            global frequency
            #SF = [7,8,9,10,11,12]
            self.nodeid = nodeid
            self.txpow = Ptx
            #self.sf = random.choice(SF)
            self.sf = 12
            self.cr = 1 ##CODING RATE
            self.bw = 125
            # transmission range, needs update XXX
            self.transRange = 150
            self.pl = packetlen
            self.symTime = (2.0**self.sf)/self.bw
            self.arriveTime = 0
            self.rssi = Prx[nodeid,:]
            self.freq = int(random.choice(frequency))
            self.rectime = airtime(self.sf,self.cr,self.pl,self.bw) ##RECTIME IS THE RECEPTION TIME (ie AIRTIME)
            self.proptime = distance[nodeid,:]*(1/c)
            #print ("rectime node ", self.nodeid, "  ", self.rectime)
            #print ("Airtime for node {} is {} [seconds]".format(self.nodeid,self.rectime)) #from https://www.loratools.nl/#/airtime
            # denote if packet is collided
            self.collided = 0
            self.processed = 0
            self.lost = bool

    def airtime(sf,cr,pl,bw):
        H = 0        # implicit header disabled (H=0) or not (H=1)
        DE = 0       # low data rate optimization enabled (=1) or not (=0)
        Npream = 8   # number of preamble symbol (12.25  from Utz paper)

        if bw == 125 and sf in [11, 12]:
            # low data rate optimization mandated for BW125 with SF11 and SF12
            DE = 1
        if sf == 6:
            # can only have implicit header with SF6
            H = 1

        Tsym = (2.0**sf)/bw
        Tpream = (Npream + 4.25)*Tsym
        #print ("PARAMS FOR TRANSMISION: sf", sf, " cr", cr, "pl", pl, "bw", bw)
        payloadSymbNB = 8 + max(math.ceil((8.0*pl-4.0*sf+28+16-20*H)/(4.0*(sf-2*DE)))*(cr+4),0)
        Tpayload = payloadSymbNB * Tsym
        return ((Tpream + Tpayload)/1000) ##IN SECS

    def selectDR (env,node):
        global DR
        dr = [8,9,10,11]
        node.dr = random.choice(dr)
        node.header.dr = node.dr
        node.intraPacket.dr = node.dr
        if node.dr == 11:
            node.sensi = DR[3]
            carriers = list(range(688))
            random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
            node.freqHop = carriers[0:86]
            node.header.freqHopHeader = node.freqHop[0:2]
            node.intraPacket.freqHopIntraPacket = node.freqHop [2:]
        elif node.dr == 10:
            node.sensi = DR[2]
            carriers = list(range(688))
            random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
            node.freqHop = carriers[0:86]
            node.header.freqHopHeader = node.freqHop[0:3]
            node.intraPacket.freqHopIntraPacket = node.freqHop [3:]
        elif node.dr == 9:
            node.sensi = DR[1]
            carriers = list(range(280))
            random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
            node.freqHop = carriers[0:35]
            node.header.freqHopHeader = node.freqHop[0:2]
            node.intraPacket.freqHopIntraPacket = node.freqHop [2:]
        else:
            node.sensi = DR[0]
            carriers = list(range(280))
            random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
            node.freqHop = carriers[0:35]
            node.header.freqHopHeader = node.freqHop[0:3]
            node.intraPacket.freqHopIntraPacket = node.freqHop [3:]

    def transmit(env,node):
        #while nodes[node.nodeid].buffer > 0.0:
        global wait_min
        global wait_max
        global back_off
        global beacon_time
        global logs
        global nodesToSend
        global DR
        while node.buffer > 0.0:
            #######STARTS TRANSMISSION AS DR8
            node.dr = 8
            node.sensi = -137
            carriers = list(range(280))
            random.shuffle(carriers) #TO CHOOSE THE HOPPING JUMPS
            node.freqHop = carriers[0:35]
            node.header.freqHopHeader = node.freqHop[0:3]
            node.intraPacket.freqHopIntraPacket = node.freqHop [3:]
            #######
            yield env.timeout(node.packet.rectime + float(node.packet.proptime[math.ceil(env.now)])) ##GIVE TIME TO RECEIVE BEACON

            if node in packetsAtBS:
                print ("{:3.5f} || ERROR: packet is already in...".format(env.now))
            else:
                sensibility = sensi[12 - 7, [125,250,500].index(node.packet.bw) + 1]
                if node.packet.rssi[math.ceil(env.now)] < sensibility: #HERE WE ARE CONSIDERING RSSI AT TIME ENV.NOW
                    print ("{:3.5f} || Node {}: Can not reach beacon due Lpl".format(env.now,node.nodeid))
                    wait =0 ##LETS WAIT FOR NEXT BEACON
                    node.header.lost = False
                    node.intraPacket.lost = False
                    trySend = False
                    nIntraPackets = 0
                else:
                    selectDR(env,node)
                    nodesToSend.append(node.nodeid)
                    wait = random.uniform(1,back_off - node.packet.rectime - float(node.packet.proptime[math.ceil(env.now)])) ##TRIGGER BACK-OFF TIME
                    yield env.timeout(wait)
                    #print ("{:3.5f} || Node {} begins to transmit a packet".format(env.now,node.nodeid))
                    trySend = True
                    node.sent = node.sent + 1
                    node.buffer = node.buffer - node.packetlen
                    if node in packetsAtBS:
                        print ("{} || ERROR: packet is already in...".format(env.now))
                    else:
                        #sensibility = sensi[node.packet.sf - 7, [125,250,500].index(node.packet.bw) + 1]
                        sensibility = node.sensi
                        #print ("------Sensi is: ",sensibility)
                        if node.packet.rssi[math.ceil(env.now)] < sensibility: #HERE WE ARE CONSIDERING RSSI AT TIME ENV.NOW
                            print ("{:3.5f} || Node {}: The Packet will be Lost due Lpl".format(env.now,node.nodeid))
                            node.header.lost = True ## LOST ONLY CONSIDERING Lpl
                            node.intraPacket.lost = True ## LOST ONLY CONSIDERING Lpl
                            #nIntraPackets = 0
                            print ("###############lost !!!!!!!!")
                        else:
                            node.header.lost = False ## LOST ONLY CONSIDERING Lpl
                            node.intraPacket.lost = False ## LOST ONLY CONSIDERING Lpl
                            #print ("{:3.5f} || Prx for node {} is {:3.2f} dB".format(env.now, node.nodeid, node.packet.rssi[math.ceil(env.now)]))
                            #print ("Prx for node",node.nodeid, "is: ",node.packet.rssi[math.ceil(env.now)],"at time",env.now)

                        for i in range(len(node.header.freqHopHeader)):
                            ###print ("{:3.5f} || Sending Header replica {} node {}...".format(env.now,i,node.nodeid))
                            ###print ("{:3.5f} || Let's try if there are collisions...".format(env.now))
                            node.header.subCh = node.header.freqHopHeader[i]
                            #print ("SUBCHANELLLL: ",node.header.subCh)
                            node.header.sentIntra +=1
                            isLost =0
                            if node.packet.rssi[math.ceil(env.now)] < sensibility:
                                node.header.Nlost +=1
                                isLost =1
                            if (checkcollision(node.header)==1):
                                #pass
                                if node.header.col == 1:
                                    if isLost == 0:
                                        node.header.collided +=1
                                #node.packet.collided = 1
                                #print ("---{:3.5f} || Collision for Header replica {} node {} !!!".format(env.now,i,node.nodeid))
                                #node.packet.collided = 1
                                #node.header.collided +=1 #ALREADY COUNTED IN FUNCTION
                            else:
                                ###print ("{:3.5f} || ...No Collision for Header replica {} node {}!".format(env.now,i,node.nodeid))
                                #node.packet.collided = 0
                                node.header.noCollided = 1 ##ALMOST ONE HEADER IS OK, THEN HEADER IS OK
                            packetsAtBS.append(node)
                            node.packet.addTime = env.now
                            isLost =0
                            yield env.timeout(node.header.rectime)
                            if (node in packetsAtBS):
                                packetsAtBS.remove(node)

                        ##CALCULATE N OF INTRAPACKETS BASED ON PACKETLEN
                        #airtime(12,1,node.packetlen,125)
                        if node.dr == 8 or node.dr == 10:
                            payloadTime = 1.85 - 0.233*3
                        elif node.dr == 9 or node.dr ==11:
                            payloadTime = 1.07 - 0.233*2

                        nIntraPackets = math.ceil(payloadTime / 50e-3)
                        #print ("NUMBER OF INTRA PACKETSSSS",nIntraPackets)

                        for j in range (nIntraPackets):
                            ###print ("{:3.5f} || Sending intra-packet {} of {} for node {}...".format(env.now,j,nIntraPackets-1,node.nodeid))
                            ###print ("{:3.5f} || Let's try if there are collisions...".format(env.now))
                            node.intraPacket.subCh = node.intraPacket.freqHopIntraPacket[j]
                            node.intraPacket.sentIntra +=1
                            isLost =0
                            if node.packet.rssi[math.ceil(env.now)] < sensibility:
                                node.intraPacket.Nlost +=1
                                isLost =1
                            #print ("INTRA-PACKT SUB CHANNELLLL", node.intraPacket.subCh)
                            if (checkcollision(node.intraPacket)==1):
                                #pass
                                if node.intraPacket.col ==1:
                                    if isLost == 0:
                                        node.intraPacket.collided+=1
                                #print ("---{:3.5f} || Collision for intra-packet {} for node {} !!!".format(env.now,j,node.nodeid))
                                #node.intraPacket.collided+=1 #ALREADY COUNTED ON FUNCTION
                            else:
                                ###print ("{:3.5f} || ...No Collision for intra-packet {} for node {}!".format(env.now,j,node.nodeid))
                                node.intraPacket.noCollided +=1
                                pass
                            packetsAtBS.append(node)
                            node.packet.addTime = env.now
                            yield env.timeout(node.intraPacket.rectime)
                            if (node in packetsAtBS):
                                packetsAtBS.remove(node)
                            #print ("INTRA-PACKET NO-PROCESEDDD",node.intraPacket.noProcessed)
                            isLost =0

            node.header.noCollided = len(node.header.freqHopHeader)-node.header.Nlost-node.header.collided
            node.intraPacket.noCollided = nIntraPackets-node.intraPacket.Nlost-node.intraPacket.collided

            if node.header.noCollided <0:
                node.header.noCollided = 0
            if node.intraPacket.noCollided <0:
                node.intraPacket.noCollided = 0

            if trySend == 1:
                #print ("----count intra-packet collided", node.intraPacket.collided)
                if node.header.lost or node.intraPacket.lost:
                    logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},PL,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))
                else:
                    if node.dr ==8 or node.dr==10:
                        if node.header.collided == 3:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},PCh,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))
                        elif node.intraPacket.collided > (1/3)*nIntraPackets:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},PCp,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))
                        elif node.header.noProcessed == 3:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},NP,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))
                        elif node.intraPacket.noProcessed > (1/3)*nIntraPackets:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},NP,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))
                        else:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},PE,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))

                    elif node.dr==9 or node.dr==11:
                        if node.header.collided == 2:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},PCh,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))
                        elif node.intraPacket.collided > (2/3)*nIntraPackets:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},PCp,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))
                        elif node.header.noProcessed == 2:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},NP,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))
                        elif node.intraPacket.noProcessed > (2/3)*nIntraPackets:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},NP,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))
                        else:
                            logs.append("{:3.3f},{},{:3.3f},{:3.3f},{},PE,#{},#{},#{},#{}".format(env.now,node.nodeid,node.dist[math.ceil(env.now)],node.elev[math.ceil(env.now)],node.dr,nIntraPackets,node.intraPacket.noCollided,len(node.header.freqHopHeader),node.header.noCollided))

            ##RESET
            node.header.collided = 0
            node.header.processed = 0
            node.header.noProcessed = 0
            node.header.lost = False
            node.header.noCollided =0
            node.intraPacket.nrColl = 0
            node.intraPacket.collided = 0
            node.intraPacket.processed = 0
            node.intraPacket.noProcessed = 0
            node.intraPacket.lost = False
            node.intraPacket.noCollided = 0
            node.header.sentIntra = 0
            node.intraPacket.sentIntra = 0
            node.header.Nlost =0
            node.intraPacket.Nlost = 0

            if trySend:
                #print ("BEACON TIMEEE",beacon_time)
                #print ("WAITTT",wait)
                #print ("NODE HEADER TIME",node.header.rectime)
                #print ("ONE INTRA-PACKET TIMEE",node.intraPacket.rectime)
                #yield env.timeout(beacon_time-wait)
                yield env.timeout(beacon_time-wait-2*3*node.header.rectime-2*nIntraPackets*node.intraPacket.rectime)
            else:
                nIntraPackets = 0
                yield env.timeout(beacon_time-wait-3*node.header.rectime-nIntraPackets*node.intraPacket.rectime)

    def beacon (env):
        global beacon_time
        global nodesToSend
        global logs
        i = 0
        while True:
            if i == 0:
                yield env.timeout(0)
            else:
                yield env.timeout(beacon_time-2)
            i=i+1
            print ("{:3.5f} || ***A new beacon has been sended from Satellite***".format(env.now))
            yield env.timeout(2)
            logs.append("{:3.3f},B,{}".format(env.now,nodesToSend))
            nodesToSend = []

    env.process(beacon(env)) ##BEACON SENDER

    ### THIS IS GOING TO CREATE NODES AND DO TRAMSMISIONS. IS THE MAIN PROGRAM ###
    for i in range(nrNodes):
        node = myNode(i,bsId, avgSendTime, packetlen, total_data)
        nodes.append(node)
        env.process(transmit(env,node))

    env.run(until=600*2)

    sent = sum(n.sent for n in nodes)
    return ([sent,nrCollFullPacket,None,None,nrReceived],logs)

#########################################################################
if chan == 1:
    ###SCENARIO 1 CHANNEL###
    channel = [0]

    nodes = [] ###EACH NODE WILL BE APPENDED TO THIS VARIABLE
    nrLost = 0 ### TOTAL OF LOST PACKETS DUE Lpl
    nrCollisions = 0 ##TOTAL OF COLLIDED PACKETS
    nrProcessed = 0 ##TOTAL OF PROCESSED PACKETS
    nrReceived = 0 ###TOTAL OF RECEIVED PACKETS
    nrNoProcessed = 0 ##TOTAL OF INTRA-PACKETS NO PROCESSED
    nrIntraTot = 0
    nrLostMaxRec = 0
    nrCollFullPacket = 0
    nrSentIntra = 0 ##TOTAL OF SENT INTRA-PACKTES
    nrReceivedIntra = 0 ##TOTAL OF RECEIVED INTRA-PACKETS

    i =0
    scenario_1ch = np.zeros((len(multi_nodes),5))
    results = []
    ## WHERE:
        ## scenario_1ch[i,j]:
            ## i --> the node i
            ## j --> [sent, nrCollisions, nrLost, nrProcessed, nrReceived]

    for nrNodes in multi_nodes:
        print ("\n\n***NEW SCENARIO BEGINS***\n")
        logs = []
        results,logs = simulate_scenario(nrNodes)
        scenario_1ch[i,:] = results
        folder = name+'_1CH_s'+str(RANDOM_SEED)+'_p'+str(packetsToSend)
        if not os.path.exists(folder):
            os.makedirs(folder)
        fname = "./"+folder+"/" + str(name+"_"+str(nrNodes)+"_1CH_"+str(maxBSReceives)+"_s"+str(RANDOM_SEED)+"_p"+str(packetsToSend)) + ".csv"
        with open(fname,"w") as myfile:
            myfile.write("\n".join(logs))
        myfile.close()
        i=i+1
        if not mode_debbug:
            nodes = [] ###EACH NODE WILL BE APPENDED TO THIS VARIABLE
        nrLost = 0 ### TOTAL OF LOST PACKETS DUE Lpl
        nrCollisions = 0 ##TOTAL OF COLLIDED PACKETS
        nrProcessed = 0 ##TOTAL OF PROCESSED PACKETS
        nrReceived = 0 ###TOTAL OF RECEIVED PACKETS
        nrNoProcessed = 0 ##TOTAL OF INTRA-PACKETS NO PROCESSED
        nrIntraTot = 0
        nrLostMaxRec = 0
        nrCollFullPacket = 0
        nrSentIntra = 0 ##TOTAL OF SENT INTRA-PACKTES
        nrReceivedIntra = 0 ##TOTAL OF RECEIVED INTRA-PACKETS

#########################################################################
if chan ==3:
    ###SCENARIO 3 CHANNELS###
    channel = [0,1,2]
    nodes = [] ###EACH NODE WILL BE APPENDED TO THIS VARIABLE
    nrLost = 0 ### TOTAL OF LOST PACKETS DUE Lpl
    nrCollisions = 0 ##TOTAL OF COLLIDED PACKETS
    nrProcessed = 0 ##TOTAL OF PROCESSED PACKETS
    nrReceived = 0 ###TOTAL OF RECEIVED PACKETS
    nrNoProcessed = 0 ##TOTAL OF INTRA-PACKETS NO PROCESSED
    nrIntraTot = 0
    nrLostMaxRec = 0
    nrCollFullPacket = 0
    nrSentIntra = 0 ##TOTAL OF SENT INTRA-PACKTES
    nrReceivedIntra = 0 ##TOTAL OF RECEIVED INTRA-PACKETS
    i =0
    scenario_3ch = np.zeros((len(multi_nodes),5))
    results = []

    for nrNodes in multi_nodes:
        print ("\n\n***NEW SCENARIO BEGINS***\n")
        logs = []
        results,logs = simulate_scenario(nrNodes)
        scenario_3ch[i,:] = results
        folder = name+'_3CH_s'+str(RANDOM_SEED)+'_p'+str(packetsToSend)
        if not os.path.exists(folder):
            os.makedirs(folder)
        fname = "./"+folder+"/" + str(name+"_"+str(nrNodes)+"_3CH_"+str(maxBSReceives)+"_s"+str(RANDOM_SEED)+"_p"+str(packetsToSend)) + ".csv"
        with open(fname,"w") as myfile:
            myfile.write("\n".join(logs))
        myfile.close()
        i=i+1
        if not mode_debbug:
            nodes = [] ###EACH NODE WILL BE APPENDED TO THIS VARIABLE
        nrLost = 0 ### TOTAL OF LOST PACKETS DUE Lpl
        nrCollisions = 0 ##TOTAL OF COLLIDED PACKETS
        nrProcessed = 0 ##TOTAL OF PROCESSED PACKETS
        nrReceived = 0 ###TOTAL OF RECEIVED PACKETS
        nrNoProcessed = 0 ##TOTAL OF INTRA-PACKETS NO PROCESSED
        nrIntraTot = 0
        nrLostMaxRec = 0
        nrCollFullPacket = 0
        nrSentIntra = 0 ##TOTAL OF SENT INTRA-PACKTES
        nrReceivedIntra = 0 ##TOTAL OF RECEIVED INTRA-PACKETS

if not mode_debbug:
    sys.stdout = old_stdout
    print("done ",name)


import keras
import os
import numpy as np
import glob
import shutil
import tensorflow as tf

from keras.applications.inception_resnet_v2 import InceptionResNetV2
from keras.applications.resnet50 import ResNet50
from keras.applications.densenet import DenseNet169
from keras.applications.nasnet import NASNetLarge
from keras.models import Model
from keras.models import load_model
from keras.layers import Dense, GlobalAveragePooling2D, LeakyReLU, concatenate, Input, Dropout
from keras import backend as K
from keras.preprocessing import image
from keras.applications.nasnet import preprocess_input
from keras.preprocessing.image import ImageDataGenerator

def get_session():
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    return tf.Session(config=config)

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
# set the modified tf session as backend in keras
keras.backend.tensorflow_backend.set_session(get_session())

dataset_path = '/home/ustc/jql/jsrtdata/'
# model_path = '/home/ustc/jql/all_lung_image/'

# def prepare_train(cls='abnormal'):
#
#     file_names = glob.glob(dataset_path+cls+'/*.png')
#     train_cls_path = dataset_path+'train/'+cls
#     test_cls_path = dataset_path+'test/'+cls
#
#     if not os.path.exists(train_cls_path):
#         os.mkdir(train_cls_path)
#     if not os.path.exists(test_cls_path):
#         os.mkdir(test_cls_path)
#     for index, file in enumerate(file_names):
#         dst_name = file.split('/')[-1]
#         if index % 30 == 0:
#             print(dst_name)
#             shutil.copyfile(file, os.path.join(test_cls_path, dst_name))
#         else:
#             shutil.copyfile(file, os.path.join(train_cls_path, dst_name))

def train(lr=1e-3, batch=16, epoch=1):
    train_flow = ImageDataGenerator(rotation_range=20, width_shift_range=0.2, shear_range=0.2,
                                    height_shift_range=0.2, zoom_range=0.2, horizontal_flip=True)
    train_generator = train_flow.flow_from_directory(dataset_path+'train', target_size=(299, 299),
                                                     classes=['normal', 'nodule'],
                                                     interpolation='bilinear',
                                                     batch_size=batch)
    num_train = train_generator.samples
    print(num_train)
    print(train_generator.class_indices)

    val_flow = ImageDataGenerator()
    val_generator = val_flow.flow_from_directory(dataset_path+'test', target_size=(299, 299),
                                                 classes=['normal', 'nodule'],
                                                 interpolation='bilinear',
                                                 batch_size=batch)
    num_val = val_generator.samples
    # base_model1 = NASNetLarge(weights='imagenet', include_top=True, input_tensor=input)
    base_model = InceptionResNetV2(weights='imagenet', include_top=False, pooling='avg')
    x = base_model.output
    # x = Dense(1024, activation='relu')(x)
    # x = Dropout(0.3)(x)
    prediction = Dense(2, activation='softmax')(x)

    model = Model(inputs=base_model.input, outputs=prediction)
    # model.load_weights('imagenet', by_name=True, skip_mismatch=True)

    for layer in base_model.layers:
        layer.trainable = False

    adam = keras.optimizers.Adam(lr=lr)
    model.compile(optimizer=adam, loss='categorical_crossentropy', metrics=['acc'])

    checkpoint = keras.callbacks.ModelCheckpoint(dataset_path+'cxr8_jsrt.h5',
                                                 monitor='val_loss', verbose=1, save_best_only=True)
    early_stopping = keras.callbacks.EarlyStopping(monitor='val_loss')

    model.fit_generator(train_generator, steps_per_epoch=num_train, verbose=1, epochs=epoch,
                        validation_data=val_generator, validation_steps=int(num_val/batch),
                        callbacks=[checkpoint, early_stopping], shuffle=True)

    for layer in base_model.layers:
        layer.trainable = True
    model.compile(optimizer=adam, loss='categorical_crossentropy', metrics=['acc'])

    model.fit_generator(train_generator, steps_per_epoch=num_train, verbose=1, epochs=epoch,
                        validation_data=val_generator, validation_steps=int(num_val / batch),
                        callbacks=[checkpoint, early_stopping], shuffle=True)

if __name__ == '__main__':
    # prepare_train(cls='normal')
    train(lr=1e-4, batch=1, epoch=5)


from utils import MolecularImageDataset, collate_fn
from torch.utils.data import DataLoader
import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
import numpy as np
from unet import UNet
import matplotlib.pyplot as plt
import os
from meter import AverageMeter
import torch.distributed as dist
import torch.multiprocessing as mp
import argparse

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
parser = argparse.ArgumentParser(description='multi-gpu-training')
parser.add_argument('--data',metavar='dir',default='train_data',help='path to dataset')
parser.add_argument('--epoch',default=15,type=int,metavar='N',help='number of total epochs to run')
parser.add_argument('-b','--batch-size',default=60,type=int,metavar='N',help='mini batch size')
parser.add_argument('--lr','--learning-rate',default=2.5e-4,type=float,
                    metavar='LR',help='initial learning rate',dest='lr')

def reduce_mean(tensor,nprocs):
    rt = tensor.clone()
    dist.all_reduce(rt,op=dist.ReduceOp.SUM)
    rt /= nprocs
    return rt

def main():
    print('start')
    args = parser.parse_args()
    args.nprocs = torch.cuda.device_count()

    print(args.nprocs)
    mp.spawn(main_worker,nprocs=args.nprocs,args=(args.nprocs,args))
    print('done')

def main_worker(local_rank,nprocs,args):
    print('enter')
    args.local_rank = local_rank

    dist.init_process_group(backend='nccl',init_method='tcp://127.0.0.1:12345',
                            world_size=nprocs,rank=local_rank)

    df = pd.read_csv('../data/indigo_train_data/processed_chembl.csv')
    train_df1 = df[:80000].copy().reset_index(drop=True)
    test_df1 = df[90000:91000].copy().reset_index(drop=True)

    df = pd.read_csv('../data/train_data/processed_chembl.csv')
    train_df2 = df[:80000].copy().reset_index(drop=True)
    test_df2 = df[90000:91000].copy().reset_index(drop=True)

    train_df = pd.concat([train_df1, train_df2], axis=0).reset_index(drop=True)
    test_df = pd.concat([test_df1, test_df2], axis=0).reset_index(drop=True)

    epoch_nums = 15
    amount = 0.1

    train_dataset = MolecularImageDataset(train_df, amount=amount)
    test_dataset = MolecularImageDataset(test_df, amount=amount)

    # imgs, atom_targets, atom_types, atom_charges, Hs, bond_targets, bond_types, bond_rhos, bond_omega_types = train_dataset[0]
    # plt.subplot(131)
    # plt.imshow(imgs[0])
    # plt.subplot(132)
    # plt.imshow(atom_targets[0])
    # plt.subplot(133)
    # plt.imshow(bond_targets[0])
    # for x, y in zip(*((bond_targets[0]==1).nonzero())):
    #     rho = bond_rhos[:, x, y]
    #     omega = bond_omega_types[:, x, y].argmax() * (2*np.pi/60) + np.pi/60 - np.pi/2
    #
    #     plt.plot([y - rho* np.sin(omega), y +  rho*np.sin(omega)], [x -  rho*np.cos(omega), x +  rho*np.cos(omega)])
    #
    # plt.show()

    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)

    train_dataloader = DataLoader(train_dataset,batch_size=args.batch_size,collate_fn=collate_fn,num_workers=3,prefetch_factor=10,sampler=train_sampler)
    test_dataloader = DataLoader(test_dataset,batch_size=args.batch_size,collate_fn=collate_fn,num_workers=3,prefetch_factor=10,sampler=test_sampler)

    model = UNet(in_channels=1, heads=[1, 14, 3, 2, 1, 360, 60, 60])
    torch.cuda.set_device(local_rank)
    model.cuda(local_rank)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

    if local_rank==0:
        torch.save(model.state_dict(),'temp.pth')

    dist.barrier()

    model.load_state_dict(torch.load('temp.pth',map_location={'cuda:0':'cuda:{}'.format(local_rank)}))

    optimizer = optim.Adam(model.parameters(), lr=2.5e-4, weight_decay=1e-8)

    torch.cuda.empty_cache()

    train_atom_targets_precision = AverageMeter()
    train_atom_targets_recall = AverageMeter()
    train_atom_targets_precision3 = AverageMeter()
    train_atom_targets_recall3 = AverageMeter()

    train_atom_types_acc = AverageMeter()
    train_atom_charges_acc = AverageMeter()
    train_atom_hs_acc = AverageMeter()

    train_bond_targets_precision = AverageMeter()
    train_bond_targets_recall = AverageMeter()
    train_bond_targets_precision3 = AverageMeter()
    train_bond_targets_recall3 = AverageMeter()

    train_bond_types_acc = AverageMeter()
    train_bond_rhos_mae = AverageMeter()

    train_bond_omega_precision = AverageMeter()
    train_bond_omega_recall = AverageMeter()
    train_bond_omega_precision3 = AverageMeter()
    train_bond_omega_recall3 = AverageMeter()

    for epoch in range(epoch_nums):
        if epoch == int(epoch_nums / 3):
            optimizer = optim.Adam(model.parameters(), lr=2.5e-5, weight_decay=1e-8)
        for i, (imgs, atom_targets, atom_types, atom_charges, atom_hs,
                bond_targets, bond_types, bond_rhos, bond_omega_types) in enumerate(train_dataloader):

            model.train()
            imgs = imgs.to(device)
            atom_targets, atom_types, atom_charges, atom_hs = atom_targets.to(device), atom_types.to(
                device), atom_charges.to(device), atom_hs.to(device)
            bond_targets, bond_types, bond_rhos, bond_omega_types = bond_targets.to(device), bond_types.to(
                device), bond_rhos.to(device), bond_omega_types.to(device)

            atom_targets_pred, atom_types_pred, atom_charges_pred, atom_hs_pred, bond_targets_pred, bond_types_pred, bond_rhos_pred, bond_omega_types_pred = model(
                imgs)
            atom_targets_pred = torch.clamp(torch.sigmoid(atom_targets_pred), 1e-5, 1 - 1e-5)
            atom_types_pred = torch.clamp(torch.softmax(atom_types_pred, dim=1), 1e-5, 1 - 1e-5)
            atom_charges_pred = torch.clamp(torch.softmax(atom_charges_pred, dim=1), 1e-5, 1 - 1e-5)
            atom_hs_pred = torch.clamp(torch.softmax(atom_hs_pred, dim=1), 1e-5, 1 - 1e-5)

            bond_targets_pred = torch.clamp(torch.sigmoid(bond_targets_pred), 1e-5, 1 - 1e-5)
            bond_types_pred = torch.clamp(torch.softmax(bond_types_pred.view(-1, 6, 60, 128, 128), dim=1), 1e-5, 1 - 1e-5)

            bond_omega_types_pred = torch.clamp(torch.sigmoid(bond_omega_types_pred), 1e-5, 1 - 1e-5)

            bond_rhos_pred = torch.abs(bond_rhos_pred)

            atom_targets_loss = torch.sum(
                -(atom_targets == 1).float() * (1 - atom_targets_pred) ** 2 * torch.log(atom_targets_pred)
                - (1 - atom_targets) ** 4 * (atom_targets_pred) ** 2 * torch.log(1 - atom_targets_pred)) / torch.sum(
                atom_targets == 1)
            atom_types_loss = torch.sum(-atom_types * (1 - atom_types_pred) ** 2 * torch.log(atom_types_pred)) / torch.sum(
                atom_types)

            atom_charges_loss = torch.sum(
                -atom_charges * (1 - atom_charges_pred) ** 2 * torch.log(atom_charges_pred)) / torch.sum(atom_charges)

            # atom_hs_weights = torch.FloatTensor([0.25,1]).reshape(1,2,1,1).to(device)
            atom_hs_loss = torch.sum(-atom_hs * (1 - atom_hs_pred) ** 2 * torch.log(atom_hs_pred)) / (
                        torch.sum(atom_hs) + 0.1)

            bond_targets_loss = torch.sum(
                -(bond_targets == 1).float() * (1 - bond_targets_pred) ** 2 * torch.log(bond_targets_pred)
                - (1 - bond_targets) ** 4 * (bond_targets_pred) ** 2 * torch.log(1 - bond_targets_pred)) / torch.sum(
                bond_targets == 1)

            bond_types_loss = torch.sum(-bond_types * (1 - bond_types_pred) ** 2 * torch.log(bond_types_pred)) / torch.sum(
                bond_types)

            bond_rhos_loss = torch.sum(torch.abs(bond_rhos_pred - bond_rhos) * torch.sum(bond_types, dim=1)) / torch.sum(
                bond_types)

            bond_omega_types_loss = -torch.sum(torch.sum(bond_omega_types, dim=1, keepdim=True) * (
                        (bond_omega_types == 1) * ((1 - bond_omega_types_pred) ** 2) * torch.log(bond_omega_types_pred) +
                        (1 - bond_omega_types) ** 4 * (bond_omega_types_pred ** 2) * torch.log(
                    1 - bond_omega_types_pred))) / torch.sum(bond_omega_types)

            atom_targets_loss *= torch.exp(-model.module.s[0]) + model.module.s[0]
            bond_targets_loss *= torch.exp(-model.module.s[1]) + model.module.s[1]
            atom_types_loss *= torch.exp(-model.module.s[2]) + model.module.s[2]
            atom_charges_loss *= torch.exp(-model.module.s[3]) + model.module.s[3]

            bond_types_loss *= torch.exp(-model.module.s[4]) + model.module.s[4]
            bond_rhos_loss *= 0.5 * torch.exp(-model.module.s[6]) + model.module.s[6]
            bond_omega_types_loss *= torch.exp(-model.module.s[7]) + model.module.s[7]
            atom_hs_loss *= torch.exp(-model.module.s[9]) + model.module.s[9]

            loss = atom_targets_loss + bond_targets_loss + atom_types_loss + atom_charges_loss + bond_types_loss + bond_rhos_loss + bond_omega_types_loss + atom_hs_loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            ####metrics

            temp = torch.nn.functional.max_pool2d(atom_targets_pred, kernel_size=3,
                                                  stride=1, padding=1)
            atom_targets_pred = ((temp == atom_targets_pred) * (atom_targets_pred > 0.25)).float()

            temp = torch.nn.functional.max_pool2d(bond_targets_pred, kernel_size=3,
                                                  stride=1, padding=1)
            bond_targets_pred = ((temp == bond_targets_pred) * (bond_targets_pred > 0.25)).float()

            atom_targets = (atom_targets == 1).float()
            train_atom_targets_precision.update(
                ((atom_targets_pred * atom_targets).sum() / atom_targets_pred.sum()).cpu().detach().numpy(),
                atom_targets_pred.sum().cpu().detach().numpy())
            train_atom_targets_precision3.update(
                ((atom_targets_pred * torch.nn.functional.max_pool2d(atom_targets, kernel_size=3, padding=1,
                                                                     stride=1)).sum() / atom_targets_pred.sum()).cpu().detach().numpy(),
                atom_targets_pred.sum().cpu().detach().numpy())

            train_atom_targets_recall.update(
                ((atom_targets * atom_targets_pred).sum() / atom_targets.sum()).cpu().detach().numpy(),
                atom_targets.sum().cpu().detach().numpy())
            train_atom_targets_recall3.update(
                ((atom_targets * torch.nn.functional.max_pool2d(atom_targets_pred, kernel_size=3, padding=1,
                                                                stride=1)).sum() / atom_targets.sum()).cpu().detach().numpy(),
                atom_targets.sum().cpu().detach().numpy())

            train_atom_types_acc.update((torch.sum(torch.sum(atom_types, dim=1) * ((atom_types.argmax(1) ==
                                                                                atom_types_pred.argmax(
                                                                                1)).float())) / torch.sum(
                atom_types)).cpu().detach().numpy(), torch.sum(atom_types).cpu().detach().numpy())

            train_atom_charges_acc.update((torch.sum(torch.sum(atom_charges, dim=1) * ((atom_charges.argmax(1) ==
                                                                                atom_charges_pred.argmax(
                                                                                1)).float())) / torch.sum(
                atom_charges)).cpu().detach().numpy(), torch.sum(atom_charges).cpu().detach().numpy())

            train_atom_hs_acc.update((torch.sum(torch.sum(atom_hs, dim=1) * ((atom_hs.argmax(1) ==
                                                                              atom_hs_pred.argmax(1)).float())) / (
                                                  0.01 + torch.sum(atom_hs))).cpu().detach().numpy(),
                                     (0.01 + torch.sum(atom_hs)).cpu().detach().numpy())

            bond_targets = (bond_targets == 1).float()

            train_bond_targets_precision.update(
                ((bond_targets_pred * bond_targets).sum() / bond_targets_pred.sum()).cpu().detach().numpy(),
                bond_targets_pred.sum().cpu().detach().numpy())
            train_bond_targets_precision3.update(
                ((bond_targets_pred * torch.nn.functional.max_pool2d(bond_targets, kernel_size=3, padding=1,
                                                                     stride=1)).sum() / bond_targets_pred.sum()).cpu().detach().numpy(),
                bond_targets_pred.sum().cpu().detach().numpy())

            train_bond_targets_recall.update(
                ((bond_targets * bond_targets_pred).sum() / bond_targets.sum()).cpu().detach().numpy(),
                bond_targets.sum().cpu().detach().numpy())
            train_bond_targets_recall3.update(
                ((bond_targets * torch.nn.functional.max_pool2d(bond_targets_pred, kernel_size=3, padding=1,
                                                                stride=1)).sum() / bond_targets.sum()).cpu().detach().numpy(),
                bond_targets.sum().cpu().detach().numpy())

            train_bond_types_acc.update((torch.sum(
                torch.sum(bond_types, dim=1) * ((bond_types.argmax(1) == bond_types_pred.argmax(1)).float())) / torch.sum(
                bond_types)).cpu().detach().numpy(), torch.sum(bond_types).cpu().detach().numpy())

            train_bond_rhos_mae.update((torch.sum(torch.abs(bond_rhos_pred - bond_rhos) * torch.sum(bond_types, dim=1)) / \
                                        torch.sum(bond_types)).cpu().detach().numpy(),
                                       torch.sum(bond_types).cpu().detach().numpy())

            temp = torch.cat([bond_omega_types_pred[:, 59:], bond_omega_types_pred, bond_omega_types_pred[:, :1]],
                             dim=1).permute(0, 2, 3, 1).reshape(-1, 128 * 128, 62)

            temp = ((torch.nn.functional.max_pool1d(temp, stride=1, kernel_size=3, padding=0).reshape(-1, 128, 128,
                                                                                60).permute(0, 3, 1,
                                                                                2) == bond_omega_types_pred) * \
                    (bond_omega_types_pred > 0.25)).float() * bond_targets

            bond_omega_types = (bond_omega_types == 1)
            train_bond_omega_precision.update((((bond_omega_types) * temp).sum() / (temp).sum()).cpu().detach().numpy(),
                                              temp.sum().cpu().detach().numpy())

            temp2 = torch.cat([temp[:, 59:], temp, temp[:, :1]],
                              dim=1).permute(0, 2, 3, 1).reshape(-1, 128 * 128, 62)
            temp2 = (torch.nn.functional.max_pool1d(temp2.float(), stride=1, kernel_size=3, padding=0).reshape(-1, 128, 128,
                                                                                60).permute(
                0, 3, 1, 2)).float()

            train_bond_omega_recall3.update(
                (((bond_omega_types) * temp2).sum() / (bond_omega_types).sum()).cpu().detach().numpy(),
                bond_omega_types.sum().cpu().detach().numpy())

            train_bond_omega_recall.update(
                (((bond_omega_types) * temp).sum() / (bond_omega_types).sum()).cpu().detach().numpy(),
                bond_omega_types.sum().cpu().detach().numpy())

            temp3 = torch.cat([bond_omega_types[:, 59:], bond_omega_types, bond_omega_types[:, :1]],
                              dim=1).permute(0, 2, 3, 1).reshape(-1, 128 * 128, 62)
            temp3 = (torch.nn.functional.max_pool1d(temp3.float(), stride=1, kernel_size=3, padding=0).reshape(-1, 128, 128,
                                                                                60).permute(
                0, 3, 1, 2)).float()
            train_bond_omega_precision3.update((((temp3) * temp).sum() / (temp).sum()).cpu().detach().numpy(),
                                               temp.sum().cpu().detach().numpy())

            if i % 100 == 0:
                model.eval()
                print(epoch, i)
                print(loss.cpu().detach().numpy())
                with torch.no_grad():
                    if local_rank==0:
                        print('loss_________________')
                        print('atom_targets_loss:', atom_targets_loss.cpu().detach().numpy())
                        print('atom_types_loss:', atom_types_loss.cpu().detach().numpy())
                        print('atom_charges_loss:', atom_charges_loss.cpu().detach().numpy())
                        print('atom_hs_loss:', atom_hs_loss.cpu().detach().numpy())
                        print('bond_targets_loss:', bond_targets_loss.cpu().detach().numpy())
                        print('bond_types_loss:', bond_types_loss.cpu().detach().numpy())

                        print('bond_rho_loss:', bond_rhos_loss.cpu().detach().numpy())
                        print('bond_omega_types:', bond_omega_types_loss.cpu().detach().numpy())
                        print('train_________________')

                        print('atom_target_precision:', train_atom_targets_precision.avg)
                        print('atom_target_recall:', train_atom_targets_recall.avg)
                        print('atom_target_precision3:', train_atom_targets_precision3.avg)
                        print('atom_target_recall3:', train_atom_targets_recall3.avg)

                        print('atom_types_acc:', train_atom_types_acc.avg)
                        print('atom_charges_acc:', train_atom_charges_acc.avg)
                        print('atom_hs_acc:', train_atom_hs_acc.avg)

                        print('bond_target_precision:', train_bond_targets_precision.avg)
                        print('bond_target_recall:', train_bond_targets_recall.avg)
                        print('bond_target_precision3:', train_bond_targets_precision3.avg)
                        print('bond_target_recall3:', train_bond_targets_recall3.avg)

                        print('bond_types_acc:', train_bond_types_acc.avg)
                        print('bond_rhos_mae:', train_bond_rhos_mae.avg)

                        print('bond_omega_precision:', train_bond_omega_precision.avg)
                        print('bond_omega_recall:', train_bond_omega_recall.avg)
                        print('bond_omega_precision3:', train_bond_omega_precision3.avg)
                        print('bond_omega_recall3:', train_bond_omega_recall3.avg)

                    train_atom_targets_precision.reset()
                    train_atom_targets_recall.reset()
                    train_atom_targets_precision3.reset()
                    train_atom_targets_recall3.reset()

                    train_atom_types_acc.reset()
                    train_atom_charges_acc.reset()
                    train_atom_hs_acc.reset()

                    train_bond_targets_precision.reset()
                    train_bond_targets_recall.reset()
                    train_bond_targets_precision3.reset()
                    train_bond_targets_recall3.reset()

                    train_bond_types_acc.reset()
                    train_bond_rhos_mae.reset()

                    train_bond_omega_precision.reset()
                    train_bond_omega_recall.reset()
                    train_bond_omega_precision3.reset()
                    train_bond_omega_recall3.reset()

                    train_steps = 0

                    test_atom_targets_precision = AverageMeter()
                    test_atom_targets_recall = AverageMeter()
                    test_atom_targets_precision3 = AverageMeter()
                    test_atom_targets_recall3 = AverageMeter()

                    test_atom_types_acc = AverageMeter()
                    test_atom_charges_acc = AverageMeter()
                    test_atom_hs_acc = AverageMeter()

                    test_bond_targets_precision = AverageMeter()
                    test_bond_targets_recall = AverageMeter()
                    test_bond_targets_precision3 = AverageMeter()
                    test_bond_targets_recall3 = AverageMeter()

                    test_bond_types_acc = AverageMeter()

                    test_bond_rhos_mae = AverageMeter()

                    test_bond_omega_precision = AverageMeter()
                    test_bond_omega_recall = AverageMeter()
                    test_bond_omega_precision3 = AverageMeter()
                    test_bond_omega_recall3 = AverageMeter()

                    count = 0
                    for j, (imgs, atom_targets, atom_types, atom_charges, atom_hs,
                            bond_targets, bond_types, bond_rhos, bond_omega_types) in enumerate(test_dataloader):
                        count += 1
                        imgs = imgs.to(device)
                        atom_targets, atom_types, atom_charges, atom_hs = atom_targets.to(device), atom_types.to(
                            device), atom_charges.to(device), atom_hs.to(device)
                        bond_targets, bond_types, bond_rhos, bond_omega_types = bond_targets.to(
                            device), bond_types.to(device), bond_rhos.to(
                            device), bond_omega_types.to(device)

                        atom_targets_pred, atom_types_pred, atom_charges_pred, atom_hs_pred, bond_targets_pred, bond_types_pred, bond_rhos_pred, bond_omega_types_pred = model(
                            imgs)
                        atom_targets_pred = torch.clamp(torch.sigmoid(atom_targets_pred), 1e-5, 1 - 1e-5)
                        atom_types_pred = torch.clamp(torch.softmax(atom_types_pred, dim=1), 1e-5, 1 - 1e-5)
                        atom_charges_pred = torch.clamp(torch.softmax(atom_charges_pred, dim=1), 1e-5, 1 - 1e-5)

                        atom_hs_pred = torch.clamp(torch.softmax(atom_hs_pred, dim=1), 1e-5, 1 - 1e-5)

                        bond_targets_pred = torch.clamp(torch.sigmoid(bond_targets_pred), 1e-5, 1 - 1e-5)
                        bond_types_pred = torch.clamp(torch.softmax(bond_types_pred.view(-1, 6, 60, 128, 128), dim=1), 1e-5,
                                                      1 - 1e-5)

                        bond_omega_types_pred = torch.clamp(torch.sigmoid(bond_omega_types_pred.view(-1, 60, 128, 128)),
                                                            1e-5, 1 - 1e-5)
                        bond_rhos_pred = torch.abs(bond_rhos_pred)

                        temp = torch.nn.functional.max_pool2d(atom_targets_pred, kernel_size=3,
                                                              stride=1, padding=1)
                        atom_targets_pred = (temp == atom_targets_pred) * (atom_targets_pred > 0.25).float()

                        temp = torch.nn.functional.max_pool2d(bond_targets_pred, kernel_size=3,
                                                              stride=1, padding=1)
                        bond_targets_pred = (temp == bond_targets_pred) * (bond_targets_pred > 0.25).float()

                        atom_targets = (atom_targets == 1).float()
                        test_atom_targets_precision.update(
                            ((atom_targets_pred * atom_targets).sum() / atom_targets_pred.sum()).cpu().detach().numpy(),
                            atom_targets_pred.sum().cpu().detach().numpy())
                        test_atom_targets_precision3.update(
                            ((atom_targets_pred * torch.nn.functional.max_pool2d(atom_targets, kernel_size=3, padding=1,
                                                                                stride=1)).sum() / atom_targets_pred.sum()).cpu().detach().numpy(),
                            atom_targets_pred.sum().cpu().detach().numpy())

                        test_atom_targets_recall.update(
                            ((atom_targets * atom_targets_pred).sum() / atom_targets.sum()).cpu().detach().numpy(),
                            atom_targets.sum().cpu().detach().numpy())
                        test_atom_targets_recall3.update(
                            ((atom_targets * torch.nn.functional.max_pool2d(atom_targets_pred, kernel_size=3, padding=1,
                                                                            stride=1)).sum() / atom_targets.sum()).cpu().detach().numpy(),
                            atom_targets.sum().cpu().detach().numpy())

                        test_atom_types_acc.update((torch.sum(torch.sum(atom_types, dim=1) * ((atom_types.argmax(1) ==
                                                                                atom_types_pred.argmax(
                                                                                1)).float())) / torch.sum(
                            atom_types)).cpu().detach().numpy(), torch.sum(atom_types).cpu().detach().numpy())

                        test_atom_charges_acc.update((torch.sum(torch.sum(atom_charges, dim=1) * ((atom_charges.argmax(1) ==
                                                                                atom_charges_pred.argmax(
                                                                                1)).float())) / torch.sum(
                            atom_charges)).cpu().detach().numpy(), torch.sum(atom_charges).cpu().detach().numpy())

                        test_atom_hs_acc.update((torch.sum(torch.sum(atom_hs, dim=1) * ((atom_hs.argmax(1) ==
                                                                                atom_hs_pred.argmax(
                                                                                1)).float())) / (
                                                             0.01 + torch.sum(atom_hs))).cpu().detach().numpy(),
                                                (0.01 + torch.sum(atom_hs)).cpu().detach().numpy())

                        bond_targets = (bond_targets == 1).float()
                        test_bond_targets_precision.update(
                            ((bond_targets_pred * bond_targets).sum() / bond_targets_pred.sum()).cpu().detach().numpy(),
                            bond_targets_pred.sum().cpu().detach().numpy())
                        test_bond_targets_precision3.update(
                            ((bond_targets_pred * torch.nn.functional.max_pool2d(bond_targets, kernel_size=3, padding=1,
                                                                                stride=1)).sum() / bond_targets_pred.sum()).cpu().detach().numpy(),
                            bond_targets_pred.sum().cpu().detach().numpy())

                        test_bond_targets_recall.update(
                            ((bond_targets * bond_targets_pred).sum() / bond_targets.sum()).cpu().detach().numpy(),
                            bond_targets.sum().cpu().detach().numpy())
                        test_bond_targets_recall3.update(
                            ((bond_targets * torch.nn.functional.max_pool2d(bond_targets_pred, kernel_size=3, padding=1,
                                                                            stride=1)).sum() / bond_targets.sum()).cpu().detach().numpy(),
                            bond_targets.sum().cpu().detach().numpy())

                        test_bond_types_acc.update((torch.sum(torch.sum(bond_types, dim=1) * (
                            (bond_types.argmax(1) == bond_types_pred.argmax(1)).float())) / torch.sum(
                            bond_types)).cpu().detach().numpy(), torch.sum(bond_types).cpu().detach().numpy())

                        test_bond_rhos_mae.update(
                            (torch.sum(torch.abs(bond_rhos_pred - bond_rhos) * torch.sum(bond_types, dim=1)) / \
                             torch.sum(bond_types)).cpu().detach().numpy(), torch.sum(bond_types).cpu().detach().numpy())

                        temp = torch.cat(
                            [bond_omega_types_pred[:, 59:], bond_omega_types_pred, bond_omega_types_pred[:, :1]],
                            dim=1).permute(0, 2, 3, 1).reshape(-1, 128 * 128, 62)
                        temp = ((torch.nn.functional.max_pool1d(temp, stride=1, kernel_size=3, padding=0).reshape(-1, 128,
                                                                                128,
                                                                                60).permute(
                            0, 3, 1, 2) == bond_omega_types_pred) * \
                                (bond_omega_types_pred > 0.25)).float() * bond_targets

                        bond_omega_types = (bond_omega_types == 1)
                        test_bond_omega_precision.update(
                            (((bond_omega_types) * temp).sum() / (temp).sum()).cpu().detach().numpy(),
                            temp.sum().cpu().detach().numpy())

                        temp2 = torch.cat([temp[:, 59:], temp, temp[:, :1]],
                                          dim=1).permute(0, 2, 3, 1).reshape(-1, 128 * 128, 62)
                        temp2 = (
                            torch.nn.functional.max_pool1d(temp2.float(), stride=1, kernel_size=3, padding=0).reshape(-1,
                                                                                128,
                                                                                128,
                                                                                60).permute(
                                0, 3, 1, 2)).float()

                        test_bond_omega_recall3.update(
                            (((bond_omega_types) * temp2).sum() / (bond_omega_types).sum()).cpu().detach().numpy(),
                            bond_omega_types.sum().cpu().detach().numpy())

                        test_bond_omega_recall.update(
                            (((bond_omega_types) * temp).sum() / (bond_omega_types).sum()).cpu().detach().numpy(),
                            bond_omega_types.sum().cpu().detach().numpy())

                        temp3 = torch.cat([bond_omega_types[:, 59:], bond_omega_types, bond_omega_types[:, :1]],
                                          dim=1).permute(0, 2, 3, 1).reshape(-1, 128 * 128, 62)
                        temp3 = (
                            torch.nn.functional.max_pool1d(temp3.float(), stride=1, kernel_size=3, padding=0).reshape(-1,
                                                                                128,
                                                                                128,
                                                                                60).permute(
                                0, 3, 1, 2)).float()

                        test_bond_omega_precision3.update((((temp3) * temp).sum() / (temp).sum()).cpu().detach().numpy(),
                                                          temp.sum().cpu().detach().numpy())

                    if local_rank==0:
                        print('test_________________')

                        print('atom_target_precision:', test_atom_targets_precision.avg)
                        print('atom_target_recall:', test_atom_targets_recall.avg)
                        print('atom_target_precision3:', test_atom_targets_precision3.avg)
                        print('atom_target_recall3:', test_atom_targets_recall3.avg)

                        print('atom_types_acc:', test_atom_types_acc.avg)
                        print('atom_charges_acc:', test_atom_charges_acc.avg)
                        print('atom_hs_acc:', test_atom_hs_acc.avg)

                        print('bond_target_precision:', test_bond_targets_precision.avg)
                        print('bond_target_recall:', test_bond_targets_recall.avg)
                        print('bond_target_precision3:', test_bond_targets_precision3.avg)
                        print('bond_target_recall3:', test_bond_targets_recall3.avg)

                        print('bond_types_acc:', test_bond_types_acc.avg)
                        print('bond_rhos_mae:', test_bond_rhos_mae.avg)

                        print('bond_omega_precision:', test_bond_omega_precision.avg)
                        print('bond_omega_recall:', test_bond_omega_recall.avg)
                        print('bond_omega_precision3:', test_bond_omega_precision3.avg)
                        print('bond_omega_recall3:', test_bond_omega_recall3.avg)

        if local_rank==0:
            torch.save(model.state_dict(), 'weights2_{}/unet_model_weights{}.pkl'.format(amount, epoch))

if __name__ == '__main__':
    import time
    start=time.time()
    main()
    print('run time:{}'.format(time.time()-start))


'''
This class manipulate OpenGL drawing

camera control and perspective control are inspired by https://learnopengl.com/#!Getting-started/Camera
'''

import wx
from wx import glcanvas
from OpenGL.GL import *
from OpenGL.GLUT import *
from OpenGL.GLU import *
import numpy as np
from BeeVeeH.bvh_helper import BVHNode
from BeeVeeH.bvh_render import RENDER_CONFIG, render

def normalize(v):
    norm = np.linalg.norm(v)
    if norm == 0:
       return v
    return v / norm

class BeeVeeHCanvas(glcanvas.GLCanvas):

    context = None

    def __init__(self, parent_window):
        super().__init__(parent_window, -1)
        self.parent_window = parent_window

        self.init = False
        self.context = glcanvas.GLContext(self)

        self.lastx = self.x = 30
        self.lasty = self.y = 30
        self.size = None

        self.SetBackgroundStyle(wx.BG_STYLE_PAINT)

        self.Bind(wx.EVT_SIZE, self.OnSize)
        self.Bind(wx.EVT_PAINT, self.OnPaint)
        self.Bind(wx.EVT_LEFT_DOWN, self.OnLeftMouseDown)
        self.Bind(wx.EVT_LEFT_UP, self.OnLeftMouseUp)
        self.Bind(wx.EVT_RIGHT_DOWN, self.OnRightMouseDown)
        self.Bind(wx.EVT_RIGHT_UP, self.OnRightMouseUp)
        self.Bind(wx.EVT_MOTION, self.OnMouseMotion)
        self.Bind(wx.EVT_MOUSEWHEEL, self.OnMouseWheel)

        self.bvh_roots = []

        # camera control
        self.camera_position = np.array([200.0, 200.0, 200.0])
        self.camera_lookat = [0.0, 0.0, 0.0]
        self.camera_up = [0.0, 1.0, 0.0]
        self.camera_lookat_md = [0.0, 0.0, 0.0] # mouse down
        self.camera_up_md = [0.0, 1.0, 0.0] # mouse down

        # perspective control
        self.fov = 45.0

        # mouse info
        self.last_x = -1;
        self.last_y = -1;
        self.yaw = -135.0
        self.pitch = -45.0

        self.RENDER_CONFIG = RENDER_CONFIG
        self.ground_texture = None

        self.SculptureInterval = 3
        self.SculptureMaxFrames = 20

    def OnSize(self, event):
        self.size = self.GetClientSize()
        wx.CallAfter(self.DoSetViewport)
        wx.CallAfter(self.update_perspective)
        event.Skip()

    def DoSetViewport(self):
        if self.size is None:
            self.size = self.GetClientSize()
        self.SetCurrent(self.context)
        glViewport(0, 0, self.size.width, self.size.height)

    def OnPaint(self, event):
        dc = wx.PaintDC(self)
        self.SetCurrent(self.context)
        if not self.init:
            self.InitGL()
            self.init = True
        if not self.ground_texture:
            self.init_ground_texture()
        self.OnDraw()

    def calculate_camera(self):
        front = np.array([np.cos(np.radians(self.yaw)) * np.cos(np.radians(self.pitch)),
                          np.sin(np.radians(self.pitch)),
                          np.sin(np.radians(self.yaw)) * np.cos(np.radians(self.pitch))])
        self.camera_front = normalize(front)

    def update_camera(self):
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        gluLookAt(self.camera_position[0],
                  self.camera_position[1],
                  self.camera_position[2],
                  self.camera_position[0] + self.camera_front[0],
                  self.camera_position[1] + self.camera_front[1],
                  self.camera_position[2] + self.camera_front[2],
                  self.camera_up[0],
                  self.camera_up[1],
                  self.camera_up[2])

    def update_perspective(self):
        if self.size is None:
            self.size = self.GetClientSize()
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(self.fov, self.size.width / self.size.height, 10, 800.0);
        glMatrixMode(GL_MODELVIEW) # reset to default

    def InitGL(self):
        glutInit()

        # set viewing projection
        glShadeModel(GL_SMOOTH)

        self.update_perspective()
        self.calculate_camera()
        self.update_camera()

        glEnable(GL_DEPTH_TEST)
        glEnable(GL_LIGHTING)
        glEnable(GL_LIGHT0)

    def init_ground_texture(self):
        self.SetCurrent(self.context)

        dimension = 1200
        data = np.ndarray((dimension, dimension))
        for i in range(dimension):
            for j in range(dimension):
                if (i // (dimension / 20) + j // (dimension / 20)) % 2 == 0:
                    data[i, j] = 255
                else:
                    data[i, j] = 127
        data = np.uint8(data)
        print(data)
        print(data.shape)
        w, h = data.shape

        self.ground_texture = glGenTextures(1)
        glBindTexture(GL_TEXTURE_2D, self.ground_texture)

        glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_MODULATE)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)

        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, w, h, 0,
                     GL_LUMINANCE, GL_UNSIGNED_BYTE, data)

    def draw_ground(self):
        glEnable(GL_TEXTURE_2D)
        glBindTexture(GL_TEXTURE_2D, self.ground_texture)
        glColor3f(1, 1, 1)

        length = 400.0

        glBegin(GL_QUADS)
        glTexCoord2f(0, 1)
        glVertex3f(-length / 2, 0, length / 2)
        glTexCoord2f(0, 0)
        glVertex3f(-length / 2, 0, -length / 2)
        glTexCoord2f(1, 0)
        glVertex3f(length / 2, 0, -length / 2)
        glTexCoord2f(1, 1)
        glVertex3f(length / 2, 0, length / 2)
        glEnd()

        glDisable(GL_TEXTURE_2D)

    def OnDraw(self):
        # clear color and depth buffers
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

        self.draw_ground()

        n_frame = len(self.bvh_roots)
        if n_frame > 0:
            sculpture_roots = self.bvh_roots[::self.SculptureInterval]
            for bvh_root in sculpture_roots:
                render(bvh_root)
            if sculpture_roots[-1] != self.bvh_roots[-1]:
                render(self.bvh_roots[-1])

        if self.size is None:
            self.size = self.GetClientSize()

        self.SwapBuffers()

    def OnLeftMouseDown(self, event):
        self.CaptureMouse()
        self.last_x, self.last_y = event.GetPosition()
        self.mouse_button = 0;

    def OnLeftMouseUp(self, event):
        self.ReleaseMouse()
        self.mouse_button = -1

    def OnRightMouseDown(self, event):
        self.CaptureMouse()
        self.last_x, self.last_y = event.GetPosition()
        self.mouse_button = 1;

    def OnRightMouseUp(self, event):
        self.ReleaseMouse()
        self.mouse_button = -1

    def OnMouseMotion(self, event):
        if event.Dragging() and self.mouse_button >= 0:
            x, y = event.GetPosition()
            offset_x = x - self.last_x
            offset_y = y - self.last_y
            self.last_x, self.last_y = x, y
            if self.mouse_button == 0:
                sensitivity = 0.05
                offset_x *= sensitivity
                offset_y *= sensitivity
                self.yaw   -= offset_x
                self.pitch += offset_y
                if self.pitch > 89.0:
                    self.pitch = 89.0
                if self.pitch < -89.0:
                    self.pitch = -89.0
                self.calculate_camera()
            elif self.mouse_button == 1:
                camera_speed = 0.2
                self.camera_position -= normalize(np.cross(self.camera_front, self.camera_up)) * camera_speed * offset_x
                self.camera_position -= normalize(np.cross(self.camera_front, np.cross(self.camera_front, self.camera_up))) * camera_speed * offset_y
            self.update_camera()

    def OnMouseWheel(self, event):
        sensitivity = 0.1
        self.camera_position += self.camera_front * event.GetWheelRotation()
        self.update_camera()

    def show_bvh_frame(self, bvh_root, clean=True):
        if clean:
            self.bvh_roots = []
        if len(self.bvh_roots) > 0 and self.bvh_roots[-1] == bvh_root:
            return
        self.bvh_roots.append(bvh_root)
        if len(self.bvh_roots) >= self.SculptureInterval * (1 + self.SculptureMaxFrames):
            self.bvh_roots = self.bvh_roots[-self.SculptureInterval * self.SculptureMaxFrames:]


import math
import numpy as np
from arcade import shader, FRAGMENT_SHADER, VERTEX_SHADER
from arcade import Sprite as ArcadeSprite
from collections import deque
from PIL import Image, ImageFilter
from ..events import Event
from typing import Tuple

class GraphicDrawable:
    def __init__(self):
        self.visible = True

    def register_z_order(self, z_index, relative=True):
        self.z_order = z_index
        self.z_order_relative = relative

class GLShader:
    _PROGRAM = None

    @staticmethod
    def get_program():
        if GLShader._PROGRAM is None:
            GLShader._PROGRAM = shader.program(vertex_shader=VERTEX_SHADER,
                                               fragment_shader=FRAGMENT_SHADER)
        return GLShader._PROGRAM

    @staticmethod
    def get_triangle_strip_vertices():
        return np.array([-1.0, -1.0, 0.0, 0.0,
                         -1.0, 1.0, 0.0, 1.0,
                         1.0, -1.0, 1.0, 0.0,
                         1.0, 1.0, 1.0, 1.0],
                        dtype=np.float32)

class DrawableLayer:
    def __init__(self, maximum=500):
        self.effects = None
        self.visible = False
        self.texture_change = Event(self)
        self._queue = deque(maxlen=maximum)
        self._group_queue = deque(maxlen=maximum)
        self._tex_changed_flag = False
        self._merged = False
        self.vao = None
        self._img_info = None
        self.sprite_data = None
        self.gl_data_buf = None
        self.gl_texture = None
        self._tex_imgs = None
        self.comp_img = None
        self.merged_tex = []
        self.program = GLShader.get_program()

    def draw(self):
        self.visible = True

    def queue(self, item: GraphicDrawable, cached=False):
        index = len(self._queue)
        # item.register_z_order(index)
        self._queue.append(item)
        if cached:
            self.on_texture_changed("add", item)

    def queue_group(self, item, cached=False):
        item.texture_change += self.on_group_changed
        for i in item._queue:
            self.queue(i)
        if cached:
            self.on_texture_changed("add", item)

    def remove(self, item: GraphicDrawable):
        if item in self._queue:
            self._queue.remove(item)
        self.on_texture_changed("remove", item)

    def remove_group(self, item):
        item.texture_change -= self.on_group_changed
        for i in item._queue:
            self.remove(i)
#
# OpenGL texture creation wrapper functions
#

    def _extract_textures(self) -> Tuple[list, dict, list]:
        atlas_map = [(sprite, sprite._texture.name) for sprite in self._queue]
        tex_map = dict([(sprite._texture.name, sprite._texture.image) for sprite in self._queue])
        texture_images = [tex_map[tex_name] for tex_name in tex_map]
        return atlas_map, tex_map, texture_images

    def _process_textures(self, images: list):

        def _merge_down_to_image(total_width, max_height, images):
            #
            # SLUGGISH REGION: REDUCE AS MANY IMAGES AS POSSIBLE!
            #
            new_image = Image.new('RGBA', (total_width, max_height))  # Make a composite image
            x_offset = 0
            for image in images:
                new_image.paste(image, (x_offset, 0))
                x_offset += image.size[0]
            return new_image
            #
            #
            #

        widths, heights = zip(*(i.size for i in images))
        total_width = sum(widths)
        max_height = max(heights)
        return total_width, max_height, self._apply_effects(_merge_down_to_image(total_width, max_height, images))

    def _map_textures_to_atlas(self, total_width, max_height, atlas: list, tex_map, normalized=True):
        tex_coords = {}
        start_x = 0.0
        if not normalized:
            total_width, max_height = 1, 1
        for tex_name, image in tex_map.items():
            end_x = start_x + (image.width / total_width)
            normalized_width = image.width / total_width
            start_height = 1 - (image.height / max_height)
            normalized_height = image.height / max_height
            tex_coords[tex_name] = [start_x, start_height, normalized_width, normalized_height]
            start_x = end_x

        # Go through each sprite and pull from the coordinate list, the proper
        # coordinates for that sprite's image.
        tex_atlas_coords = [tex_coords[tex_name] for sprite, tex_name in atlas]
        return tex_atlas_coords

    def _cache_vao(self, ref_sprites, total_width, max_height, atlas, tex_map):
        sprite_count = len(ref_sprites)
        sprite_positions = []
        sprite_colors = []
        sprite_angles = []
        sprite_sizes = []
        for sprite in ref_sprites:
            sprite_positions.append([sprite.center_x, sprite.center_y])
            sprite_colors.append(sprite.color + (sprite.alpha, ))
            sprite_angles.append(math.radians(sprite.angle))
            sprite_sizes.append([sprite.width / 2, sprite.height / 2])

        sprite_coords = self._map_textures_to_atlas(total_width, max_height, atlas, tex_map)

        #
        # Creating OpenGL data buffer
        #
        sprite_data = np.zeros(sprite_count,
                               dtype=np.dtype([('position', '2f4'), ('angle', 'f4'),
                               ('size', '2f4'), ('sub_tex_coords', '4f4'),
                               ('color', '4B')]))
        sprite_data['position'] = sprite_positions
        sprite_data['angle'] = sprite_angles
        sprite_data['size'] = sprite_sizes
        sprite_data['sub_tex_coords'] = sprite_coords
        sprite_data['color'] = sprite_colors
        self.sprite_data = sprite_data.tobytes()
        self.gl_data_buf = shader.buffer(self.sprite_data,
                                         usage='stream')
        #
        #
        #

        self.vbo_buf = shader.buffer(GLShader.get_triangle_strip_vertices().tobytes())
        vbo_buf_desc = shader.BufferDescription(
            self.vbo_buf,
            '2f 2f',
            ('in_vert', 'in_texture')
        )
        pos_angle_scale_buf_desc = shader.BufferDescription(
            self.gl_data_buf,
            '2f 1f 2f 4f 4B',
            ('in_pos', 'in_angle', 'in_scale', 'in_sub_tex_coords', 'in_color'),
            normalized=['in_color'], instanced=True)

        vao_content = [vbo_buf_desc, pos_angle_scale_buf_desc]
        self.vao = shader.vertex_array(self.program, vao_content)

    def regenerate_vao_on_the_fly(self):
        self._cache_vao(*self._img_info)

    def rasterize(self, cached=True):
        """
            Convert down from sprite outlines to a single pixelated image\n
            The image also acts as an intermediate for a VisualRenderer
        """
        sprite_count = len(self._queue)
        if sprite_count == 0:
            return
        atlas, tex_map, tex_imgs = self._extract_textures()
        reset_flag = False
        if self._img_info is not None:
            if not set(self._img_info[2]).difference(atlas) == set():
                reset_flag = True
            else:
                self.gl_texture = shader.texture((self.comp_img.width, self.comp_img.height),
                                                    4, np.asarray(self.comp_img))
                if self._merged:
                    q = self._queue
                else:
                    q = self.merged_tex
                self._cache_vao(q, self._img_info[0], self._img_info[1], atlas, tex_map)
        else:
            reset_flag = True

        if reset_flag:
            total_width, max_height, comp_img = self._process_textures(tex_imgs)
            self.gl_texture = shader.texture((comp_img.width, comp_img.height),
                                              4, np.asarray(comp_img))
            if cached:
                self.comp_img = comp_img
                self._cache_vao(self._queue, comp_img.width, comp_img.height, atlas, tex_map)
                self._img_info = comp_img.width, comp_img.height, atlas
                self._tex_imgs = tex_imgs

    def merge_down(self):
        self.rasterize()
        positions = [(sprite.center_x, sprite.center_y) for sprite in self._queue]
        range = [(sprite.left, sprite.right) for sprite in self._queue]
        max_left = max([l for l, r in range])
        max_right = max([r for l, r in range])
        pos_x = [pos[0] for pos in positions]
        pos_y = [pos[1] for pos in positions]
        min_pos = min([x for x, y in positions]), min([y for x, y in positions])
        max_pos = max([x for x, y in positions]), max([y for x, y in positions])
        ind = -1

        tex_map = dict([(sprite._texture.name, sprite._texture.image) for sprite in self._queue])
        tex_coords = self._map_textures_to_atlas(self.comp_img.width, self.comp_img.height, self._img_info[2],
                                                 tex_map, normalized=False)
        x_offset, y_offset = 0, 0
        if min_pos[0] < 0:
            x_offset = abs(min_pos[0])
        if min_pos[1] < 0:
            y_offset = abs(min_pos[0])

        width = int(max_pos[0] + self._queue[pos_x.index(max_pos[0])].width - min_pos[0])
        height = int(max_pos[1] + self._queue[pos_x.index(max_pos[1])].height - min_pos[1])

        new_image = Image.new('RGBA', (width, height))  # Make a composite image
        for sprite in self._queue:
            new_image.paste(sprite.texture.image, (sprite.center_x + x_offset, sprite.center_y + y_offset), sprite.texture.image)
        new_image.save('test.png')

        temp_spr = Sprite()
        temp_spr.set_position(min_pos[0], min_pos[1])
        self.merged_tex = []
        self.merged_tex.append(temp_spr)

    def _apply_effects(self, image: Image) -> Image:
        return image.filter(ImageFilter.GaussianBlur(1.25))

    def update(self):
        pass

    def on_group_changed(self, sender, *args):
        command, item = args[0], args[1]
        if command == "add":
            sender.queue(args[1])
        elif command == "remove":
            sender.remove(args[1])

    def on_texture_changed(self, *args):
        self.texture_change(*args)
        self.rasterize()

class Sprite(GraphicDrawable, ArcadeSprite):
    def __init__(self, filename: str=None, scale: float=1, image_x: float=0, image_y: float=0,
                 image_width: float=0, image_height: float=0,
                 center_x: float=0, center_y: float=0,
                 repeat_count_x=1, repeat_count_y=1):
        GraphicDrawable.__init__(self)
        ArcadeSprite.__init__(self, filename, scale,
                              image_x, image_y,
                              image_width, image_height,
                              center_x, center_y,
                              repeat_count_x, repeat_count_y)