import random
from collections import namedtuple, deque

import numpy as np
import torch
import torch.nn
import torch.optim as optim
import torch.nn.functional as F
from Model import QNetwork

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class Agent():
    def __init__(self, bufferSize, batchSize, gamma, tau, lr, updateEvery, state_size, action_size, seed):
        print(device)
        # hyper parameters
        self.BufferSize = bufferSize
        self.BatchSize = batchSize
        self.Gamma = gamma
        self.Tau = tau
        self.LR = lr
        self.UpdateEvery = updateEvery
        self.State_size = state_size
        self.Action_size = action_size
        self.Seed = random.seed(seed)

        # Q network
        self.QNetwork_local = QNetwork(state_size, action_size, seed).to(device)
        self.QNetwork_target = QNetwork(state_size, action_size, seed).to(device)
        self.Optimizer = optim.Adam(self.QNetwork_local.parameters(), lr=self.LR)

        # Replay memory
        self.memory = ReplayBuffer(action_size, self.BufferSize, self.BatchSize, seed)
        # Initialize time step (for updating every UPDATE_EVERY steps)
        self.t_step = 0

    def act(self, state, eps=0.):
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        self.QNetwork_local.eval()
        with torch.no_grad():
            action_values = self.QNetwork_local(state)
        self.QNetwork_local.train()

        # Epsilon-greedy action selection
        if random.random() > eps:
            return np.argmax(action_values.cpu().data.numpy())
        else:
            return random.choice(np.arange(self.Action_size))

    def step(self, state, action, reward, next_state, done):
        # Save experience in replay memory
        self.memory.add(state, action, reward, next_state, done)

        # Learn every UPDATE_EVERY time steps.
        self.t_step = (self.t_step + 1) % self.UpdateEvery
        if self.t_step == 0:
            # If enough samples are available in memory, get random subset and learn
            if len(self.memory) > self.BatchSize:
                experiences = self.memory.sample()
                self.learn(experiences, self.Gamma)

    def learn(self, experiences, gamma):
        states, actions, rewards, next_states, dones = experiences
        q_targets_next = self.QNetwork_target(next_states).detach().max(1)[0].unsqueeze(1)
        # Compute Q targets for current states
        q_targets = rewards + (gamma * q_targets_next * (1 - dones))

        # Get expected Q values from local model
        q_expected = self.QNetwork_local(states).gather(1, actions)

        # Compute loss
        loss = F.mse_loss(q_expected, q_targets)
        # Minimize the loss
        self.Optimizer.zero_grad()
        loss.backward()
        self.Optimizer.step()

        # ------------------- update target network ------------------- #
        self.soft_update(self.QNetwork_local, self.QNetwork_target, self.Tau)

    def soft_update(self, local_model, target_model, tau):
        for target_param, local_param in zip(target_model.parameters(), local_model.parameters()):
            target_param.data.copy_(tau*local_param.data + (1.0-tau)*target_param.data)

class ReplayBuffer:
    def __init__(self, action_size, buffer_size, batch_size, seed):
        self.Action_size = action_size
        self.Memory = deque(maxlen=buffer_size)
        self.Batch_size = batch_size
        self.Experience = namedtuple("Experience", field_names=["state", "action", "reward", "next_state", "done"])
        self.Seed = random.seed(seed)

    def add(self, state, action, reward, next_state, done):
        e = self.Experience(state, action, reward, next_state, done)
        self.Memory.append(e)

    def sample(self):
        experiences = random.sample(self.Memory, k=self.Batch_size)
        states = torch.from_numpy(np.vstack([e.state for e in experiences if e is not None])).float().to(device)
        actions = torch.from_numpy(np.vstack([e.action for e in experiences if e is not None])).long().to(device)
        rewards = torch.from_numpy(np.vstack([e.reward for e in experiences if e is not None])).float().to(device)
        next_states = torch.from_numpy(np.vstack([e.next_state for e in experiences if e is not None])).float().to(
            device)
        dones = torch.from_numpy(np.vstack([e.done for e in experiences if e is not None]).astype(np.uint8)).float().to(
            device)
        return states, actions, rewards, next_states, dones

    def __len__(self):
        return len(self.Memory)


#judges a picture that is taken and stored

import numpy as np
import cv2
import random

def detectHappyOrPeace():
    pass

def detectMood(data, imageInfo):
    face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
    eye_cascade = cv2.CascadeClassifier('haarcascade_eye.xml')
    smile_cascade = cv2.CascadeClassifier('haarcascade_smile.xml')

    #picture stored in test.png
    img = cv2.imread('test.png')
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    faces = face_cascade.detectMultiScale(gray, 1.3, 5)

    if len(faces) != 1:  #make sure only one face detected
        # print("toomanyfaces")
        return None

    x = faces[0][0]
    y = faces[0][1]
    w = faces[0][2]
    h = faces[0][3]
    #limits search for smile/eyes to face
    img = cv2.rectangle(img,(x,y),(x+w,y+h),(255,0,0),2)
    roi_gray = gray[y:y+h, x:x+w]
    roi_color = img[y:y+h, x:x+w]

    #uses gray face to find smile
    smile = smile_cascade.detectMultiScale(roi_gray, scaleFactor= 1.7,
        minNeighbors=22, minSize=(25, 25))
    if len(smile) > 0:
        sx = smile[0][0]
        sy = smile[0][1]
        sw = smile[0][2]
        sh = smile[0][3]
        smileImageColor = roi_color[sy:sy+sh, sx:sx+sw]
        # cv2.imshow('img',smileImageColor)

        #boundary of teeth color
        lower = np.array([150, 150, 150], dtype = "uint8")
        upper = np.array([255, 255, 255], dtype = "uint8")

        mask = cv2.inRange(smileImageColor, lower, upper)

        # cv2.imshow('mask',mask)
        # cv2.imshow('image', smileImageColor)
        output = cv2.bitwise_and(smileImageColor, smileImageColor, mask = mask)
        # cv2.imshow("images", np.hstack([smileImageColor, output]))

        #if there is white, it is a chance of happiness or peacefulness.
        if np.any(mask):
            # print("teeth might be detected")
            guess = random.randint(0, 10)
            if guess < 8:
                return 0
                # print("happiness guessed")
            else:
                return 3
                # print("peace guessed")
        else:
            # print("teeth not detected")
            guess = random.randint(0, 10)
            if guess < 8:
                return 3
            else:
                return 0

    #uses gray face to find eyes
    eyes = eye_cascade.detectMultiScale(roi_gray)
    sumEyeDist = 0
    for (ex,ey,ew,eh) in eyes:
        #find the two main eyes based on similarity of size and y level
        #find height of eyes as related to the picture
        cv2.rectangle(roi_color,(ex,ey),(ex+ew,ey+eh),(0,255,0),2)
        sumEyeDist += ey
    avgEyeDist = ey/2
    totalEyeDistFromTop = avgEyeDist + y + eh//2
    if totalEyeDistFromTop > data.height//2:
        # print("Sadness detected")
        return 1
    else:
        # print("Powerfulness detected")
        return 2


# coding: utf-8
from __future__ import division

import models
import data

import theano
import sys
import codecs

import theano.tensor as T
import numpy as np

MAX_SUBSEQUENCE_LEN = 200

def to_array(arr, dtype=np.int32):
    # minibatch of 1 sequence as column
    return np.array([arr], dtype=dtype).T

def convert_punctuation_to_readable(punct_token):
    if punct_token == data.SPACE:
        return " "
    else:
        return punct_token[0]

def restore_with_pauses(output_file, text, pauses, word_vocabulary, reverse_punctuation_vocabulary, predict_function):
    i = 0
    with codecs.open(output_file, 'w', 'utf-8') as f_out:
        while True:

            subsequence = text[i:i+MAX_SUBSEQUENCE_LEN]
            subsequence_pauses = pauses[i:i+MAX_SUBSEQUENCE_LEN]

            if len(subsequence) == 0:
                break

            converted_subsequence = [word_vocabulary.get(w, word_vocabulary[data.UNK]) for w in subsequence]

            y = predict_function(to_array(converted_subsequence), to_array(subsequence_pauses, dtype=theano.config.floatX))

            f_out.write(subsequence[0])

            last_eos_idx = 0
            punctuations = []
            for y_t in y:

                p_i = np.argmax(y_t.flatten())
                punctuation = reverse_punctuation_vocabulary[p_i]

                punctuations.append(punctuation)

                if punctuation in data.EOS_TOKENS:
                    last_eos_idx = len(punctuations) # we intentionally want the index of next element

            if subsequence[-1] == data.END:
                step = len(subsequence) - 1
            elif last_eos_idx != 0:
                step = last_eos_idx
            else:
                step = len(subsequence) - 1

            for j in range(step):
                f_out.write(" " + punctuations[j] + " " if punctuations[j] != data.SPACE else " ")
                if j < step - 1:
                    f_out.write(subsequence[1+j])

            if subsequence[-1] == data.END:
                break

            i += step

def restore(output_file, text, word_vocabulary, reverse_punctuation_vocabulary, predict_function):
    i = 0
    with codecs.open(output_file, 'w', 'utf-8') as f_out:
        while True:

            subsequence = text[i:i+MAX_SUBSEQUENCE_LEN]

            if len(subsequence) == 0:
                break

            converted_subsequence = [word_vocabulary.get(w, word_vocabulary[data.UNK]) for w in subsequence]

            y = predict_function(to_array(converted_subsequence))

            f_out.write(subsequence[0])

            last_eos_idx = 0
            punctuations = []
            for y_t in y:

                p_i = np.argmax(y_t.flatten())
                punctuation = reverse_punctuation_vocabulary[p_i]

                punctuations.append(punctuation)

                if punctuation in data.EOS_TOKENS:
                    last_eos_idx = len(punctuations) # we intentionally want the index of next element

            if subsequence[-1] == data.END:
                step = len(subsequence) - 1
            elif last_eos_idx != 0:
                step = last_eos_idx
            else:
                step = len(subsequence) - 1

            for j in range(step):
                f_out.write(" " + punctuations[j] + " " if punctuations[j] != data.SPACE else " ")
                if j < step - 1:
                    f_out.write(subsequence[1+j])

            if subsequence[-1] == data.END:
                break

            i += step

if __name__ == "__main__":

    if len(sys.argv) > 1:
        model_file = sys.argv[1]
    else:
        sys.exit("Model file path argument missing")

    if len(sys.argv) > 2:
        output_file = sys.argv[2]
    else:
        sys.exit("Output file path argument missing")

    use_pauses = len(sys.argv) > 3 and bool(int(sys.argv[3]))

    x = T.imatrix('x')

    if use_pauses:

        p = T.matrix('p')

        print "Loading model parameters..."
        net, _ = models.load(model_file, 1, x, p)

        print "Building model..."
        predict = theano.function(
            inputs=[x, p],
            outputs=net.y
        )

    else:

        print "Loading model parameters..."
        net, _ = models.load(model_file, 1, x)

        print "Building model..."
        predict = theano.function(
            inputs=[x],
            outputs=net.y
        )

    word_vocabulary = net.x_vocabulary
    punctuation_vocabulary = net.y_vocabulary

    reverse_word_vocabulary = {v:k for k,v in word_vocabulary.items()}
    reverse_punctuation_vocabulary = {v:k for k,v in punctuation_vocabulary.items()}

    input_text = codecs.getreader('utf-8')(sys.stdin).read()

    if len(input_text) == 0:
        sys.exit("Input text from stdin missing.")

    text = [w for w in input_text.split() if w not in punctuation_vocabulary and w not in data.PUNCTUATION_MAPPING and not w.startswith(data.PAUSE_PREFIX)] + [data.END]
    pauses = [float(s.replace(data.PAUSE_PREFIX,"").replace(">","")) for s in input_text.split() if s.startswith(data.PAUSE_PREFIX)]

    if not use_pauses:
        restore(output_file, text, word_vocabulary, reverse_punctuation_vocabulary, predict)
    else:
        if not pauses:
            pauses = [0.0 for _ in range(len(text)-1)]
        restore_with_pauses(output_file, text, pauses, word_vocabulary, reverse_punctuation_vocabulary, predict)


from os import listdir
from os.path import join, isfile
import random
from scipy import io as sio
import numpy as np
import copy
from ipdb import set_trace as st
from math import ceil

class BRATS_sub():
    def __init__(self,opt,phase):
        super(BRATS_sub, self).__init__()
        random.seed(0)
        self.dataroot = opt.dataroot
        self.root   = join(self.dataroot,phase)
        self.flist  = []

        self.root   = self.dataroot

        ## numpy file is a list of filename paths for each slice of brain images.
        ## For fast data loading, the brain images are saved and loaded by slice-by-slice .
        self.flist = np.load(join(self.dataroot, phase+'_flist_main_z.npy'))

        self.N = 4
        self.nZ      = 3  #150
        self.nCh_seg = 1
        self.nCh_in  = self.N
        self.nCh_out = self.N
        self.nY      = 240
        self.nX      = 240
        self.len     = len(self.flist)
        self.use_aug = (phase=='train') and opt.AUG
        self.tumor = opt.tumor

    def get_info(self,opt):
        opt.nCh_in = self.nCh_in
        opt.nCh_out= self.nCh_out
        opt.nCh_seg= self.nCh_seg
        opt.nY     = self.nY
        opt.nX     = self.nX
        opt.nZ     = self.nZ
        return opt

    def getBatch(self, start, end):
        nB = end-start
        end   = min([end,self.len])
        start = end-nB
        batch = self.flist[start:end]

        sz_p   = [end-start,self.nCh_in*2, self.nY, self.nX]
        sz_a   = [end-start,  self.nCh_in, self.nY, self.nX]
        sz_M   = [end-start, self.nCh_seg, self.nY, self.nX]

        input_pre = np.empty(sz_p, dtype=np.float32)
        input_img = np.empty(sz_a, dtype=np.float32)
        input_post = np.empty(sz_p,dtype=np.float32)
        target_img = np.empty(sz_M, dtype=np.float32)

        for iB, aBatch in enumerate(batch):
            i_str = len(aBatch.rstrip("0123456789"))
            z_s = int( aBatch[i_str:] )
            _Batch = aBatch[:i_str]

            input_pre[iB,0,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-2)+'O.mat'))
            input_pre[iB,1,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-2)+'C.mat'))
            input_pre[iB,2,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-2)+'T.mat'))
            input_pre[iB,3,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-2)+'F.mat'))
            input_pre[iB,4,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-1)+'O.mat'))
            input_pre[iB,5,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-1)+'C.mat'))
            input_pre[iB,6,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-1)+'T.mat'))
            input_pre[iB,7,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-1)+'F.mat'))

            input_img[iB,0,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s)+'O.mat'))
            input_img[iB,1,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s)+'C.mat'))
            input_img[iB,2,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s)+'T.mat'))
            input_img[iB,3,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s)+'F.mat'))

            input_post[iB,0,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+1)+'O.mat'))
            input_post[iB,1,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+1)+'C.mat'))
            input_post[iB,2,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+1)+'T.mat'))
            input_post[iB,3,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+1)+'F.mat'))
            input_post[iB,4,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+2)+'O.mat'))
            input_post[iB,5,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+2)+'C.mat'))
            input_post[iB,6,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+2)+'T.mat'))
            input_post[iB,7,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+2)+'F.mat'))

            tmp_S = self.read_mat(join(self.root,_Batch+str(z_s)+'S.mat'))
            target_img[iB,:,:,:] = tmp_S[:,:,self.tumor]

        if self.use_aug:
            scale = np.random.uniform(low=0.9, high=1.1)
            input_pre = input_pre*scale
            input_img = input_img*scale
            input_post= input_post*scale

            if random.randint(0,1):
                input_pre = np.flip(input_pre,axis=3)
                input_img = np.flip(input_img,axis=3)
                input_post = np.flip(input_post,axis=3)
                target_img = np.flip(target_img,axis=3)

        return input_pre, input_img, input_post, target_img

    def getBatch_WTTCEC(self, start, end):
        nB = end-start
        end   = min([end,self.len])
        start = end-nB
        batch = self.flist[start:end]

        sz_p   = [end-start,self.nCh_in*2, self.nY, self.nX]
        sz_a   = [end-start,  self.nCh_in, self.nY, self.nX]
        sz_M   = [end-start, self.nCh_seg, self.nY, self.nX]

        input_pre = np.empty(sz_p, dtype=np.float32)
        input_img = np.empty(sz_a, dtype=np.float32)
        input_post = np.empty(sz_p,dtype=np.float32)
        target_img = np.empty(sz_M, dtype=np.float32)

        for iB, aBatch in enumerate(batch):
            i_str = len(aBatch.rstrip("0123456789"))
            z_s = int( aBatch[i_str:] )
            _Batch = aBatch[:i_str]

            input_pre[iB,0,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-2)+'O.mat'))
            input_pre[iB,1,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-2)+'C.mat'))
            input_pre[iB,2,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-2)+'T.mat'))
            input_pre[iB,3,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-2)+'F.mat'))
            input_pre[iB,4,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-1)+'O.mat'))
            input_pre[iB,5,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-1)+'C.mat'))
            input_pre[iB,6,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-1)+'T.mat'))
            input_pre[iB,7,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s-1)+'F.mat'))

            input_img[iB,0,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s)+'O.mat'))
            input_img[iB,1,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s)+'C.mat'))
            input_img[iB,2,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s)+'T.mat'))
            input_img[iB,3,:,:]  = self.read_mat(join(self.root, _Batch+str(z_s)+'F.mat'))

            input_post[iB,0,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+1)+'O.mat'))
            input_post[iB,1,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+1)+'C.mat'))
            input_post[iB,2,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+1)+'T.mat'))
            input_post[iB,3,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+1)+'F.mat'))
            input_post[iB,4,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+2)+'O.mat'))
            input_post[iB,5,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+2)+'C.mat'))
            input_post[iB,6,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+2)+'T.mat'))
            input_post[iB,7,:,:] = self.read_mat(join(self.root, _Batch+str(z_s+2)+'F.mat'))

            tmp_S = self.read_mat(join(self.root,_Batch+str(z_s)+'S.mat'))
            target_img[iB,:,:,:] = np.transpose(tmp_S,(2,0,1))

        if self.use_aug:
            scale = np.random.uniform(low=0.9, high=1.1)
            input_pre = input_pre*scale
            input_img = input_img*scale
            input_post= input_post*scale

            if random.randint(0,1):
                input_pre = np.flip(input_pre,axis=3)
                input_img = np.flip(input_img,axis=3)
                input_post = np.flip(input_post,axis=3)
                target_img = np.flip(target_img,axis=3)

        return input_pre, input_img, input_post, target_img

    def shuffle(self, seed=0):
        random.seed(seed)
        random.shuffle(self.flist)

    def name(self):
        return 'BRATSseg dataset'

    def __len__(self):
        return self.len

    @staticmethod
    def read_mat(filename, var_name="img"):
        mat = sio.loadmat(filename)
        return mat[var_name]


import numpy as np

from config.atlas_config import WBCConfig, PnCConfig
from pnc.wbc.tci_container import TCIContainer
from pnc.wbc.basic_task import BasicTask
from pnc.wbc.basic_contact import SurfaceContact

class AtlasTCIContainer(TCIContainer):
    def __init__(self, robot):
        super(AtlasTCIContainer, self).__init__(robot)

        # ======================================================================
        # Initialize Task
        # ======================================================================
        # COM Task
        self._com_task = BasicTask(robot, "COM", 3, 'com', PnCConfig.SAVE_DATA)
        self._com_task.kp = WBCConfig.KP_COM
        self._com_task.kd = WBCConfig.KD_COM
        self._com_task.w_hierarchy = WBCConfig.W_COM
        # Pelvis Task
        self._pelvis_ori_task = BasicTask(robot, "LINK_ORI", 3, "pelvis_com",
                                          PnCConfig.SAVE_DATA)
        self._pelvis_ori_task.kp = WBCConfig.KP_PELVIS
        self._pelvis_ori_task.kd = WBCConfig.KD_PELVIS
        self._pelvis_ori_task.w_hierarchy = WBCConfig.W_PELVIS
        # Non leg related joints (18, so 18 + 6*2 = n_a)
        selected_joint = [
            "back_bkx", "back_bky", "back_bkz", "l_arm_elx", "l_arm_ely",
            "l_arm_shx", "l_arm_shz", "l_arm_wrx", "l_arm_wry", "l_arm_wry2",
            "neck_ry", "r_arm_elx", "r_arm_ely", "r_arm_shx", "r_arm_shz",
            "r_arm_wrx", "r_arm_wry", "r_arm_wry2"
        ]
        self._upper_body_task = BasicTask(robot, "SELECTED_JOINT",
                                          len(selected_joint), selected_joint,
                                          PnCConfig.SAVE_DATA)
        self._upper_body_task.kp = np.array([WBCConfig.KP_UPPER_BODY] *
                                            self._upper_body_task.dim)
        self._upper_body_task.kd = np.array([WBCConfig.KD_UPPER_BODY] *
                                            self._upper_body_task.dim)
        self._upper_body_task.w_hierarchy = WBCConfig.W_UPPER_BODY
        # Rfoot Pos Task
        self._rfoot_pos_task = BasicTask(robot, "LINK_XYZ", 3, "r_sole",
                                         PnCConfig.SAVE_DATA)
        self._rfoot_pos_task.kp = WBCConfig.KP_FOOT_POS
        self._rfoot_pos_task.kd = WBCConfig.KD_FOOT_POS
        self._rfoot_pos_task.w_hierarchy = WBCConfig.W_CONTACT_FOOT
        # Lfoot Pos Task
        self._lfoot_pos_task = BasicTask(robot, "LINK_XYZ", 3, "l_sole",
                                         PnCConfig.SAVE_DATA)
        self._lfoot_pos_task.kp = WBCConfig.KP_FOOT_POS
        self._lfoot_pos_task.kd = WBCConfig.KD_FOOT_POS
        self._lfoot_pos_task.w_hierarchy = WBCConfig.W_CONTACT_FOOT
        # Rfoot Ori Task
        self._rfoot_ori_task = BasicTask(robot, "LINK_ORI", 3, "r_sole",
                                         PnCConfig.SAVE_DATA)
        self._rfoot_ori_task.kp = WBCConfig.KP_FOOT_ORI
        self._rfoot_ori_task.kd = WBCConfig.KD_FOOT_ORI
        self._rfoot_ori_task.w_hierarchy = WBCConfig.W_CONTACT_FOOT
        # Lfoot Ori Task
        self._lfoot_ori_task = BasicTask(robot, "LINK_ORI", 3, "l_sole",
                                         PnCConfig.SAVE_DATA)
        self._lfoot_ori_task.kp = WBCConfig.KP_FOOT_ORI
        self._lfoot_ori_task.kd = WBCConfig.KD_FOOT_ORI
        self._lfoot_ori_task.w_hierarchy = WBCConfig.W_CONTACT_FOOT

        self._task_list = [
            self._com_task, self._pelvis_ori_task, self._upper_body_task,
            self._rfoot_pos_task, self._lfoot_pos_task, self._rfoot_ori_task,
            self._lfoot_ori_task
        ]

        # ======================================================================
        # Initialize Contact
        # ======================================================================
        # Rfoot Contact
        self._rfoot_contact = SurfaceContact(robot, "r_sole", 0.11, 0.065, 0.3,
                                             PnCConfig.SAVE_DATA)
        self._rfoot_contact.rf_z_max = 1e-3  # Initial rf_z_max
        # self._rfoot_contact.rf_z_max = 1  # Initial rf_z_max
        # Lfoot Contact
        self._lfoot_contact = SurfaceContact(robot, "l_sole", 0.11, 0.065, 0.3,
                                             PnCConfig.SAVE_DATA)
        self._lfoot_contact.rf_z_max = 1e-3  # Initial rf_z_max
        # self._lfoot_contact.rf_z_max = 1  # Initial rf_z_max

        self._contact_list = [self._rfoot_contact, self._lfoot_contact]

    @property
    def com_task(self):
        return self._com_task

    @property
    def pelvis_ori_task(self):
        return self._pelvis_ori_task

    @property
    def upper_body_task(self):
        return self._upper_body_task

    @property
    def rfoot_pos_task(self):
        return self._rfoot_pos_task

    @property
    def lfoot_pos_task(self):
        return self._lfoot_pos_task

    @property
    def rfoot_ori_task(self):
        return self._rfoot_ori_task

    @property
    def lfoot_ori_task(self):
        return self._lfoot_ori_task

    @property
    def rfoot_contact(self):
        return self._rfoot_contact

    @property
    def lfoot_contact(self):
        return self._lfoot_contact


import numpy as np
import os
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
from .agent import Agent
from ..utils.evaluation import entropy
from ..networks.critic import Critic
from ..networks.actor import Actor

class BCQ(Agent):

    def __init__(self,
                 obs_space,
                 action_space,
                 discount,
                 lr=1e-4,
                 seed=None):
        super(BCQ, self).__init__(obs_space, action_space, discount, lr, seed)

        # epsilon decay
        self.initial_eps = 1.0
        self.end_eps = 1e-2
        self.eps_decay_period = 1000
        self.slope = (self.end_eps - self.initial_eps) / self.eps_decay_period
        self.eval_eps = 0.

        # loss function
        self.huber = nn.SmoothL1Loss()
        self.nll = nn.NLLLoss()

        # threshold for actions, unlikely under the behavioral policy
        self.threshold = 0.3

        # Number of training iterations
        self.iterations = 0

        # After how many training steps 'snap' target to main network?
        self.target_update_freq = 100

        # Q-Networks
        self.Q = Critic(self.obs_space, self.action_space, seed).to(self.device)
        self.Q_target = copy.deepcopy(self.Q)
        # BCQ has a separate Actor, as I have no common vision network
        self.actor = Actor(self.obs_space, self.action_space, seed).to(self.device)

        # Optimization
        self.optimizer = torch.optim.Adam(params=list(self.Q.parameters())+list(self.actor.parameters()), lr=self.lr)

    def policy(self, state, eval=False):

        # set networks to eval mode
        self.actor.eval()
        self.Q.eval()

        if eval:
            eps = self.eval_eps
        else:
            eps = max(self.slope * self.iterations + self.initial_eps, self.end_eps)

        # epsilon greedy policy
        if self.rng.uniform(0, 1) > eps:
            with torch.no_grad():
                state = torch.FloatTensor(state).to(self.device)

                q_val = self.Q(state).cpu()
                actions = self.actor(state).cpu()

                sm = F.log_softmax(actions, dim=1).exp()
                mask = ((sm / sm.max(1, keepdim=True)[0]) > self.threshold).float()

                # masking non-eligible values with -9e9 to be sure they are not sampled
                return int((mask * q_val + (1. - mask) * -9e9).argmax(dim=1)), \
                       q_val, entropy(sm)
        else:
            return self.rng.integers(self.action_space), np.nan, np.nan

    def train(self, buffer, writer, minimum=None, maximum=None, use_probas=False):
        # Sample replay buffer
        state, action, next_state, reward, not_done = buffer.sample(minimum, maximum, use_probas)

        # set networks to train mode
        self.actor.train()
        self.Q.train()
        self.Q_target.train()

        with torch.no_grad():
            q_val = self.Q(next_state)
            actions = self.actor(next_state)

            sm = F.log_softmax(actions, dim=1).exp()
            sm = (sm / sm.max(1, keepdim=True)[0] > self.threshold).float()
            next_action = (sm * q_val + (1. - sm) * -9e9).argmax(dim=1, keepdim=True)

            q_val = self.Q_target(next_state)
            target_Q = reward + not_done * self.discount * q_val.gather(1, next_action).reshape(-1, 1)

        # Get current Q estimate and actor decisions on actions
        current_Q = self.Q(state).gather(1, action)
        actions = self.actor(state)

        # Compute Q loss (Huber loss)
        Q_loss = self.huber(current_Q, target_Q)
        A_loss = self.nll(F.log_softmax(actions, dim=1), action.reshape(-1))
        # third term is
        loss = Q_loss + A_loss + 1e-2 * actions.pow(2).mean()

        # log temporal difference error
        if self.iterations % 100 == 0:
            writer.add_scalar("train/TD-error", torch.mean(Q_loss).detach().cpu().item(), self.iterations)
            writer.add_scalar("train/CE-loss", torch.mean(A_loss).detach().cpu().item(), self.iterations)

        # Optimize the Q
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        self.iterations += 1
        # Update target network by full copy every X iterations.
        if self.iterations % self.target_update_freq == 0:
            self.Q_target.load_state_dict(self.Q.state_dict())

    def get_name(self) -> str:
        return "BatchConstrainedQLearning"

    def save_state(self) -> None:
        torch.save(self.Q.state_dict(), os.path.join("models", self.get_name() + "_Q.pt"))
        torch.save(self.actor.state_dict(), os.path.join("models", self.get_name() + "_actor.pt"))
        torch.save(self.optimizer.state_dict(), os.path.join("models", self.get_name() + "_optim.pt"))

    def load_state(self) -> None:
        self.Q.load_state_dict(torch.load(os.path.join("models", self.get_name() + "_Q.pt")))
        self.Q_target = copy.deepcopy(self.Q)
        self.actor.load_state_dict(torch.load(os.path.join("models", self.get_name() + "_actor.pt")))
        self.optimizer.load_state_dict(torch.load(os.path.join("models", self.get_name() + "_optim.pt")))


from .base import BaseDataset
import os
import cv2
import numpy as np
from .DOTA_devkit.ResultMerge_multi_process import mergebypoly

class DOTA(BaseDataset):
    def __init__(self, data_dir, phase, input_h=None, input_w=None, down_ratio=None):
        super(DOTA, self).__init__(data_dir, phase, input_h, input_w, down_ratio)
        self.category = ['Dry-Cargo-Ship', 'Passenger-Ship', 'Fishing-Boat', 'Engineering-Ship', 'Motorboat', 'A220', 'Boeing737', 'Liquid-Cargo-Ship',\
 'Cargo-Truck', 'Small-Car', 'Dump-Truck', 'Van', 'Excavator', 'Intersection', 'other-vehicle', 'A321', 'Tennis-Court', 'Basketball-Court', \
 'other-airplane', 'Boeing787', 'Warship', 'Tugboat', 'other-ship', 'Tractor', 'Bus', 'Roundabout', 'ARJ21', 'Boeing747', 'Football-Field', \
 'Trailer', 'Truck-Tractor', 'Bridge', 'Baseball-Field', 'A330', 'A350', 'Boeing777', 'C919']
        self.color_pans = [(204,78,210),
                           (0,192,255),
                           (0,131,0),
                           (240,176,0),
                           (254,100,38),
                           (0,0,255),
                           (182,117,46),
                           (185,60,129),
                           (204,153,255),
                           (80,208,146),
                           (0,0,204),
                           (17,90,197),
                           (0,255,255),
                           (102,255,102),
                           (255,255,0)]
        self.num_classes = len(self.category)
        self.cat_ids = {cat:i for i,cat in enumerate(self.category)}
        self.img_ids = self.load_img_ids()
        self.image_path = os.path.join(data_dir, 'images')
        self.label_path = os.path.join(data_dir, 'labelTxt')

    def load_img_ids(self):
        if self.phase == 'train':
            image_set_index_file = os.path.join(self.data_dir, 'trainval.txt')
        else:
            image_set_index_file = os.path.join(self.data_dir, self.phase + '.txt')
        assert os.path.exists(image_set_index_file), 'Path does not exist: {}'.format(image_set_index_file)
        with open(image_set_index_file, 'r') as f:
            lines = f.readlines()
        image_lists = [line.strip() for line in lines]
        return image_lists

    def load_image(self, index):
        img_id = self.img_ids[index]
        imgFile = os.path.join(self.image_path, img_id+'.jpg')
        assert os.path.exists(imgFile), 'image {} not existed'.format(imgFile)
        img = cv2.imread(imgFile)
        return img

    def load_annoFolder(self, img_id):
        return os.path.join(self.label_path, img_id+'.txt')

    def load_annotation(self, index):
        image = self.load_image(index)
        h,w,c = image.shape
        valid_pts = []
        valid_cat = []
        valid_dif = []
        with open(self.load_annoFolder(self.img_ids[index]), 'r') as f:
            for i, line in enumerate(f.readlines()):
                obj = line.split(' ')  # list object
                if len(obj)>8:
                    x1 = min(max(float(obj[0]), 0), w - 1)
                    y1 = min(max(float(obj[1]), 0), h - 1)
                    x2 = min(max(float(obj[2]), 0), w - 1)
                    y2 = min(max(float(obj[3]), 0), h - 1)
                    x3 = min(max(float(obj[4]), 0), w - 1)
                    y3 = min(max(float(obj[5]), 0), h - 1)
                    x4 = min(max(float(obj[6]), 0), w - 1)
                    y4 = min(max(float(obj[7]), 0), h - 1)
                    # TODO: filter small instances
                    xmin = max(min(x1, x2, x3, x4), 0)
                    xmax = max(x1, x2, x3, x4)
                    ymin = max(min(y1, y2, y3, y4), 0)
                    ymax = max(y1, y2, y3, y4)
                    if ((xmax - xmin) > 10) and ((ymax - ymin) > 10):
                        valid_pts.append([[x1,y1], [x2,y2], [x3,y3], [x4,y4]])
                        valid_cat.append(self.cat_ids[obj[8]])
                        valid_dif.append(int(obj[9]))
        f.close()
        annotation = {}
        annotation['pts'] = np.asarray(valid_pts, np.float32)
        annotation['cat'] = np.asarray(valid_cat, np.int32)
        annotation['dif'] = np.asarray(valid_dif, np.int32)
        # pts0 = np.asarray(valid_pts, np.float32)
        # img = self.load_image(index)
        # for i in range(pts0.shape[0]):
        #     pt = pts0[i, :, :]
        #     tl = pt[0, :]
        #     tr = pt[1, :]
        #     br = pt[2, :]
        #     bl = pt[3, :]
        #     cv2.line(img, (int(tl[0]), int(tl[1])), (int(tr[0]), int(tr[1])), (0, 0, 255), 1, 1)
        #     cv2.line(img, (int(tr[0]), int(tr[1])), (int(br[0]), int(br[1])), (255, 0, 255), 1, 1)
        #     cv2.line(img, (int(br[0]), int(br[1])), (int(bl[0]), int(bl[1])), (0, 255, 255), 1, 1)
        #     cv2.line(img, (int(bl[0]), int(bl[1])), (int(tl[0]), int(tl[1])), (255, 0, 0), 1, 1)
        #     cv2.putText(img, '{}:{}'.format(valid_dif[i], self.category[valid_cat[i]]), (int(tl[0]), int(tl[1])), cv2.FONT_HERSHEY_TRIPLEX, 0.6,
        #                 (0, 0, 255), 1, 1)
        # cv2.imshow('img', np.uint8(img))
        # k = cv2.waitKey(0) & 0xFF
        # if k == ord('q'):
        #     cv2.destroyAllWindows()
        #     exit()
        return annotation

    def merge_crop_image_results(self, result_path, merge_path):
        mergebypoly(result_path, merge_path)


##----------------------------------------------------
## Created by Ivan Luiz de Oliveira
## Unesp Sorocaba Dez. 2018
##----------------------------------------------------

##---------------------------------------------------------
##---------------------------------------------------------
## THIS CODE ONLY WORKS ON THE ST-ONE BOARD (RASPBERRY PI)
##---------------------------------------------------------
##---------------------------------------------------------

from keras.models import load_model
import numpy as np
import cv2
import RPi.GPIO as GPIO
from time import sleep

##----------------------------------------------------
## Read config file.
##----------------------------------------------------
config=open('config.cfg','r')
config_data=config.read()
config.close()
config_data=config_data.split('::')
config_data=config_data[1].split('\n')
for i in range(1,len(config_data)):
    config_data[i] = config_data[i].split('==')
cameraPort = int(config_data[1][1])
capture_width = int(config_data[2][1])
capture_height = int(config_data[3][1])
x_partCenter = int(config_data[4][1])
y_partCenter = int(config_data[5][1])
storage_imgSize = int(config_data[6][1])
in_address = config_data[7][1]
out_address = config_data[8][1].split(']')
out_address = out_address[0].split('[')
out_address = out_address[1].split(',')
enable_collor = int(config_data[9][1])
enable_quality = int(config_data[10][1])
training_imgsize = int(config_data[12][1])
enable_preview = int(config_data[13][1])
enable_startup_run = int(config_data[14][1])

partCenter = [x_partCenter,y_partCenter]
storage_imageSize = [storage_imgSize , storage_imgSize]
training_imagesize = [training_imgsize,training_imgsize]
testCamera = [False,]

##---------------------------------------------------------
## Setup ST-One IOs.
##---------------------------------------------------------
IO_mapping = {"OUT1":13,"OUT2":12,"OUT3":25,"OUT4":24,"IN1":19,"IN2":6,"IN3":5,"IN4":22}
in_address=IO_mapping[in_address]
out_address[0]=IO_mapping[out_address[0]]
out_address[1]=IO_mapping[out_address[1]]

GPIO.setwarnings(False)
GPIO.setmode(GPIO.BCM)
GPIO.setup(in_address, GPIO.IN, pull_up_down = GPIO.PUD_DOWN)
GPIO.setup(out_address[0], GPIO.OUT, initial=GPIO.LOW)
GPIO.setup(out_address[1], GPIO.OUT, initial=GPIO.LOW)

##---------------------------------------------------------
## Setup operation mode and test camera
##---------------------------------------------------------

if enable_startup_run==1:
    enable_preview=0

if enable_collor==1:
        model_collor = load_model('models/color/color_ANN_model.h5')
if enable_quality==1:
        model_quality = load_model('models/quality/quality_CNN_model.h5')

while(testCamera[0]==False):
        capture = cv2.VideoCapture(cameraPort)
        testCamera=capture.read()
        if(testCamera[0]==False):
            input('Camera not found!\nPress any key to try again...')
capture.set(3,capture_width)#640
capture.set(4,capture_height)#480

##---------------------------------------------------------
## Operation loop
##---------------------------------------------------------
print("Program Running...")
while True:
    camera_trigger = GPIO.input(in_address)
    while camera_trigger==0:
        sleep(0.1)
        camera_trigger=GPIO.input(in_address)
    print("trigger!")
    ret, frame = capture.read()
    for i in range(10): ret, frame = capture.read()
    current_image=frame[int(partCenter[0]-(storage_imageSize[0]/2)):int(partCenter[0]+(storage_imageSize[0]/2)),int(partCenter[1]-(storage_imageSize[1]/2)):int(partCenter[1]+(storage_imageSize[1]/2))]
    current_image = cv2.resize(current_image, (training_imagesize[0], training_imagesize[1]))
    if enable_preview==1:
        cv2.imshow('Preview Image!',current_image)
        print('Preview the image!')
        waitInput = cv2.waitKey(0)
        cv2.destroyAllWindows()
    current_image = np.array(current_image) / 255.0
    current_image = np.expand_dims(current_image,aaxis=0)

    if enable_collor==1:
        collor_predict = model_collor.predict_classes(current_image)
        print(collor_predict)
        if collor_predict[0]==0:
            print("Part color: BLACK")
            GPIO.output(out_address[0], GPIO.HIGH)
            GPIO.output(out_address[1], GPIO.HIGH)
        elif collor_predict[0]==1:
            print("Part color: RED")
            GPIO.output(out_address[0], GPIO.LOW)
            GPIO.output(out_address[1], GPIO.HIGH)
        elif collor_predict[0]==2:
            print("Part color: SILVER")
            GPIO.output(out_address[0], GPIO.HIGH)
            GPIO.output(out_address[1], GPIO.LOW)

    if enable_quality==1:
        quality_predict = model_quality.predict_classes(current_image)
        print(quality_predict)
        if quality_predict[0]==0:
            GPIO.output(out_address[0], GPIO.LOW)
            GPIO.output(out_address[1], GPIO.HIGH)
        elif quality_predict[0]==1:
            GPIO.output(out_address[0], GPIO.HIGH)
            GPIO.output(out_address[1], GPIO.LOW)

    print("pallet liberado!")
    while camera_trigger==1:
        sleep(0.1)
        camera_trigger=GPIO.input(in_address)
    GPIO.output(out_address[0], GPIO.LOW)
    GPIO.output(out_address[1], GPIO.LOW)
    print('pallet deixou a estacao!')

capture.release()
GPIO.cleanup()
print("-----------------------------\nEnd.")


import sys, os
import numpy as np

"""
Determination of the 4-digit HLA type based on the 2-digit determination by seq2HLA.py
"""

#hoz=1

def determine4digits_main(file,bgVec,runName,HLAclass,hla_classification):
	if hla_classification=="nonclassical":
		result=determine4digitsMinorAlleles(file, bgVec, runName)
	else:
		dbmhc_file=""
		dbmhc=[]
		if HLAclass==1:
			dbmhc_file=os.path.abspath(os.path.dirname(sys.argv[0]))+"/HLAI.dbmhc"
		else:
			dbmhc_file=os.path.abspath(os.path.dirname(sys.argv[0]))+"/HLAII.dbmhc"
		dbmhc=createdbMHCdict(dbmhc,dbmhc_file)
		result=determine4digits(dbmhc, file, bgVec, runName)
	return result

def createdbMHCdict(dbmhc,dbmhc_file):
	population_dict={}
	linenumber=0
	for line in open(dbmhc_file,"r"):
		l=line.split("\t")
		if linenumber==0:
			#header
			for i in range(1,len(l)):
				population_dict[i]=l[i]
			linenumber=1
		else:
			if ":" in l[0]:
				fourdigit=str(l[0].split(":")[0])+":"+str(l[0].split(":")[1])
				dbmhc.append(fourdigit)

	return dbmhc

def meanPanPopFreq(freq_dict):
	pan_list=[]
	for entry in freq_dict:
		pan_list.append(float(freq_dict[entry]))
	return np.mean(pan_list)

def determine4digits(dbmhc, file, bgVec,runName):

	result_dict={}
	result_dict2={}
	readcount=[]
	ambiguity_dict={}
	result=[]
	allele=""
	for line in open(file,"r"):
		l=line.split(":")
		for i in range(0,len(l)-1):
			allele+=l[i]+":"
		result_dict[allele[0:-1]]=int(l[len(l)-1])
		allele=""
		readcount.append(int(l[len(l)-1]))

	writehandle=open(file+".solutions","w")
	writehandle.write("#Full allele\tp-value\n")

	readcount=[x for x in readcount if x != 0]

	top=readcount[0]
	readcount_copy=readcount[1:]
	readcount_copy_str=str(readcount_copy)[1:-1].replace(" ","")

	#call R-script "commmand_fourdigit.R" to calculate the confidence of the top-scoring allele (four-digits)
	routput=os.popen("R --vanilla < "+os.path.abspath(os.path.dirname(sys.argv[0]))+"/command_fourdigit.R --args "+str(top)+" "+str(top)+","+readcount_copy_str).read()
	parseOutput=routput.split("\n")

	entries = []
	repeat=1
	for entry in parseOutput:
		if entry[0:3]=="[1]":
			entries.append(str(entry[4:len(entry)]))
	#print entries[0]
	if not entries[0]=="NA":
		if float(entries[0]) <0.001 and float(entries[0]) > 0.0:
			for item in result_dict:
				if result_dict[item]==top:
					writehandle.write(item+"\t"+str(calculateOutlier(result_dict[item],bgVec))+"\n")
					repeat=0
					mostProbableAllele=str(item.split(":")[0])+":"+str(item.split(":")[1])
					numbersolutions=1

	percentile=95
	while repeat==1:
		if percentile==0:
			return "no,0"
		percentile-=5
		perc=int(np.percentile(readcount,percentile))
		if perc==0:
			perc+=0.1
		count=0
		for item in result_dict:
			if result_dict[item]>=perc:
				count+=1

		#there is only 1 solution:
		if count==1:
			count=0
			for item in result_dict:
				if result_dict[item]>=perc:
					writehandle.write(item+"\t"+str(calculateOutlier(result_dict[item],bgVec))+"\n")
					repeat=0
					mostProbableAllele=str(item.split(":")[0])+":"+str(item.split(":")[1])
					numbersolutions=1
		#there are more possible solutions
		else:
			count=0
			#check if allele (4-digits) occurs in dbMHC table. If not, is it a very rare allele in all populations
			for item in result_dict:
				item4digit=str(item.split(":")[0])+":"+str(item.split(":")[1])
				if result_dict[item]>=perc and item4digit in dbmhc:
					result.append(item4digit)
					result_dict2[item]=result_dict[item]
			#only 1 solution (at 4-digit level) left
			if len(set(result))==1:
				#for item in result_dict:
				for item in result_dict2:
					writehandle.write(item+"\t"+str(calculateOutlier(result_dict2[item],bgVec))+"\n")
					repeat=0
					mostProbableAllele=str(item.split(":")[0])+":"+str(item.split(":")[1])
					numbersolutions=1
			#more than one 2 solution possible (at 4-digit level).
			elif len(set(result))>1:
				numbersolutions=len(set(result))
				max=0.0

				for item in result_dict2:
					item4digit=str(item.split(":")[0])+":"+str(item.split(":")[1])
					if not item4digit in ambiguity_dict:
						ambiguity_dict[item4digit]=result_dict2[item]
					elif ambiguity_dict[item4digit]<result_dict2[item]:
						ambiguity_dict[item4digit]=result_dict2[item]
					if result_dict2[item] > max:
						max=result_dict2[item]
						mostProbableAllele=item4digit
				ambiguity_handle=open(runName+".ambiguity", "a")
				ambiguity_handle.write("#################################################################################################\n")
				ambiguity_handle.write("#Ambiguity:\n")
				ambiguity_handle.write("#Based on the RNA-Seq reads and the dbMHC table, the following 4-digits alleles are possible:\n")
				for ambi4DigitAllele in ambiguity_dict:
					ambiguity_handle.write(ambi4DigitAllele+"\t"+str(calculateOutlier(ambiguity_dict[ambi4DigitAllele],bgVec))+"\n")

				ambiguity_handle.write("#However, by taking into account the read data, the most probable 4-digit allele is:\n")
				ambiguity_handle.write(mostProbableAllele+"\n\n")
				ambiguity_handle.close

				for item in result_dict2:
					item4digit=str(item.split(":")[0])+":"+str(item.split(":")[1])
					if result_dict2[item]==max:
						writehandle.write(item+"\t"+str(calculateOutlier(result_dict2[item],bgVec))+"\n")
						repeat=0
	writehandle.close()
	return mostProbableAllele+","+str(numbersolutions)

def calculateOutlier(top,bgVec):
	#call R-script "commmand_fourdigit.R" to calculate the confidence of the top-scoring allele (four-digits)
	routput=os.popen("R --vanilla < "+os.path.abspath(os.path.dirname(sys.argv[0]))+"/command_fourdigit.R --args "+str(top)+" "+bgVec).read()
	parseOutput=routput.split("\n")

	entries = []
	repeat=1
	for entry in parseOutput:
		if entry[0:3]=="[1]":
			entries.append(str(entry[4:len(entry)]))
	return entries[0]

#method for determining the 4 digit HLA type for minor HLA alleles, for which no entry in dbMHC is available
def determine4digitsMinorAlleles(file, bgVec,runName):

	result_dict={}
	result_dict2={}
	readcount=[]
	ambiguity_dict={}
	result=[]
	allele=""
	for line in open(file,"r"):
		l=line.split(":")
		for i in range(0,len(l)-1):
			allele+=l[i]+":"
		result_dict[allele[0:-1]]=int(l[len(l)-1])
		allele=""
		readcount.append(int(l[len(l)-1]))

	writehandle=open(file+".solutions","w")
	writehandle.write("#Full allele\tp-value\n")

	readcount=[x for x in readcount if x != 0]

	top=readcount[0]
	readcount_copy=readcount[1:]
	readcount_copy_str=str(readcount_copy)[1:-1].replace(" ","")

	#call R-script "commmand_fourdigit.R" to calculate the confidence of the top-scoring allele (four-digits)
	routput=os.popen("R --vanilla < "+os.path.abspath(os.path.dirname(sys.argv[0]))+"/command_fourdigit.R --args "+str(top)+" "+str(top)+","+readcount_copy_str).read()
	parseOutput=routput.split("\n")

	entries = []
	repeat=1
	for entry in parseOutput:
		if entry[0:3]=="[1]":
			entries.append(str(entry[4:len(entry)]))
	#print entries[0]
	if not entries[0]=="NA":
		if float(entries[0]) <0.001 and float(entries[0]) > 0.0:
			for item in result_dict:
				if result_dict[item]==top:
					writehandle.write(item+"\t"+str(calculateOutlier(result_dict[item],bgVec))+"\n")
					repeat=0
					mostProbableAllele=str(item.split(":")[0])+":"+str(item.split(":")[1])
					numbersolutions=1

	percentile=95
	while repeat==1:
		if percentile==0:
			return "no,0"
		percentile-=5
		perc=int(np.percentile(readcount,percentile))
		if perc==0:
			perc+=0.1
		count=0
		for item in result_dict:
			if result_dict[item]>=perc:
				count+=1

		#there is only 1 solution:
		if count==1:
			count=0
			for item in result_dict:
				if result_dict[item]>=perc:
					writehandle.write(item+"\t"+str(calculateOutlier(result_dict[item],bgVec))+"\n")
					repeat=0
					mostProbableAllele=str(item.split(":")[0])+":"+str(item.split(":")[1])
					numbersolutions=1
		#there are more possible solutions
		else:
			count=0
			#check if allele (4-digits) is in the defined percentile range
			for item in result_dict:
				item4digit=str(item.split(":")[0])+":"+str(item.split(":")[1])
				if result_dict[item]>=perc:
					result.append(item4digit)
					result_dict2[item]=result_dict[item]
			#only 1 solution (at 4-digit level) left
			if len(set(result))==1:
				#for item in result_dict:
				for item in result_dict2:
					writehandle.write(item+"\t"+str(calculateOutlier(result_dict2[item],bgVec))+"\n")
					repeat=0
					mostProbableAllele=str(item.split(":")[0])+":"+str(item.split(":")[1])
					numbersolutions=1
			#more than one 2 solution possible (at 4-digit level).
			elif len(set(result))>1:
				numbersolutions=len(set(result))
				max=0.0

				for item in result_dict2:
					item4digit=str(item.split(":")[0])+":"+str(item.split(":")[1])
					if not item4digit in ambiguity_dict:
						ambiguity_dict[item4digit]=result_dict2[item]
					elif ambiguity_dict[item4digit]<result_dict2[item]:
						ambiguity_dict[item4digit]=result_dict2[item]
					if result_dict2[item] > max:
						max=result_dict2[item]
					#if meanPanPopFreq(dbmhc_prob[item4digit]) > max:
					#	max = meanPanPopFreq(dbmhc_prob[item4digit])
						mostProbableAllele=item4digit
					#if np.mean(dbmhc_prob[item4digit]) > max:
					#	max = np.mean(dbmhc_prob[item4digit])
					#	mostProbableAllele=item4digit
				ambiguity_handle=open(runName+".ambiguity", "a")
				ambiguity_handle.write("#################################################################################################\n")
				ambiguity_handle.write("#Ambiguity:\n")
				ambiguity_handle.write("#Based on the RNA-Seq reads and the dbMHC table, the following 4-digits alleles are possible:\n")
				for ambi4DigitAllele in ambiguity_dict:
					ambiguity_handle.write(ambi4DigitAllele+"\t"+str(calculateOutlier(ambiguity_dict[ambi4DigitAllele],bgVec))+"\n")

				ambiguity_handle.write("#However, by taking into account the read data, the most probable 4-digit allele is:\n")
				ambiguity_handle.write(mostProbableAllele+"\n\n")
				ambiguity_handle.close

				for item in result_dict2:
					item4digit=str(item.split(":")[0])+":"+str(item.split(":")[1])
					#if np.mean(dbmhc_prob[item4digit]) == max:
					#if meanPanPopFreq(dbmhc_prob[item4digit]) == max:
					if result_dict2[item]==max:
						writehandle.write(item+"\t"+str(calculateOutlier(result_dict2[item],bgVec))+"\n")
						repeat=0
	writehandle.close()
	return mostProbableAllele+","+str(numbersolutions)


# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""export"""
import os
import numpy as np

from mindspore import Tensor, export, load_checkpoint
from mindspore import context

from src.dien import DIEN
from src.dataset import DataIterator
from src.config import parse_args

args_opt = parse_args()

EMBEDDING_DIM = 18
HIDDEN_SIZE = 18 * 2
ATTENTION_SIZE = 18 * 2
best_auc = 0.0

def export_DIEN(train_file, uid_voc, mid_voc, cat_voc, meta_path, review_path, batch_size, max_len, ckpt_path):
    """export"""

    train_data = DataIterator(train_file, uid_voc, mid_voc, cat_voc, meta_path, review_path, batch_size, maxlen,
                              shuffle_each_epoch=False)

    n_uid, n_mid, n_cat = train_data.get_n()

    model = DIEN(n_uid, n_mid, n_cat, embedding_size=EMBEDDING_DIM)

    # load the parameter into net
    load_checkpoint(ckpt_path, net=model)
    uids = Tensor(np.random.rand(batch_size).astype(np.int32))
    mids = Tensor(np.random.rand(batch_size).astype(np.int32))
    cats = Tensor(np.random.rand(batch_size).astype(np.int32))
    mid_his = Tensor(np.random.rand(batch_size, max_len).astype(np.int32))
    cat_his = Tensor(np.random.rand(batch_size, max_len).astype(np.int32))
    mid_mask = Tensor(np.random.rand(batch_size, max_len).astype(np.int32))
    noclk_mids = Tensor(np.random.rand(batch_size, max_len, 5).astype(np.int32))
    noclk_cats = Tensor(np.random.rand(batch_size, max_len, 5).astype(np.int32))
    export(model, mid_mask, uids, mids, cats, mid_his, cat_his,
           noclk_mids, noclk_cats, file_name=args_opt.MINDIR_file_name, file_format='MINDIR')

def main():
    context.set_context(mode=context.GRAPH_MODE, device_target='Ascend')

    device_id = args_opt.device_id

    context.set_context(device_id=device_id)

    dataset_file_path = args_opt.dataset_file_path
    train_file = args_opt.train_mindrecord_path
    uid_voc = os.path.join(dataset_file_path, "uid_voc.pkl")
    mid_voc = os.path.join(dataset_file_path, "mid_voc.pkl")
    cat_voc = os.path.join(dataset_file_path, "cat_voc.pkl")
    meta_path = os.path.join(dataset_file_path, "item-info")
    review_path = os.path.join(dataset_file_path, "reviews-info")
    batch_size = args_opt.batch_size
    max_len = args_opt.max_len
    ckpt_path = args_opt.save_checkpoint_path
    export_DIEN(train_file, uid_voc, mid_voc, cat_voc, meta_path, review_path, batch_size, max_len, ckpt_path)
    print('export success!')

if __name__ == "__main__":
    main()


import copy
import time
from collections import Counter
import cv2
import pydicom
import torch
import matplotlib.pyplot as plt
from scipy import signal

import ltr.data.processing_utils as prutils
from ltr.data.processing_utils import compute_distance_of_two_point
from pydoctor.doctors.base import BaseDoctor
import torch.nn.functional as F
import numpy as np
import SimpleITK as sitk
import torch.nn.functional as F

def stack_tensors(x):
    if isinstance(x, (list, tuple)) and isinstance(x[0], torch.Tensor):
        return torch.stack(x)
    return x

def compute_distance_list(keypoint_list):
    keypoint_list.sort(key=lambda x:x[1], reverse=True)
    distance_list = []
    for i, label in enumerate(keypoint_list):
        if i ==0:
            continue
        else:
            distance = compute_distance_of_two_point(keypoint_list[i],keypoint_list[i-1])
            distance_list.append(int(distance))
    return distance_list

class DoctorWang(BaseDoctor):

    def initialize_feature(self):
        if not getattr(self, 'features_initialized', False):
            self.params.body_loc_net.initialize()
            self.params.body_cls_net.initialize()
            self.params.disc_loc_net.initialize()
            self.params.disc_cls_net.initialize()
        self.features_initialized = True

    def initialize(self):
        if not self.params.has('device'):
            self.params.device = 'cuda' if self.params.use_gpu else 'cpu'
        self.initialize_feature()
        self.body_loc_net = self.params.body_loc_net
        self.body_cls_net = self.params.body_cls_net
        self.disc_loc_net = self.params.disc_loc_net
        self.disc_cls_net = self.params.disc_cls_net
        return True

    def diagnose(self, frame, index, name):
        self.image_sz = torch.tensor([frame.shape[1], frame.shape[2]])
        self.support_sz = self.params.image_resize_size
        self.resized_image, self.resize_factor = self.resize_image(frame, self.support_sz)

        # compute the location stage.
        location_dict = self.get_body_disc_keypoint_dict(frame,name)
        scale_factor = {'body':[1.0,1.1],'disc':[0.9,1.4]}
        # crop the patches from frame with first stage result and put it in dict
        frame_tensor = torch.from_numpy(frame).float()
        patches = self.crop_from_frame(location_dict,frame_tensor,scale_factor=scale_factor)

        with torch.no_grad():
            # extract_backbone features
            body_cls_init_backbone_feat = self.body_cls_net.extract_backbone(patches['body_patches'])
            disc_cls_init_backbone_feat = self.disc_cls_net.extract_backbone(patches['disc_patches'])
            # get the cls feat (FC layer)
            body_cls_backbone_feat = self.body_cls_net.get_backbone_reg_feat(body_cls_init_backbone_feat)
            disc_cls_backbone_feat = self.disc_cls_net.get_backbone_reg_feat(disc_cls_init_backbone_feat)
            # get the class output
            body_cls_pred = self.body_cls_net.cls_regressor(body_cls_backbone_feat[0])
            disc_cls_pred = self.disc_cls_net.cls_regressor(disc_cls_backbone_feat[0])
            # return the predict label.
            body_index_list = torch.argmax(body_cls_pred, dim=-1).view(5,-1).tolist()
            disc_index_list = torch.argmax(disc_cls_pred, dim=-1).view(6,-1).tolist()

        self.body_id = {'4':'L1','3':'L2','2':'L3','1':'L4','0':'L5'}
        self.body_class = {'0':'v1','1':'v2'}
        self.disc_id = {'5':'T12-L1','4':'L1-L2','3':'L2-L3','2':'L3-L4','1':'L4-L5','0':'L5-S1'}
        self.disc_class = {'0':'v1','1':'v2','2':'v3','3':'v4','4':'v5'}
        point = []
        for idx,body_point in enumerate(body_index_list):
            body_cls = self.body_class[str(Counter(body_point).most_common(1)[0][0])]
            coord = location_dict['body']['coord'][idx]
            tag = {'identification':self.body_id[str(idx)],
                   'vertebra':str(body_cls)}
            point_dict = {'tag':tag,'coord':coord,'zIndex':index}
            point.append(point_dict)
        for idx2 ,disc_point in enumerate(disc_index_list):
            disc_cls = self.disc_class[str(Counter(disc_point).most_common(1)[0][0])]
            coord = location_dict['disc']['coord'][idx2]
            tag = {'identification':self.disc_id[str(idx2)],
                   'disc':str(disc_cls)}
            point_dict = {'tag':tag,'coord':coord,'zIndex':index}
            point.append(point_dict)
        annotation = [{'annotator':0,'data':{'point':point}}]
        return annotation

    def resize_image(self, image, resize_sz):
        """
            Resize a pair data ( image and key_point mask_label) into a resize_sz.
            :param image: Src data image
            :param resize_sz: A unified size for a batch.
            :return: return a pair re_sized data.
            """
        image = torch.from_numpy(image).float().unsqueeze(0)
        _, channel, width, high = image.shape
        image_output = (F.interpolate(image, (resize_sz, resize_sz), mode='bilinear',
                                      align_corners=False))
        resize_factor = [width / resize_sz, high / resize_sz]
        return image_output, resize_factor

    def locate_keypoint_from_mask(self, pred_mask, point_num, threshold,name):
        """
        Locate the key point from the predicted mask.
        :param threshold: the threshold to compute
        :param pred_mask: mask for analyse
        :param max_order: threshold for adjust.
        :param resize_factor: resize factor to restore the coordinate.
        :param total key point:the number key point to located.
        :return:
        """

        keypoint_result, Flag = self.result_analyse_mask(predict=pred_mask, threshold=threshold, keypoint_num=point_num)
        if Flag is not True or name == 'study252':
            if name =='study252':
                distance_list = compute_distance_list(keypoint_result)
                if max(distance_list) > 1.5 * min(distance_list):
                    index_max = distance_list.index(max(distance_list))
                    x = (keypoint_result[index_max][0] - keypoint_result[index_max+1][0])//2
                    y = (keypoint_result[index_max][1] - keypoint_result[index_max+1][1])//2
                    new_point = [keypoint_result[index_max][0] - x, keypoint_result[index_max][1] - y]
                    keypoint_result.insert((index_max+1),new_point)
                    # print('Lost at inner.')

            if len(keypoint_result) < point_num:
                # print('Lost {} point'.format(point_num-len(keypoint_result)))
                #  compute the distance_between_linear_point_list
                distance_list = compute_distance_list(keypoint_result)
                if max(distance_list) > 1.5 * min(distance_list):
                    pass
                    # print('Lost at inner.')

                else:
                    # print('Lost at out')
                    lost_number = point_num - len(keypoint_result)
                    for i in range(0,lost_number):
                        x = keypoint_result[-2][0]-keypoint_result[-1][0]
                        y = keypoint_result[-2][1]-keypoint_result[-1][1]
                        new_point=[keypoint_result[-1][0]-x,keypoint_result[-1][1]-y]
                        keypoint_result.append(new_point)

            if len(keypoint_result)> point_num:
                keypoint_result = keypoint_result[:point_num]

        test_mask = torch.zeros_like(pred_mask)
        for point in keypoint_result:
            y,x = point
            test_mask[x][y] = torch.tensor(1)
        kernel = np.ones((5,5))
        score_mask = torch.tensor(signal.convolve2d(test_mask.numpy(), kernel, boundary='symm', mode='same'))

        return keypoint_result, int(np.mean(compute_distance_list(keypoint_result))),score_mask

    def result_analyse_mask(self, predict, threshold,  keypoint_num):
        '''
        Analize the mask of network output
        :param predict:  The predict of network output
        :param threshold: the threshold
        :param resize_factor: the list of resize_factor to restore the location
        :return:
        '''
        # get mask
        mask = (predict > predict.max() * threshold).float().cpu().numpy().astype(np.uint8)

        # get contours from mask
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

        # compute the area for the contours
        cnt_area = [cv2.contourArea(cnt) for cnt in contours]
        if len(contours) > keypoint_num:
            # remove the little noise with area < 20.
            for idx, area  in enumerate(cnt_area):
                if area <=10:
                    contours.pop(idx)
            # other method need to do .
        location_output = []
        for cont in contours:
            polygon = cont.reshape(-1, 2)
            # with a most suitable circle to fit the polygon of contour
            # return the center and radius
            (loc_x, loc_y), radius = cv2.minEnclosingCircle(polygon)

            location_output.append([int(loc_x), int(loc_y)])
        flag = len(location_output) == keypoint_num
        location_output.sort(key=lambda x: x[1], reverse=True)
        return location_output,flag

    def get_body_disc_keypoint_dict(self,image,name):
        with torch.no_grad():
            # extract_backbone features
            body_loc_init_backbone_feat = self.body_loc_net.extract_backbone(self.resized_image)
            disc_loc_init_backbone_feat = self.disc_loc_net.extract_backbone(self.resized_image)
            # get the reg feat
            body_keypoint_backbone_feat = self.body_loc_net.get_backbone_kpreg_feat(body_loc_init_backbone_feat)
            disc_keypoint_backbone_feat = self.disc_loc_net.get_backbone_kpreg_feat(disc_loc_init_backbone_feat)
            # compute the mask for location
            body_pred = self.body_loc_net.kp_regressor.forward(body_keypoint_backbone_feat)
            disc_pred = self.disc_loc_net.kp_regressor.forward(disc_keypoint_backbone_feat)
        # reshape the pred for processing.
        body_pred = F.interpolate(body_pred, (self.image_sz[0], self.image_sz[1]), mode='bilinear', align_corners=True)
        disc_pred = F.interpolate(disc_pred, (self.image_sz[0], self.image_sz[1]), mode='bilinear', align_corners=True)
        body_pred = body_pred.view(*body_pred.shape[-2:])
        disc_pred = disc_pred.view(*disc_pred.shape[-2:])
        body_loc = body_pred.cpu().detach().numpy().astype(np.uint8)
        disc_loc = disc_pred.cpu().detach().numpy().astype(np.uint8)
        # self.visdom.register(torch.tensor(body_loc + image[0]), 'heatmap', 2, 'body_pred_heatmap')
        # self.visdom.register(torch.tensor(disc_loc + image[0]), 'heatmap', 2, 'disc_pred_heatmap')
        # for visualize
        # plt.figure(1)
        # plt.subplot(1, 1, 1)
        # plt.imshow(body_loc+image[0])
        # plt.savefig('/home/adminer/SPARK/PyDoctor/pydoctor/result_plots/body_loc/' + name + '.png')
        # plt.show()
        # #
        # plt.figure(1)
        # plt.subplot(1, 1, 1)
        # plt.imshow(disc_loc+image[0])
        # plt.savefig('/home/adminer/SPARK/PyDoctor/pydoctor/result_plots/disc_loc/' + name + '.png')
        # plt.show()

        body_keypoint_list,body_average,result_mask_body = self.locate_keypoint_from_mask(body_pred, self.params.body_keypoint,
                                                              self.params.body_threshold,name='None')
        disc_keypoint_list,disc_average,result_mask_disc = self.locate_keypoint_from_mask(disc_pred, self.params.disc_keypoint,
                                                              self.params.disc_threshold,name)
        # self.visdom.register(torch.tensor(image[0])+result_mask_body, 'heatmap', 2, 'body_pred_heatmap_process')
        # self.visdom.register(torch.tensor(image[0])+result_mask_disc, 'heatmap', 2, 'disc_pred_heatmap_process')
        # for visualize
        # plt.figure(1)
        # plt.subplot(1, 1, 1)
        # plt.imshow(result_mask_body.numpy()+image[0])
        # plt.savefig('/home/adminer/SPARK/PyDoctor/pydoctor/result_plots/body_loc_B/' + name + '.png')
        # plt.show()
        # #
        # plt.figure(1)
        # plt.subplot(1, 1, 1)
        # plt.imshow(result_mask_disc.numpy()+image[0])
        # plt.savefig('/home/adminer/SPARK/PyDoctor/pydoctor/result_plots/disc_loc_B/' + name + '.png')
        # plt.show()

        return {'body':{'coord':body_keypoint_list,'size':body_average},
                'disc':{'coord':disc_keypoint_list,'size':disc_average}}

    def crop_from_frame(self, location_dict,frame,scale_factor):
        body_coord,body_size = location_dict['body']['coord'],location_dict['body']['size']
        disc_coord,disc_size = location_dict['disc']['coord'],location_dict['disc']['size']
        body_patches = self.generate_bbox_label(body_coord,scale_factor['body'],body_size,11,frame,resize_sz=64)
        disc_patches = self.generate_bbox_label(disc_coord,scale_factor['disc'],disc_size,11,frame,resize_sz=64)
        # self.visdom.register(body_patches,'images',2,'body_patches')
        # self.visdom.register(disc_patches,'images',2,'disc_patches')
        return {'body_patches':body_patches,'disc_patches':disc_patches}

    def generate_bbox_label(self, location_list, scale_factor,size,argument_number,frame,resize_sz=64):
        bbox_size = [scale_factor[0]*size, scale_factor[1]*size]
        bbox_list = []
        for point in location_list:

            x_coord,y_coord = point[1],point[0]
            bbox = torch.tensor([x_coord - bbox_size[0]//2,y_coord - bbox_size[1]//2,bbox_size[0],bbox_size[1]]).round()
            bbox_argue = [bbox]
            patches = []
            for i in range(argument_number-1):
                if True:
                    argued_bbox, iou = prutils.perturb_box(bbox)

                bbox_argue.append(argued_bbox)
            for crop_bbox in bbox_argue:
                x,y,w,h = crop_bbox.tolist()
                x,y,w,h = int(x),int(y),int(w),int(h),
                patch = (F.interpolate((frame[:,x:x+w,y:y+h]).unsqueeze(0),(resize_sz,resize_sz),mode='bilinear',align_corners=False)).squeeze(0)
                patches.append(patch)
            # for idx, pt in enumerate(patches):
            #     self.visdom.register(pt[0],'heatmap',2,str(idx)+'_image')
            bbox_argue = torch.stack(patches,0)
            bbox_list.append(bbox_argue)
        bbox_list = torch.stack(bbox_list,0)
        patches_return = bbox_list.view(-1,3,64,64)
        return  patches_return

    def read_dicom_image(self, study):
        image = np.squeeze(sitk.GetArrayFromImage(sitk.ReadImage(study.frame_path)))
        key_image = np.expand_dims(np.uint8((image - image.min()) / (image.max() - image.min()) * 255.0),0).repeat(3, axis=0)
        pydicom_file = pydicom.read_file(study.frame_path)
        studyUid = pydicom_file.get(0x0020000D)._value
        seriesUid = pydicom_file.get(0x0020000E)._value
        instanceUid = pydicom_file.get(0x00080018)._value
        uid_info = {'studyUid':str(studyUid),
                    'seriesUid':str(seriesUid),
                    'instanceUid':str(instanceUid)}
        return key_image/255.0, uid_info


#!/usr/bin/env python
import os, sys, argparse, errno, yaml, time, datetime
import rospy, rospkg
import torch, torchvision, cv2, PIL.Image
import numpy as np
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import CompressedImage, Image
from road_following.msg import Inference
from jetcam_ros.utils import bgr8_to_jpeg

class Inference_Model_Node(object):
    def __init__(self):
        self.package = "road_following"
        self.node_name = rospy.get_name()
        self.veh_name = self.node_name.split("/")[1]
        self.start = rospy.wait_for_message("/" + self.veh_name +"/jetson_camera/image/raw", Image)
        rospy.loginfo("[{}]  Initializing road_model_inference.py......".format(self.node_name))

        ## set/get ros param
        self.model_name  = self.setup_parameter("~model_pth", "best_steering_model_xy.pth")
        self.use_cuda  = self.setup_parameter("~use_cuda", True)

        ## read information
        self.recording = self.read_param_from_file(package=self.package, folder="model", file_name="recording.yaml")
        self.inference_parameter = self.read_param_from_file(package=self.package, folder="param", file_name="inference_model.yaml")

        ## check model
        if self.check_model_exist(self.model_name):
            ## configure model
            self.model_struct = self.recording[self.model_name]["train"]["model"]
            self.model = self.load_model(model=self.model_struct, param_pretrained=False) # configure: self.model
            self.cuda(use=self.use_cuda) # configure: self.device, self.model

            ## setup parameter with processing image
            self.process_img_mean  = torch.Tensor(self.inference_parameter["preprocess_mean"]).cuda().half()
            self.process_img_stdev = torch.Tensor(self.inference_parameter["preprocess_stdev"]).cuda().half()

            ## setup parameter with inference
            self.angle = 0.0
            self.angle_last = 0.0

            ## CV_bridge
            self.bridge = CvBridge()

            ## configure subscriber
            self.first_sub = True
            self.sub_msg = rospy.Subscriber("~image/raw", Image, self.convert_image_to_cv2,queue_size=1)

            ## configure Publisher
            self.pub_msg = rospy.Publisher("~inference", Inference, queue_size=1)
        else:
            self.on_shutdown()

    def getFilePath(self, package, folder, file_name):
        rospack = rospkg.RosPack()
        return rospack.get_path(package) + "/" + folder + "/" + file_name

    def read_param_from_file(self, package, folder, file_name):
        fname = self.getFilePath(package, folder, file_name)
        with open(fname, 'r') as in_file:
            try:
                yaml_dict = yaml.load(in_file)
            except yaml.YAMLError as exc:
                rospy.loginfo("[{}] YAML syntax  error. File: {}".format(self.node_name, fname))
        return yaml_dict

    def check_model_exist(self, name):
        if not name in self.recording.keys():
            rospy.logwarn("[{}] {} does not exist in folder [model]. Please check your model_pth in [{}].".format(self.node_name, name, self.getFilePath(package=self.package, folder="model", file_name="inference_model.yaml") ))
            rospy.logwarn("[{}] There are model_pth you can use below:".format(self.node_name))
            rospy.logwarn("[{}] {}".format(self.node_name, list(self.recording.keys())))
            return False
        return True

    def load_model(self, model="resnet18", param_pretrained=False, kind_of_classifier=2):
        # reference : https://pytorch.org/docs/stable/torchvision/models.html
        model_list = [
                      "resnet18", "alexnet", "squeezenet", "vgg16",
                      "densenet", "inception", "googlenet", "shufflenet",
                      "mobilenet", "resnet34", "wide_resnet50_2", "mnasnet"
                     ]

        if model in model_list:
            rospy.logwarn("[{0}] [{1}]] use [{2}]. Need some time to load model [{1}]...".format(self.node_name, self.model_name, model))
            start_time = rospy.get_time()
            if model == "resnet18":
                self.model = torchvision.models.resnet18(pretrained=param_pretrained)
                self.model.fc = torch.nn.Linear(512, 2)
            elif model == "alexnet":
                self.model = torchvision.models.alexnet(pretrained=param_pretrained)
                self.model.classifier[-1] = torch.nn.Linear(self.model.classifier[-1].in_features, int(kind_of_classifier))
            elif model == "squeezenet":
                self.model = torchvision.models.squeezenet1_1(pretrained=param_pretrained)
                self.model.classifier[1] = torch.nn.Conv2d(self.model.classifier[1].in_features, kind_of_classifier, kernel_size=1)
                self.num_classes = kind_of_classifier
            elif model == "vgg16":
                self.model = torchvision.models.vgg16(pretrained=param_pretrained)
            elif model == "densenet":
                self.model = torchvision.models.densenet161(pretrained=param_pretrained)
            elif model == "inception":
                self.model = torchvision.models.inception_v3(pretrained=param_pretrained)
            elif model == "googlenet":
                self.model = torchvision.models.googlenet(pretrained=param_pretrained)
            elif model == "shufflenet":
                self.model = torchvision.models.shufflenet_v2_x1_0(pretrained=param_pretrained)
            elif model == "mobilenet":
                self.model = torchvision.models.mobilenet_v2(pretrained=param_pretrained)
            elif model == "resnext50_32x4d":
                self.model = torchvision.models.resnext50_32x4d(pretrained=param_pretrained)
            elif model == "resnet34":
                self.model = torchvision.models.resnet34(pretrained=param_pretrained)
                self.model.fc = torch.nn.Linear(512, kind_of_classifier)
            elif model == "wide_resnet50_2":
                self.model = torchvision.models.wide_resnet50_2(pretrained=param_pretrained)
            elif model == "mnasnet":
                self.model = torchvision.models.mnasnet1_0(pretrained=param_pretrained)
            model_pth = self.getFilePath(package=self.package, folder="model", file_name=self.model_name)
            self.model.load_state_dict(torch.load(model_pth))
            interval = rospy.get_time() - start_time
            rospy.loginfo("[{}] Loading model done! Use {:.2f} seconds.".format(self.node_name, interval))
            return self.model
        else:
            rospy.loginfo("[{}] Your classifier is wrong. Please check out model struct!".format(self.node_name))
            self.on_shutdown()

    def cuda(self,use=False):
        if use == True:
            rospy.loginfo("[{}] Using cuda! Need some time to start...".format(self.node_name))
            self.device = torch.device('cuda')
            start_time = rospy.get_time()
            self.model = self.model.to(self.device)
            self.model = self.model.eval().half()
            interval = rospy.get_time() - start_time
            rospy.loginfo("[{}] Done with starting! Can use cuda now! Use {:.2f} seconds to start.".format(self.node_name, interval))
        else:
            rospy.loginfo("[{}] Do not use cuda!".format(self.node_name))

    def convert_image_to_cv2(self, img_msg):
        try:
            # Convert your ROS Image ssage to opencv2
            cv2_img = self.bridge.imgmsg_to_cv2(img_msg, desired_encoding="bgr8")
            #jpeg_img = cv2.resize(bgr8_to_jpeg(cv2_img), (224, 224))
            self.inference(img=cv2_img)
        except CvBridgeError as e:
            print(e)

    def preprocess(self, camera_value):
        img = PIL.Image.fromarray(cv2.cvtColor(camera_value, cv2.COLOR_BGR2RGB))
        img = torchvision.transforms.functional.to_tensor(img).to(self.device).half()
        img.sub_(self.process_img_mean[:, None, None]).div_(self.process_img_stdev[:, None, None])
        return img[None, ...]

    def inference(self, img):
        start_time = rospy.get_time()
        inference_msg = Inference()
        if self.first_sub == True:
            rospy.loginfo("[{}] Deploy model to gpu. Please wait...".format(self.node_name))
        img = self.preprocess(img)
        xy = self.model(img).detach().float().cpu().numpy().flatten()
        if self.first_sub == True:
            interval = rospy.get_time() - start_time
            self.inference_information(interval)
            self.first_sub = False
        x = xy[0]
        y = (0.5 - xy[1]) / 2.0

        self.angle = np.arctan2(x, y)
        self.angle_last = self.angle
        #print("angle: {}, last_angle: {}".format(self.angle, self.angle_last))
        inference_msg.angle = self.angle
        inference_msg.angle_last = self.angle_last
        self.pub_msg.publish(inference_msg)
        time.sleep(0.001)

    def inference_information(self, interval):
        rospy.loginfo("[{}] Deployment complete! Use {:.2f} seconds.".format(self.node_name, interval))
        rospy.loginfo("[{}] You can listen the topic to see how much the angle of head between the lane: {}".format(self.node_name, self.node_name + "/inference"))
        rospy.loginfo("[{}] More information about {} :\n{}".format(self.node_name, self.model_name, self.recording[self.model_name]))

    def on_shutdown(self):
        rospy.loginfo("[{}] Close.".format(self.node_name))
        rospy.loginfo("[{}] shutdown.".format(self.node_name))
        rospy.loginfo("[{}] Now you can press [ctrl] + [c] to close this launch file.".format(self.node_name))
        rospy.sleep(1)
        rospy.is_shutdown = True
        #sys.exit()

    def setup_parameter(self, param_name, default_value):
        value = rospy.get_param(param_name, default_value)
        # Write to parameter server for transparency
        rospy.set_param(param_name, value)
        rospy.loginfo("[%s] %s = %s " % (self.node_name, param_name, value))
        return value

if __name__ == "__main__" :
    rospy.init_node("road_model_inference", anonymous=False)
    inference_model_node = Inference_Model_Node()
    rospy.on_shutdown(inference_model_node.on_shutdown)
    rospy.spin()


import multiprocessing
import threading
import tensorflow as tf
import numpy as np
import gym
import os
import shutil
import matplotlib.pyplot as plt

from playground.utilities.utils import get_latest_run_count
from playground.a3c.worker import Worker
from playground.a3c.actorcritic import ActorCriticNet

N_WORKERS = multiprocessing.cpu_count()

GLOBAL_EP = 0
DATA_LENGTH = 600
INIT_CASH = 100

GLOBAL_NET_SCOPE = 'Global_Net'

if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument('-d', '--device', required=True, help='either "CPU" or "GPU" for training')
    args = vars(ap.parse_args())
    SESS = tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True))

    with tf.device("/{}:0".format(args['device'])):

        GLOBAL_ACTORCRITIC = ActorCriticNet(GLOBAL_NET_SCOPE, SESS)  # we only need its params

        COORD = tf.train.Coordinator()

        # Setup Tensorboard
        summary_writer = tf.summary.FileWriter('logs/' + str(get_latest_run_count()))
        summary_writer.add_graph(SESS.graph)

        workers = []
        # Create worker
        for i in range(N_WORKERS):
            i_name = 'W_%i' % i   # worker name
            workers.append(Worker(i_name, GLOBAL_ACTORCRITIC, SESS, summary_writer))

        SESS.run(tf.global_variables_initializer())
        # loading networks
        saver = tf.train.Saver()

        worker_threads = []
        for worker in workers:
            # checkpoint = tf.train.get_checkpoint_state('logs/' + str(get_latest_run_count()-2) + '/saved_networks_' + worker.name)
            # if checkpoint and checkpoint.model_checkpoint_path:
            #     saver.restore(SESS, checkpoint.model_checkpoint_path)
            #     print("Successfully loaded:", checkpoint.model_checkpoint_path)
            # else:
            #     print("Could not find old network weights")

            job = lambda: worker.work(COORD, saver)
            t = threading.Thread(target=job)
            t.start()
            worker_threads.append(t)
        COORD.join(worker_threads)


import pdb

import numpy as np
import torch
import os

from torch import optim

import network.DuRN_Pure_Conv_3_Refactor as pure
import model_zoo
# from network.pyramid_ppp import Pyramid_Net
from network.Ms_Discriminator import MsImageDis
from network.base_model import BaseModel
import network.DuRN_Pure_Conv_3_coarse_3 as coarse_3
from network.metrics import ssim, L1_loss
from network.weights_init import init_weights
from utils.torch_utils import ExponentialMovingAverage, print_network

models = {
    'default': pure.cleaner(),
    'pure': pure.cleaner(),
    'coarse': coarse_3.cleaner(),
    'FFA': model_zoo.FFA(),
    'Nested': model_zoo.NestedUNet(),
}

class Model(BaseModel):
    def __init__(self, opt):
        super(Model, self).__init__(opt)

        # self.cleaner = Pyramid_Net(3, 256).cuda(device=opt.device)
        self.cleaner = models[opt.model].cuda(device=opt.device)
        if opt.init:
            init_weights(self.cleaner, opt.init)

        print_network(self.cleaner)

        params = {
            'dim': 64,
            'norm': 'none',
            'activ': 'lrelu',  # activation function [relu/lrelu/prelu/selu/tanh]
            'n_layer': 4,  # number of layers in D
            'gan_type': 'lsgan',  # GAN loss [lsgan/nsgan]
            'num_scales': 3,  # number of scales
            'pad_type': 'reflect'
        }

        # self.discriminitor = MsImageDis(input_dim=3, params=params).cuda(device=opt.device)
        #
        # print(self.discriminitor)

        self.g_optimizer = optim.Adam(self.cleaner.parameters(), lr=opt.lr)
        # self.d_optimizer = optim.Adam(cleaner.parameters(), lr=opt.lr)

        # load networks
        if opt.load:
            pretrained_path = opt.load
            self.load_network(self.cleaner, 'G', opt.which_epoch, pretrained_path)
            # if self.training:
            #     self.load_network(self.discriminitor, 'D', opt.which_epoch, pretrained_path)

        self.avg_meters = ExponentialMovingAverage(0.95)
        self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)

    def update_G(self, x, y):

        # L1 & SSIM loss
        some_loss = 0
        cleaned = self.cleaner(x)
        ssim_loss_r = -ssim(cleaned, y)
        ssim_loss = ssim_loss_r * 1.1

        # Compute L1 loss (not used)
        l1_loss = L1_loss(cleaned, y)
        l1_loss = l1_loss * 0.75

        loss = ssim_loss + l1_loss

        # GAN loss
        # loss_gen_adv = self.discriminitor.calc_gen_loss(input_fake=cleaned)
        self.avg_meters.update({'ssim': -ssim_loss_r.item(), 'L1': l1_loss.item()})

        #loss_gen = loss + loss_gen_adv * 1.
        self.g_optimizer.zero_grad()
        loss.backward()
        self.g_optimizer.step()

        return cleaned

    def update_D(self, x, y):
        self.d_optimizer.zero_grad()
        # encode
        cleaned = cleaner(x)
        # h_b, n_b = self.gen_b.encode(x_b)
        # decode (cross domain)

        # D loss
        loss_dis = self.discriminitor.calc_dis_loss(input_fake=cleaned, input_real=y)
        self.avg_meters.update({'dis': loss_dis})

        loss_dis = loss_dis * 1.  # weights
        loss_dis.backward()
        self.d_optimizer.step()
        return cleaned

    def discriminate(self, input_label, test_image, use_pool=False):
        input_concat = torch.cat((input_label, test_image.detach()), dim=1)
        if use_pool:
            fake_query = self.fake_pool.query(input_concat)
            return self.netD.forward(fake_query)
        else:
            return self.netD.forward(input_concat)

    def forward(self, x):
        return self.cleaner(x)

    def inference(self, x, image=None):
        pass

    def save(self, which_epoch):
        self.save_network(self.cleaner, 'G', which_epoch)
        # self.save_network(self.discriminitor, 'D', which_epoch)

    def update_learning_rate(self):
        lrd = self.opt.lr / self.opt.niter_decay
        lr = self.old_lr - lrd
        for param_group in self.d_optimizer.param_groups:
            param_group['lr'] = lr
        for param_group in self.g_optimizer.param_groups:
            param_group['lr'] = lr
        if self.opt.verbose:
            print('update learning rate: %f -> %f' % (self.old_lr, lr))
        self.old_lr = lr