from sklearn.model_selection import train_test_split
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import pickle
import torch
import torch.nn as nn
import torch.optim as optim
from train_lgbm import *
import lightgbm as lgb
from  torch.utils.data import TensorDataset,DataLoader
import torch.nn.functional as F
from sklearn.model_selection import train_test_split
import math
import time
import random
from sklearn.preprocessing import *
from model_train_util_vectorized import *
import sys
import os
import os.path as osp
import time
import datetime
from itertools import repeat
from typing import Optional
import gc
import pickle
from sklearn.preprocessing import *
from sys import argv

Device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
print('using device:', Device)
dataset_name = argv[1]
use_gbm = argv[2]
use_botspotpp = argv[3]

ROOT_DIR = f'{dataset_name}/'

class BotSpotTrans(object):
    def __init__(self,train_filepath,test_filepath):

        train_file = train_filepath
        test_file = test_filepath
        print("Loading train & test file...")
        train_df = pd.read_csv(train_file) # ѵ����
        test_df = pd.read_csv(test_file) # ���Լ�
        total_df = pd.concat([train_df, test_df], axis=0)

        print("Graph Generating...")
        self.edge_index = total_df[["combin_index", "device_index", "target"]].astype(int).values
        self.edge_index_train = train_df[["combin_index", "device_index", "target"]].astype(int).values
        self.edge_index_test = test_df[["combin_index", "device_index", "target"]].astype(int).values

        # stat_columns��ʾ����ͳ����ص�������list
        # category_columns��ʾ����category��������list
        # stat_columns���������ctit��cvr_total, category_columns��һ����label encoder���channel_id

        stat_columns_file = osp.join(ROOT_DIR, "stat_columns.txt")
        category_columns_file = osp.join(ROOT_DIR, "category_columns.txt")
        stat_columns = self.pickle_load(stat_columns_file)
        category_columns = self.pickle_load(category_columns_file)

        # ����������
        feature_columns = stat_columns + category_columns

        # �������ֻ��ctit����normalization
        normalized_columns = [stat_columns[-2]]
        except_normalized_columns = [column for column in stat_columns if column not in normalized_columns]

        # channel-campagin��ص�����ѵ����������, ������ctit��ͳ������channel_id(���һ��)
        combin_feature_columns = except_normalized_columns + [category_columns[0]]
        # device�ڵ���ص�����ѵ���������У�����ctit + ʣ�µ�category��
        device_feature_columns = normalized_columns + category_columns[1:]
        #���е�channel-campaign��
        device_columns = ["device_index"] + device_feature_columns
        # ���е�device����
        combin_columns = ["combin_index"] + combin_feature_columns

        #��������Ǵ�total_df���ֳ�device_df,���������Ҫ�������train_dfҲһ��
        device_df = total_df[device_columns].sort_values(["device_index"])
        device_df.drop_duplicates(subset="device_index", keep="first", inplace=True)

        #ͬ��Ҳ�Ǵ�total_df���ֳ���combin_df,���������Ҫ���е���
        combin_df = total_df[combin_columns].sort_values(["combin_index"])
        combin_df.drop_duplicates(subset="combin_index", keep="first", inplace=True)

        # ��ctit�н��������С��һ��
        norm_data = RobustScaler().fit_transform(device_df.loc[:,normalized_columns[0]].values.reshape((-1,1)))
        device_df.loc[:,normalized_columns[0]] = norm_data.reshape(-1,)

        print("feature matrix generating...")
        self.device_matrix = device_df[device_feature_columns].values
        # �����folat�ľ���ת�����float16�����Ը��������Ҫ����
        self.combin_matrix = combin_df[combin_feature_columns].astype('float16').values

    def pickle_dump(self, data, filename):
        with open(filename, 'wb') as f:
            pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)

    def pickle_load(self, filename):
        with open(filename, 'rb') as f:
            return pickle.load(f)

class Self_Attention(nn.Module):
    # a classical self_attention module as depicted in BERT
    def __init__(self, emb_size, num_head):
        super(Self_Attention, self).__init__()
        self.Q = nn.ModuleList([])
        self.K = nn.ModuleList([])
        self.V = nn.ModuleList([])
        output_size = emb_size // num_head
        self.output_size = output_size
        self.num_head = num_head
        self.final_linear = nn.Linear(emb_size, emb_size)

        for i in range(num_head):
            self.Q.append(nn.Linear(emb_size, output_size))
            self.K.append(nn.Linear(emb_size, output_size))
            self.V.append(nn.Linear(emb_size, output_size))

    def calc_attention(self, X, Q, K, V):
        query = Q(X)
        key = K(X)
        value = V(X)
        key_ = key.transpose(2, 1).contiguous()
        attn_weights = torch.softmax(torch.bmm(query, key_) / math.sqrt(self.output_size), dim=-1)
        output = torch.bmm(attn_weights, value)
        return output

    def forward(self, X):
        outputs = []
        for i in range(self.num_head):
            q, k, v = self.Q[i], self.K[i], self.V[i]
            out = self.calc_attention(X, q, k, v)
            outputs.append(out)
        outputs = torch.cat(outputs, dim=-1)
        return self.final_linear(outputs).mean(dim=1)

class BotSpot(nn.Module):
    def __init__(self,edge_index_train,edge_index_test,channel_feats,device_feats,device_package = None,use_enhance_botspot = False,use_gbm = False,gbm_model = None):
        super(BotSpot,self).__init__()
        if use_gbm:
            leaf_dim = 20
            assert gbm_model is not None
            self.leaf_emb_models = nn.ModuleList()
            for n in range(gbm_model.n_estimators):
                self.leaf_emb_models.append(nn.Embedding(31, leaf_dim)) # 31 is the max depth of decision tree
        self.gbm_best_model = gbm_model
        self.use_gbm = use_gbm
        self.self_attn_module = Self_Attention(32, 2)
        self.use_enhanced_botspot = use_enhance_botspot
        self.edge_index_train = edge_index_train
        self.edge_index_test = edge_index_test
        self.edge_num_train = edge_index_train.shape[0]
        self.edge_num_test = edge_index_test.shape[0]
        self.channel_feats = torch.from_numpy(channel_feats).float() # feature matrix for channel-campaign nodes
        self.device_feats = torch.from_numpy(device_feats).float()  # package_name already removed in the device feature matrix
        N_chan = channel_feats.shape[1]-1 + 16 # remove one col and add embeeding_size of 16
        N_dev = 16*(device_feats.shape[1]-1)+1 # add one col of ctit, others are embeddings
        self.edge_matrix = self.gen_adj_matrix() # generate adj matrix of shape (N_device,N_channel_campaign)
        self.device_split_val = 1 # the first col is ctit for device
        self.channel_split_val = -1 # the last col is channel id
        self.channel_neibr_cache = {}  # cache device neighbors as a boolean array of channel_campaign nodes for training stage
        self.channel_neibr_test_cache = {}  # cache neighbors as boolean array of channel_campaign nodes for test stage
        self.device_idx_cache = {} # cache device indices for a channel-campaign node
        if use_enhance_botspot:
            assert device_package is not None
            self.device_package = device_package
            self.super_device_neibr_cache = {} # cache channel_campaign indices for each super device node for training set
            self.super_device_neibr_test_cache = {} # cache channel_campaign indices for each super device node for testing set

        # precompute channel neighbors and  cache it
        N = self.edge_matrix.shape[1]
        for i in range(N):
            self.channel_neibr_cache[i] = self.edge_matrix[:,i]
#         N_pack = len(device_package)
        # precompute the super device neighbors for training set if botspot++ is used
        if self.use_enhanced_botspot:
            for idx,d in enumerate(device_package):
                if d not in self.super_device_neibr_cache:
                    self.super_device_neibr_cache[d] = set()
                t_ = np.where(self.edge_matrix[idx]==1)[0]
                for v in t_:
                    self.super_device_neibr_cache[d].add(v)

            for d in self.super_device_neibr_cache:
                self.super_device_neibr_cache[d] = list(self.super_device_neibr_cache)

        self.set_adj_test_matrix(True) # modify edge_matrix to include test edges
        for i in range(N):
            self.channel_neibr_test_cache[i] = self.edge_matrix[:,i]
        # precompute super device neighbors for test set, cache them in super_device_neibr_test_cache
        if self.use_enhanced_botspot:
            for idx,d in enumerate(device_package):
                if d not in self.super_device_neibr_test_cache:
                    self.super_device_neibr_test_cache[d] = set()
                t_ = np.where(self.edge_matrix[idx]==1)[0]
                for v in t_:
                    self.super_device_neibr_test_cache[d].add(v)

            for d in self.super_device_neibr_test_cache:
                self.super_device_neibr_test_cache[d] = list(self.super_device_neibr_test_cache)
        self.set_adj_test_matrix(False) # modify the edge matrix to exclude test edges to proceed on training

        # initialze embedding matrix
        emb_size = 16
        channel_id_max = int(channel_feats[:, -1].max() + 1)
        temp = np.max(device_feats[:, 1:], axis=0) + 1
        temp = [int(i) for i in temp]
        lang, plat, os, country, carrier, device_brand, plat_os = temp  # bypass install city, be careful
        self.channel_id_emb = nn.Embedding(channel_id_max, emb_size)
        self.carrier_emb = nn.Embedding(carrier, emb_size)
        self.language_emb = nn.Embedding(lang, emb_size)
        self.device_brand_emb = nn.Embedding(device_brand, emb_size)
        self.plat_os_emb = nn.Embedding(plat_os, emb_size)
        self.plat_emb = nn.Embedding(plat, emb_size)
        self.os_emb = nn.Embedding(os, emb_size)
        self.country_emb = nn.Embedding(country, emb_size)

        # device modules if there is no super device convolution
        if not self.use_enhanced_botspot:
            self.dev_linear1 = nn.Linear(N_dev,int(0.6*N_dev)) # NOT also used in channel side for convolving device feats
            self.dev_relu1 = nn.ReLU()
            self.dev_dropout1 = nn.Dropout(0.2)
            self.dev_linear2 = nn.Linear(int(0.6*N_dev),int(0.75*0.6*N_dev))
            self.dev_relu2 = nn.ReLU()
        # channel linear and message passing modules
        self.channel_linear1 = nn.Linear(N_chan,int(0.6*N_chan))
        self.channel_msg_pass1 = nn.Linear(N_dev,int(0.6*N_dev))
        fusion_input = int(0.6*N_chan) + int(0.6*N_dev)
        self.fusion_linear1 = nn.Linear(fusion_input,int(0.6*fusion_input))
        self.fusion_relu1 = nn.ReLU()
        self.fusion_dropout1 = nn.Dropout(0.2)
        fusion_output_dim =  int(0.6*fusion_input)
        device_output_dim = int(0.75*0.6*N_dev)

        if not self.use_enhanced_botspot:
            # concat modules if no botspot++
            concat_input_dim = fusion_output_dim+device_output_dim if not self.use_gbm else fusion_output_dim+device_output_dim+leaf_dim
            self.concat_linear1 = nn.Linear(concat_input_dim,int(0.6*concat_input_dim))
            self.concat_relu1 = nn.ReLU()
            self.concat_linear2 = nn.Linear(int(0.6*concat_input_dim),int(0.5*0.6*concat_input_dim))
            self.concat_relu2 = nn.ReLU()
            self.concat_linear3 = nn.Linear(int(0.5*0.6*concat_input_dim),1)
        else:
            # device side gnn if botspot++ is used
            self.dev_linear1 = nn.Linear(N_dev,int(0.6*N_dev)) # NOT also used in channel side for convolving device feats
#             self.device_msg_passing = nn.Linear(N_chan,int(0.6*N_chan))
            in_dim = int(0.6*N_dev) + int(0.6*N_chan)
            self.sup_dev_fusion_linear1 = nn.Linear(in_dim,int(0.6*in_dim))
            self.sup_dev_fusion_relu1 = nn.ReLU()
            self.sup_dev_fusion_dropout1 = nn.Dropout(0.2)

            # concat layer for botspot++
            sup_dev_fusion_output_dim = int(0.6*in_dim)
            concat_input_dim = fusion_output_dim+sup_dev_fusion_output_dim if not self.use_gbm else fusion_output_dim+sup_dev_fusion_output_dim+ leaf_dim
            self.concat_linear1 = nn.Linear(concat_input_dim,int(0.6*concat_input_dim))
            self.concat_relu1 = nn.ReLU()
            self.concat_linear2 = nn.Linear(int(0.6*concat_input_dim),int(0.5*0.6*concat_input_dim))
            self.concat_relu2 = nn.ReLU()
            self.concat_linear3 = nn.Linear(int(0.5*0.6*concat_input_dim),1)

    def to_emb(self, arr, *models):
        '''
        :param arr: matrix for holding high-cardinality features, without one-hot encoding
        :param left: channel node if left is True else device node
        :param models: a list of embedding matrices to embed each high-cardinality feature to dense embeddings
        :return: 2-d tensor with dense embeddings for all the high-cardinality features.
        '''

        out_arr = []
#         arr = torch.from_numpy(arr)
        arr = arr.long().to(Device)
        # device node sparse features

        # N = arr.shape[0]
        num_models = len(models)
        for i in range(len(models)):
            # if num_models > 2 and i == 4:  # bypass install city, hardcoded
            #     continueed
            #             print (i,models[i])
            out_arr.append(models[i](arr[:, i]))
        return torch.cat(out_arr, dim=1)

    def concat_device_feats(self, dev_feats):  # NEED TO MODIFY IT
        '''
        this method invokes to_emb to embed device categorical features into dense embeddings
        :param dev_feats: normalized device features
        :param more_dev_feats: feature matrix with high-cardinality features
        :return: feature matrix with dense embeddings
        '''
        dev_feats = dev_feats.to(Device)
        cat_dev_feats = dev_feats[:, self.device_split_val:]
        emb_tensor = self.to_emb(cat_dev_feats, self.language_emb,
                                 self.plat_emb, self.os_emb, self.country_emb, self.carrier_emb,
                                 self.device_brand_emb, self.plat_os_emb)

        dev_emb_feats = torch.cat((dev_feats[:, :self.device_split_val], emb_tensor), dim=1).float().to(Device)
        return dev_emb_feats

    def concat_channel_feats(self, chan_feats):
        '''
        this method invokes to_emb to embed channel_campaign node's categorical feature into dense embeddings
        similar to concat_device_feats, to add dense embeddings
        '''
        chan_feats = chan_feats.to(Device)
        emb_tensor = self.to_emb(chan_feats[:, self.channel_split_val:], self.channel_id_emb)
#         print (chan_feats[:, :self.channel_split_val].shape)
#         print (emb_tensor.shape)
        return torch.cat((chan_feats[:, :self.channel_split_val], emb_tensor), dim=1).float().to(Device)

    def gen_adj_matrix(self):
        """
        this method build ajdacency matrix of shape (N_device,N_channel).
        adj_matrix[j,i] = True indicates device node j connects to channel-campaign node i
        """
        e = np.vstack((self.edge_index_train,self.edge_index_test))
        N_dev = np.max(e[:,1])+1
        N_channel = np.max(e[:,0])+1
        adj_matrix = np.zeros((N_dev,N_channel),dtype = bool)
        for i,j,_ in self.edge_index_train:
            adj_matrix[j,i] = True
        return adj_matrix

    def set_adj_test_matrix(self,set_test = True):
        """
        this method modifies edge_matrix
        when set_test = True, edge_matrix includes test edges in it, otherwise does not
        """
        count_not_in_train = 0
        for i,j,_ in self.edge_index_test:
            if set_test:
                self.edge_matrix[j,i] = True
            else:
                self.edge_matrix[j,i] = False
        return

    def sample_minibatch(self,channel_vertices,device_vertices,train_stage = True):
        """
        input:
        for a minitach of edges, channel_vertices is edge[:,0], device_vertices is edge[:,1]
        this method takes a minibatch of edges outputs:
        1) features for channel_campaign nodes
        2) features for device nodes
        3) neighboring device features for each channel_campaign node
        4) neighboring channel_campaign features for each super device node
        """
        #set number of neighbors for channel and device for different stages
        num_neibr = 50 if train_stage else 800
        sup_dev_num_neibr = 20 if train_stage else 50
        edge_matrix = self.edge_matrix

        # channel_vertices and device_vertices must be numpy array
        # original features
        neibr_feats_tensor = []
        sup_dev_neibr_feats_tensor = []
        minibatch_channel_feats = self.concat_channel_feats(self.channel_feats[channel_vertices]).to(Device)  # shape of (minibatch,feats_num_channel)
        minibatch_device_feats = self.concat_device_feats(self.device_feats[device_vertices]).to(Device) # shape of (minibatch,feats_num_device)
        channel_vertices = channel_vertices.cpu().numpy()

        # for each channel vertices, get neighbor devices and its features
        for i in channel_vertices:
            if train_stage:
                neibr_feats = self.adj_indice_to_feat_mat(self.channel_neibr_cache[i],num_neibr,i)
                neibr_feats_tensor.append(neibr_feats)
            else:
                neibr_feats = self.adj_indice_to_feat_mat(self.channel_neibr_test_cache[i],num_neibr,i)
                neibr_feats_tensor.append(neibr_feats)
        neibr_feats_tensor = torch.cat(neibr_feats_tensor,dim = 0).to(Device) # neibr_feats_tensor would be a tensor of shape: (Minibatch,num_neibr,feats_num_device)

        # if use botspot++, for each device node, retrieve its super device index and get its neighboring channel_campaign node features
        if self.use_enhanced_botspot:
            for i in device_vertices:
                sup_dev_neibr_feats = self.adj_indice_to_feat_mat_super_device(sup_dev_num_neibr,i,train_stage)
                sup_dev_neibr_feats_tensor.append(sup_dev_neibr_feats)
            sup_dev_neibr_feats_tensor = torch.cat(sup_dev_neibr_feats_tensor,dim = 0).to(Device)
            return minibatch_channel_feats,minibatch_device_feats,neibr_feats_tensor,sup_dev_neibr_feats_tensor

        return minibatch_channel_feats,minibatch_device_feats,neibr_feats_tensor,-1

    def adj_indice_to_feat_mat(self,col,neibr_num,indice):
        """
        input:
        indice: indice of channel_campaign node
        neibr_num: num of sample size
        col: a boolean array of a col given channel_campaign node of $indice
        return:
        neighboring device features
        """
        if indice in self.device_idx_cache:
            ind = self.device_idx_cache[indice]
        else:
            ind = np.where(col==True)[0]
            self.device_idx_cache[indice] = ind
        if len(ind)>neibr_num:
            np.random.shuffle(ind)
            ind_subset = ind[:neibr_num]
        else:
            ind_subset = np.random.choice(ind,size = neibr_num,replace = True)

        dev_feats = self.concat_device_feats(self.device_feats[ind_subset].to(Device))
        return dev_feats.unsqueeze(0)

    def adj_indice_to_feat_mat_super_device(self,neibr_num,device_idx,train_stage = True):
        """
        input:
        device_idx: index of device node
        neibr_num: num of sample size
        return:
        neighboring channel_campaign features
        """
        i = self.device_package[device_idx]
        if train_stage:
            sup_dev_neibr = list(self.super_device_neibr_cache[i])
        else:
            sup_dev_neibr = list(self.super_device_neibr_test_cache[i])
        if len(sup_dev_neibr)>neibr_num:
            random.shuffle(sup_dev_neibr)
            c = self.concat_channel_feats(self.channel_feats[sup_dev_neibr[:neibr_num]].to(Device))
        else:
            sup_dev_neibr = np.random.choice(np.asarray(sup_dev_neibr),size = neibr_num,replace = True)
            c = self.concat_channel_feats(self.channel_feats[sup_dev_neibr].to(Device))
        return c.unsqueeze(0)

    def get_leaf_from_light_gbm(self, left_vertices, right_vertices, use_self_attn=False):
        # get leaf indices from gbm model and embed into dense matrix
        output_leaf_emb = []
        chan_data = self.channel_feats[left_vertices]
        dev_data = self.device_feats[right_vertices]
        try:
            edge_data = np.hstack((chan_data, dev_data))
        except:
            edge_data = torch.cat((chan_data, dev_data),
                                  dim=1)  # edge feature is the concatenation of channel_node and device_node
            edge_data = edge_data.cpu().numpy()
        # N = len(left_vertices)
        if len(edge_data.shape)==1:
            edge_data = edge_data.reshape((1,-1))
        pred_leaf = self.gbm_best_model.predict_proba(edge_data, pred_leaf=True)
        pred_leaf = torch.from_numpy(pred_leaf).long().to(Device)

        for i in range(pred_leaf.shape[1]):
            # print (self.leaf_emb_models[i](pred_leaf[:, i]).shape)
            output_leaf_emb.append(self.leaf_emb_models[i](pred_leaf[:, i]).unsqueeze(1))
            # ret = torch.cat(output_leaf_emb, dim=1).to(Device)  # leaf node concatenation
        if not use_self_attn:
            ret = torch.cat(output_leaf_emb, dim=1).mean(axis=1).to(Device)  # leaf node mean pooling
            return ret
        else:
            ret = torch.cat(output_leaf_emb, dim=1).to(Device)
            out = self.self_attn_module(ret)
            return out

    def forward(self,edges,train_stage = True):
        minibatch_channel_feats,minibatch_device_feats,neibr_feats_tensor,sup_dev_neibr_feats_tensor = self.sample_minibatch(edges[:,0],edges[:,1],train_stage)
        # forward device feats
        if not self.use_enhanced_botspot:
            device_out = self.dev_relu2(self.dev_linear2(self.dev_dropout1(self.dev_relu1(self.dev_linear1(minibatch_device_feats)))))
            channel_conv = self.channel_linear1(minibatch_channel_feats)
            dev_conv = self.dev_linear1(neibr_feats_tensor).mean(dim=1) # share dev_linear1
            fuse_conv = self.fusion_dropout1(self.fusion_relu1(self.fusion_linear1(torch.cat((channel_conv,dev_conv),dim=1))))
            if not self.use_gbm:
                h = self.concat_linear3(self.concat_relu2(self.concat_linear2(self.concat_relu1(self.concat_linear1(torch.cat((fuse_conv,device_out),dim=1))))))
            else:
                leaf_out = self.get_leaf_from_light_gbm(edges[:,0],edges[:,1])
                h = self.concat_linear3(self.concat_relu2(self.concat_linear2(self.concat_relu1(self.concat_linear1(torch.cat((fuse_conv,device_out,leaf_out),dim=1))))))
            return torch.sigmoid(h)
        else:
            channel_conv = self.channel_linear1(minibatch_channel_feats)
            dev_conv = self.dev_linear1(neibr_feats_tensor).mean(dim=1) # share dev_linear1
            fuse_conv = self.fusion_dropout1(self.fusion_relu1(self.fusion_linear1(torch.cat((channel_conv,dev_conv),dim=1))))
            # device side conv:
            sup_dev_conv = self.dev_linear1(minibatch_device_feats)
            sup_channel_conv = self.channel_linear1(sup_dev_neibr_feats_tensor).mean(dim=1)
            sup_fuse_conv = self.sup_dev_fusion_dropout1(self.sup_dev_fusion_relu1(self.sup_dev_fusion_linear1(torch.cat((sup_channel_conv,sup_dev_conv),dim=1))))
            if not self.use_gbm:
                h = self.concat_linear3(self.concat_relu2(self.concat_linear2(self.concat_relu1(self.concat_linear1(torch.cat((fuse_conv,sup_fuse_conv),dim=1))))))
            else:
                leaf_out = self.get_leaf_from_light_gbm(edges[:,0],edges[:,1])
                h = self.concat_linear3(self.concat_relu2(self.concat_linear2(self.concat_relu1(self.concat_linear1(torch.cat((fuse_conv,sup_fuse_conv,leaf_out),dim=1))))))
            return torch.sigmoid(h)

if __name__ == '__main__':
    bot_preprocess = BotSpotTrans(f'{dataset_name}/train.csv', f'{dataset_name}/test.csv')
    edge_index_train = bot_preprocess.edge_index_train
    edge_index_test = bot_preprocess.edge_index_test
    chan_feats = bot_preprocess.combin_matrix
    device_feats = bot_preprocess.device_matrix
    device_package = list(device_feats[:, -1])
    device_feats = device_feats[:, :-1]
    if use_gbm:
        gbm_model = make_dataset(chan_feats,device_feats,-1,1,edge_index_train,edge_index_test)
    else:
        gbm_model  = None
    botspot_model = BotSpot(edge_index_train, edge_index_test, chan_feats, device_feats,device_package = device_package,use_enhance_botspot = use_botspotpp,use_gbm = use_gbm,gbm_model = gbm_model)

    tr_dset = TensorDataset(torch.from_numpy(edge_index_train))
    tr_dloader = DataLoader(tr_dset, batch_size=500, shuffle=True)
    test_dset = TensorDataset(torch.from_numpy(edge_index_test))
    test_dloader = DataLoader(test_dset, batch_size=500, shuffle=False)
    optimizer = optim.Adam(botspot_model.parameters(), lr=2e-4, weight_decay=3e-6)
    # model_states = torch.load('model_results_1014/botspotpp/model_checkpoints_4.pt')
    # botspot_model.load_state_dict(model_states)

    botspot_model.to(Device)
    try:
        os.mkdir('model_results')
    except:
        pass
    _ = train(botspot_model, tr_dloader, test_dloader, optimizer, 10, save_name='model_results')


from __future__ import absolute_import, print_function, unicode_literals

import argparse
import os
import sys
import yaml

# Get relative imports to work when the package is not installed on the PYTHONPATH.
if __name__ == "__main__" and __package__ is None:
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# pylint: disable=wrong-import-position
from buildscripts.ciconfig.evergreen import parse_evergreen_file
# pylint: enable=wrong-import-position

# Name of map to search for in the variables map in evergreen.yml
MONGOCRYPTD_VARIANTS = "mongocryptd_variants"
PUSH_TASK_NAME = "push"

def can_validation_be_skipped(evg_config, variant):
    """
    Determine if the given build variant needs to be validated.

    A build variant does not need to be validated if it does not run the 'push' task or
    if it does not exist in the configuration (it is dynamically created).

    :param evg_config: Evergreen configuration.
    :param variant: Build variant to check.
    :return: True if validation can be skipped.
    """
    variant_config = evg_config.get_variant(variant)
    if not variant_config:
        return True

    if PUSH_TASK_NAME not in variant_config.task_names:
        return True

    return False

def read_variable_from_yml(filename, variable_name):
    """
    Read the given variable from the given yaml file.

    :param filename: Yaml file to read from.
    :param variable_name: Variable to read from file.
    :return: Value of variable or None.
    """
    with open(filename, 'r') as fh:
        nodes = yaml.safe_load(fh)

    variables = nodes["variables"]

    for var in variables:
        if variable_name in var:
            return var[variable_name]
    return None

def main():
    # type: () -> None
    """Execute Main Entry point."""

    parser = argparse.ArgumentParser(description='MongoDB CryptD Check Tool.')

    parser.add_argument('file', type=str, help="etc/evergreen.yml file")
    parser.add_argument('--variant', type=str, help="Build variant to check for")

    args = parser.parse_args()

    expected_variants = read_variable_from_yml(args.file, MONGOCRYPTD_VARIANTS)
    if not expected_variants:
        print("ERROR: Could not find node %s in file '%s'" % (MONGOCRYPTD_VARIANTS, args.file),
              file=sys.stderr)
        sys.exit(1)

    evg_config = parse_evergreen_file(args.file)
    if can_validation_be_skipped(evg_config, args.variant):
        print(f"Skipping validation on buildvariant {args.variant}")
        sys.exit(0)

    if args.variant not in expected_variants:
        print("ERROR: Expected to find variant %s in list %s" % (args.variant, expected_variants),
              file=sys.stderr)
        print(
            "ERROR:  Please add the build variant %s to the %s list in '%s'" %
            (args.variant, MONGOCRYPTD_VARIANTS, args.file), file=sys.stderr)
        sys.exit(1)

    sys.exit(0)

if __name__ == '__main__':
    main()


""" Provides an intermediate datastructure close to the PDF, which is used
to transform the lines in the pdf into TimeTable objects. """

from __future__ import annotations

import logging
from operator import attrgetter
from pathlib import Path
from typing import Callable, TypeAlias

from pdf2gtfs.config import Config
from pdf2gtfs.datastructures.pdftable.container import (
    Column, FieldContainer, Row)
from pdf2gtfs.datastructures.pdftable.enums import (
    ColumnType, FieldType,
    RowType,
    )
from pdf2gtfs.datastructures.pdftable.lists import ColumnList, RowList
from pdf2gtfs.datastructures.timetable.table import TimeTable

logger = logging.getLogger(__name__)
Tables: TypeAlias = list["PDFTable"]
Rows: TypeAlias = list[Row]
Cols: TypeAlias = list[Column]
Splitter: TypeAlias = Callable[[Tables, list[FieldContainer]], None]

class PDFTable:
    """ Describes a table, using coordinates, rows and columns. """

    def __init__(self, rows: Rows = None, columns: Cols = None):
        self.rows = rows or []
        self.columns = columns or []

    @property
    def rows(self) -> RowList:
        """ The rows of the table. """
        return self._rows

    @rows.setter
    def rows(self, rows: Rows | RowList) -> None:
        if isinstance(rows, RowList):
            self._rows = rows
        else:
            self._rows = RowList.from_list(self, rows)

    @property
    def columns(self) -> ColumnList:
        """ The columns of the table. """
        return self._columns

    @columns.setter
    def columns(self, columns: Cols | ColumnList):
        if isinstance(columns, ColumnList):
            self._columns = columns
        else:
            self._columns = ColumnList.from_list(self, columns)

    @property
    def empty(self) -> bool:
        """ Whether either columns or rows are empty. """
        return self.columns.empty or self.rows.empty

    def generate_columns_from_rows(self) -> None:
        """ Create columns from the given rows. """

        def _generate_single_field_columns() -> Cols:
            """ Generate single-field columns from the rows. """
            field_columns = [Column.from_field(self, field)
                             for row in rows for field in row]
            return sorted(field_columns, key=attrgetter("bbox.x0"))

        def _merge_overlapping_columns(field_columns: Cols) -> Cols:
            """ Merges overlapping field_columns. """
            first_field = field_columns.pop(0).fields[0]
            cols: Cols = [Column.from_field(self, first_field)]

            for column in field_columns:
                previous_column = cols[-1]
                # Begin new column, if the current column is not overlapping.
                if previous_column.bbox.x1 <= column.bbox.x0:
                    cols.append(Column.from_field(self, column.fields[0]))
                    continue

                previous_column.add_field(column.fields[0])
            return cols

        rows = self.rows.of_types(
            [RowType.DATA, RowType.ANNOTATION, RowType.ROUTE_INFO])
        if not rows:
            return

        columns = _merge_overlapping_columns(_generate_single_field_columns())
        self.columns = columns

    def fix_split_stopnames(self) -> None:
        """ Fix stop names (indented or starting with a delimiter),
        indicating they use the same city/POI as the previous stop. """
        stop_column = self.columns.of_type(ColumnType.STOP)[0]
        first_field_idx = 0
        reference_field = None
        for i, field in enumerate(stop_column.fields):
            if field.row.type == RowType.DATA:
                reference_field = field
                first_field_idx = i
                break

        if not reference_field:
            return

        for field in stop_column.fields[first_field_idx:]:
            # Don't update the reference_field, in case field is indented.
            if not field.fix_name_if_split(reference_field):
                reference_field = field

    def to_timetable(self) -> TimeTable:
        """ Creates a TimeTable containing the values of this table. """
        return TimeTable.from_pdf_table(self)

    def get_header_from_column(self, column: Column) -> str:
        """ Returns the header text of the given column. """
        for row in self.rows.of_type(RowType.HEADER):
            for i, field in enumerate(row, 1):
                next_field = row.fields[i] if i < len(row.fields) else None
                if not next_field or next_field.bbox.x0 >= column.bbox.x1:
                    return field.text

        return ""

    def add_row_or_column(self, obj: Row | Column) -> None:
        """ Add the object to either rows or columns, based on its type. """
        if isinstance(obj, Row):
            self.rows.add(obj)
            return
        self.columns.add(obj)

    @staticmethod
    def _split_at(splitter: list, splitter_func: Splitter) -> Tables:
        tables: Tables = [PDFTable() for _ in splitter]

        splitter_func(tables, splitter)

        for table in tables:
            for row in table.rows:
                row.update_type()
            table.generate_columns_from_rows()

        return tables

    def split_at_stop_columns(self) -> Tables:
        """ Return a list of tables with each having a single stop column. """

        def splitter(tables: Tables, splitter_columns: list) -> None:
            """ Split the given tables at the given splitter_columns. """
            for row in self.rows:
                splits = row.split_at(splitter_columns)
                for table, split in zip(tables, splits):
                    if not split.fields:
                        continue
                    table.add_row_or_column(split)

        return self._split_at(self.columns.of_type(ColumnType.STOP), splitter)

    def split_at_header_rows(self) -> Tables:
        """ Return a list of tables with each having a single header row. """

        def splitter(tables: Tables, splitter_rows: list) -> None:
            """ Splits the current tables' rows such that each split starts
            with a splitter_row and assigns each split to a table. """
            rows_list = [[] for _ in splitter_rows]
            first_is_splitter = self.rows[0] in splitter_rows
            idx = -1 if first_is_splitter else 0

            for row in self.rows:
                if row in splitter_rows:
                    idx += 1
                rows_list[idx].append(row)

            for table, rows in zip(tables, rows_list, strict=True):
                table.rows = rows

        return self._split_at(self.rows.of_type(RowType.HEADER), splitter)

    def to_file(self, fname: Path) -> None:
        """ Export the PDFTable to the given path in the csv format. """
        def escape_field_text(text: str) -> str:
            """ Wrap field text that contains a comma in quotes.

            Also removes any existing quotes.
            """
            text = text.replace('"', "").strip()
            if "," in text:
                return f'"{text}"'
            return text

        def get_header_row_column_idx(idx: int = 0, *, row_field=None) -> int:
            """ Get the column index of the fields of a header row. """
            if not row_field:
                row_field = row.fields[idx]
            for i, col_ in enumerate(self.columns):
                if col_.bbox.x0 > row_field.bbox.x0:
                    return i
            # Last column.
            return len(self.columns) - 1

        row_strings = [[] for _ in self.rows]
        seen_fields = []
        for col in self.columns:
            last_row_id = -1
            for field in col:
                row_id = field.row.index
                for row in self.rows.objects[last_row_id + 1: row_id]:
                    row_strings[row.index].append("")
                row_strings[row_id].append(escape_field_text(field.text))
                seen_fields.append(field)
                last_row_id = row_id
            # Fix missing trailing commas.
            for row_string in row_strings:
                while len(row_string) < self.columns.index(col) + 1:
                    row_string.append("")
        # Add fields without a column.
        for row in self.rows:
            for field in row:
                if field.type != FieldType.HEADER or field in seen_fields:
                    continue
                idx = get_header_row_column_idx(row_field=field)
                row_strings[row.index][idx] = field.text
        # Export.
        table_str = "\n".join(
            [",".join(row_string) for row_string in row_strings
             if any(row_string)]) + "\n"
        with open(fname, "w") as fil:
            fil.write(table_str)

def split_rows_into_tables(rows: Rows) -> Tables:
    """ Split raw rows into (multiple) PDFTable. A new table is created,
    whenever the distance to the previous line is higher than the defined
    max_row_distance. Tables with too few rows are dropped with an info. """

    def log_skipped_rows() -> None:
        """ Log the rows that are dropped. """
        row_str = ",\n\t\t  ".join([str(r) for r in current_rows])
        logger.debug(f"Dropped rows:\n\tDistance: {y_distance:.2f}"
                     f"\n\tRows: {row_str}")

    tables = []
    current_rows = [rows[0]]
    for row in rows[1:]:
        if not row.fields:
            continue
        y_distance = row.y_distance(current_rows[-1])
        if y_distance > Config.max_row_distance:
            if len(current_rows) < Config.min_row_count:
                log_skipped_rows()
                current_rows = [row]
                continue
            logger.debug(f"Distance between rows: {y_distance}")
            tables.append(PDFTable(current_rows))
            current_rows = []
        current_rows.append(row)
    else:
        if len(current_rows) < Config.min_row_count:
            log_skipped_rows()
            return tables
        tables.append(PDFTable(current_rows))
    return tables

def cleanup_tables(tables: Tables) -> Tables:
    """ Fix some errors in the tables. """
    for table in tables:
        for row in table.rows:
            row.update_type()

    tables = split_tables_with_multiple_header_rows(tables)
    for table in tables:
        table.generate_columns_from_rows()
    return split_tables_with_multiple_stop_columns(tables)

def split_tables_with_multiple_header_rows(tables: Tables) -> Tables:
    """ Merge table with the previous one, if it has no header row. """
    if not tables:
        return []

    merged_tables: Tables = [tables[0]]
    for table in tables[1:]:
        header_rows = table.rows.of_type(RowType.HEADER)
        if len(header_rows) > 1:
            merged_tables += table.split_at_header_rows()
            continue
        if header_rows:
            merged_tables.append(table)
            continue
        merged_tables[-1].rows.merge(table.rows)

    return merged_tables

def split_tables_with_multiple_stop_columns(tables: Tables) -> Tables:
    """ If a table has multiple stop columns, it will be split into two
    tables with each having only one stop column. """
    split_tables = []
    for table in tables:
        stop_columns = table.columns.of_type(ColumnType.STOP)
        if len(stop_columns) <= 1:
            split_tables.append(table)
            continue
        split_tables += table.split_at_stop_columns()
    return split_tables


import torch.nn as nn
import torch.nn.functional as F

from   ...utils import CONFIG
from   ..encoders.resnet_enc import ResNet_D
from   ..ops import GuidedCxtAtten, SpectralNorm

class ResGuidedCxtAtten(ResNet_D):

    def __init__(self, block, layers, norm_layer=None, late_downsample=False):
        super(ResGuidedCxtAtten, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
        first_inplane = 3 + CONFIG.model.trimap_channel
        self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
        self.shortcut_plane = [32, self.midplanes, 64, 128, 256]

        self.shortcut = nn.ModuleList()
        for stage, inplane in enumerate(self.shortcut_inplane):
            self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))

        self.guidance_head = nn.Sequential(
            nn.ReflectionPad2d(1),
            SpectralNorm(nn.Conv2d(3, 16, kernel_size=3, padding=0, stride=2, bias=False)),
            nn.ReLU(inplace=True),
            self._norm_layer(16),
            nn.ReflectionPad2d(1),
            SpectralNorm(nn.Conv2d(16, 32, kernel_size=3, padding=0, stride=2, bias=False)),
            nn.ReLU(inplace=True),
            self._norm_layer(32),
            nn.ReflectionPad2d(1),
            SpectralNorm(nn.Conv2d(32, 128, kernel_size=3, padding=0, stride=2, bias=False)),
            nn.ReLU(inplace=True),
            self._norm_layer(128)
        )

        self.gca = GuidedCxtAtten(128, 128)

        # initialize guidance head
        for layers in range(len(self.guidance_head)):
            m = self.guidance_head[layers]
            if isinstance(m, nn.Conv2d):
                if hasattr(m, "weight_bar"):
                    nn.init.xavier_uniform_(m.weight_bar)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _make_shortcut(self, inplane, planes):
        return nn.Sequential(
            SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
            nn.ReLU(inplace=True),
            self._norm_layer(planes),
            SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
            nn.ReLU(inplace=True),
            self._norm_layer(planes)
        )

    def forward(self, x):

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.activation(out)
        out = self.conv2(out)
        out = self.bn2(out)
        x1 = self.activation(out) # N x 32 x 256 x 256
        out = self.conv3(x1)
        out = self.bn3(out)
        out = self.activation(out)

        im_fea = self.guidance_head(x[:,:3,...]) # downsample origin image and extract features
        if CONFIG.model.trimap_channel == 3:
            unknown = F.interpolate(x[:,4:5,...], scale_factor=1/8, mode='nearest')
        else:
            unknown = F.interpolate(x[:,3:,...].eq(1.).float(), scale_factor=1/8, mode='nearest')

        x2 = self.layer1(out) # N x 64 x 128 x 128
        x3= self.layer2(x2) # N x 128 x 64 x 64
        x3, offset = self.gca(im_fea, x3, unknown) # contextual attention
        x4 = self.layer3(x3) # N x 256 x 32 x 32
        out = self.layer_bottleneck(x4) # N x 512 x 16 x 16

        fea1 = self.shortcut[0](x) # input image and trimap
        fea2 = self.shortcut[1](x1)
        fea3 = self.shortcut[2](x2)
        fea4 = self.shortcut[3](x3)
        fea5 = self.shortcut[4](x4)

        return out, {'shortcut':(fea1, fea2, fea3, fea4, fea5),
                     'image_fea':im_fea,
                     'unknown':unknown,
                     'offset_1':offset}

if __name__ == "__main__":
    from networks.encoders.resnet_enc import BasicBlock
    m = ResGuidedCxtAtten(BasicBlock, [3, 4, 4, 2])
    for m in m.modules():
        print(m)


#! /usr/bin/env python3

import json
import time

import rclpy
from rclpy.node import Node
from rclpy.executors import MultiThreadedExecutor
from robot_navigator import BasicNavigator, TaskResult
from geometry_msgs.msg import PoseStamped

from std_msgs.msg import Float64MultiArray, Int32, String, Bool, Int32MultiArray

class ProtMsg:
    def __init__(self) -> None:
        self.data = 0

# Holds the current pose of the robot
current_x = 0.0
current_y = 0.0
current_yaw_angle = 0.0

# Holds the pose of the current tag
current_tag_x = 0.0
current_tag_y = 0.0
current_tag_yaw_angle = 0.0

# Holds next machine on the queue
next_machine = None

# Holds the current state of the battery
this_battery_state = 0.0
prev_battery_state = 0.0

# Flag for detecting the change in the battery state
low_battery = False
low_battery_min_threshold = 0.15

# File with the current machines positions
_file = open("Data/machine_positions.json")
machines = json.load(_file)

#PoC Variables
sim_list = [0,675,674,673,672,671,670,669,668]
idx = 0
first_run = True

nav_start = 0
pick_start = 0

class PickingNavigator(Node):
    """
    Navigates and request docking routines
    """

    def __init__(self):
        super().__init__("navigate_to_machines")

        self.current_mc_tag_pub = self.create_publisher(Int32, "/current_tag", 10)

        self.pick_feedback_sub = self.create_publisher(
            Int32MultiArray, "/picked_machine", 10
        )

        self.next_mc_request_sub = self.create_publisher(
            Bool, "/machine_number_request", 10
        )

        self.start_docking_pub = self.create_publisher(Bool, "/start_docking", 10)

        self.next_machine = self.create_subscription(
            Int32, "/next_machine_number", self.set_next_machine, 10
        )

        self.picking_feedback = self.create_subscription(
            Int32MultiArray, "/pick_feedback", self.update_feedback, 10
        )

        self.request_machine_sim()

        # Execution timer
        timer_period = 0.02
        self.timer = self.create_timer(timer_period, self.navigate_to_station)

        # Declare distance metrics in meters
        self.distance_goal_tolerance = 0.01
        self.reached_distance_goal = False

        # Declare angle metrics in radians
        self.heading_tolerance = 0.10
        self.yaw_goal_tolerance = 0.05

        # Picking related variables
        self.picking_feedback = [0, 0]
        self.picking = False

        self.nav_start = 0
        self.pick_start = 0

        # Navigation setup and feedback
        self.nav = BasicNavigator()
        self.nav.waitUntilNav2Active()
        self.done_navigating = False

        self.curr_record = [0]*4

    def set_next_machine(self, msg):
        """
        Sets the next machine in queue
        """
        global next_machine
        for mc_name, mc in machines.items():
            try:
                if mc["Number"] == msg.data:
                    next_machine = mc
            except:
                print("Machine number not found in 'machine_positions.json'")

    def update_feedback(self, msg):
        """
        Gets picking feedback from the docking node
        """
        self.picking_feedback = msg.data

    def request_machine_sim(self):
        global idx
        mc = sim_list[idx]
        idx = idx+1 if idx <= len(sim_list)-2 else 0
        msg = ProtMsg()
        msg.data = mc
        print("NEXT MACHINE: ")
        print(mc)
        self.set_next_machine(msg)

    def navigate_to_station(self):
        """
        Main routine: Goes to the selected machine and requests a docking sequence.
        Will end with "SUCCESS" flag if the docking sequence returns a "Done" value.
        Will stop with "FAIL" flag if the machine is unreachable/unconnected or if the
        docking sequence fails.
        """
        global next_machine, first_run
        if self.done_navigating:
            result = self.nav.getResult()
            if result == TaskResult.SUCCEEDED:
                if not self.picking:
                    self.curr_record[0] = 1
                    self.curr_record[1] = time.time()-self.nav_start
                    start_picking_msg = Bool()
                    start_picking_msg.data = False if (next_machine!=None and next_machine["Number"]==0) else True
                    self.start_docking_pub.publish(start_picking_msg)
                    self.picking_feedback = [0, 0]
                    self.picking = True
                    self.pick_start = time.time()
                else:
                    if self.picking_feedback == [0, 0]:
                        print("Waiting for picking feedback")
                    else:
                        if self.picking_feedback[1] == 0: fb = "Failed"
                        elif self.picking_feedback[1] == 1: fb = "Picked"
                        elif self.picking_feedback[1] == 2: fb = "At home"
                        else: fb = "Unknown feedback code"
                        self.curr_record[2] = time.time()-self.pick_start
                        self.curr_record[3] = self.picking_feedback[1]
                        with open("records.txt", "a") as file_object:
                            file_object.write(f"\n{round(self.curr_record[0],2)},{round(self.curr_record[1],2)},{round(self.curr_record[2],2)},{round(self.curr_record[3],2)}")
                        print(fb)
                        self.picking = False
                        self.done_navigating = False
                        next_machine = None
                        self.request_machine_sim()
            elif result == TaskResult.CANCELED:
                print("Goal was canceled!")
                self.done_navigating = False
                next_machine = None
                self.request_machine_sim()
            elif result == TaskResult.FAILED:
                print("Goal failed!")
                self.done_navigating = False
                next_machine = None
                self.request_machine_sim()
            else:
                print("Goal has an invalid return status!")
                self.done_navigating = False
                next_machine = None
                self.request_machine_sim()
            return None

        msg_req = Bool()
        msg_req.data = True
        self.next_mc_request_sub.publish(msg_req)

        if next_machine == None:
            return None

        current_machine = next_machine
        tag_msg = Int32()
        tag_msg.data = int(current_machine["Tag"])
        self.current_mc_tag_pub.publish(tag_msg)

        self.get_logger().info("Navigate to station")

        # You may use the navigator to clear or obtain costmaps
        self.nav.clearAllCostmaps()  # also have clearLocalCostmap() and clearGlobalCostmap()
        # global_costmap = self.nav.getGlobalCostmap()
        # local_costmap = self.nav.getLocalCostmap()

        # Set the robot's goal pose
        goal_pose = PoseStamped()
        goal_pose.header.frame_id = "map"
        goal_pose.header.stamp = self.nav.get_clock().now().to_msg()

        goal_pose = PoseStamped()
        goal_pose.header.frame_id = "map"
        goal_pose.header.stamp = self.nav.get_clock().now().to_msg()
        goal_pose.pose.position.x = current_machine["Position"][0]
        goal_pose.pose.position.y = current_machine["Position"][1]
        goal_pose.pose.position.z = 0.0
        goal_pose.pose.orientation.x = 0.0
        goal_pose.pose.orientation.y = 0.0
        if current_machine["Position"][2] == "D":
            z_pos = -0.7071068
        elif current_machine["Position"][2] == "E":
            z_pos = 0.7071068
        else:
            z_pos = 0.0
        goal_pose.pose.orientation.z = z_pos  # current_machine["Position"][2]
        goal_pose.pose.orientation.w = 0.7071068

        # Go to the goal pose
        self.nav.goToPose(goal_pose)
        self.nav_start = time.time()

        i = 0

        # Keep doing stuff as long as the robot is moving towards the goal
        while not self.nav.isTaskComplete():
            # Do something with the feedback
            i = i + 1
            feedback = self.nav.getFeedback()
            if feedback and i % 5 == 0:
                print(
                    "Distance remaining: "
                    + "{:.2f}".format(feedback.distance_remaining)
                    + " meters."
                )

                # Some navigation timeout to demo cancellation
                # if Duration.from_msg(feedback.navigation_time) > Duration(seconds=1800.0):
                # self.nav.cancelNav()
        self.done_navigating = True

class BatteryStateSubscriber(Node):
    """
    Subscriber node to the current battery state
    """

    def __init__(self):

        # Initialize the class using the constructor
        super().__init__("battery_state_subscriber")

        # Create a subscriber
        # This node subscribes to messages of type
        # sensor_msgs/BatteryState
        self.subscription_battery_state = self.create_subscription(
            Int32, "/battery_soc", self.get_battery_state, 10
        )

    def get_battery_state(self, msg):
        """
        Update the current battery state.
        """
        global this_battery_state
        global prev_battery_state
        global low_battery
        prev_battery_state = this_battery_state
        this_battery_state = float(msg.data) / 100

        # Check for low battery
        if (
            prev_battery_state >= low_battery_min_threshold
            and this_battery_state < low_battery_min_threshold
        ):
            low_battery = True

class PoseSubscriber(Node):
    """
    Subscriber node to the current 2D pose of the robot
    """

    def __init__(self):

        # Initialize the class using the constructor
        super().__init__("pose_subscriber")

        # Create a subscriber
        # This node subscribes to messages of type
        # std_msgs/Float64MultiArray
        self.subscription_pose = self.create_subscription(
            Float64MultiArray, "/bot_2d_pose", self.get_pose, 1
        )

        self.tag_pose_sub = self.create_subscription(
            Float64MultiArray, "/tag_2d_pose", self.get_tag_pose, 1
        )

    def get_pose(self, msg):
        global current_x
        global current_y
        global current_yaw_angle
        current_2d_pose = msg.data
        current_x = current_2d_pose[0]
        current_y = current_2d_pose[1]
        current_yaw_angle = current_2d_pose[2]

    def get_tag_pose(self, msg):
        global current_tag_x
        global current_tag_y
        global current_tag_yaw_angle
        current_2d_pose = msg.data
        current_tag_x = current_2d_pose[0]
        current_tag_y = current_2d_pose[1]
        current_tag_yaw_angle = current_2d_pose[2]

def main(args=None):
    rclpy.init(args=args)

    try:
        picking_navigator = PickingNavigator()
        battery_state_subscriber = BatteryStateSubscriber()
        pose_subscriber = PoseSubscriber()

        executor = MultiThreadedExecutor(num_threads=4)
        executor.add_node(picking_navigator)
        executor.add_node(battery_state_subscriber)
        executor.add_node(pose_subscriber)

        try:
            executor.spin()
        finally:
            executor.shutdown()
            picking_navigator.destroy_node()
            battery_state_subscriber.destroy_node()
            pose_subscriber.destroy_node()

    finally:
        # Shutdown
        rclpy.shutdown()

if __name__ == "__main__":
    main()


#!/usr/bin/env python

"""
A thermometer display stimulus.

This module contains a class implementing a thermometer display stimulus.

"""
from __future__ import absolute_import, print_function, division
from builtins import *

__author__ = 'Florian Krause <florian@expyriment.org>'
__version__ = ''
__revision__ = ''
__date__ = ''

import pygame

from expyriment.stimuli import Rectangle
from expyriment.stimuli._visual import Visual

class ThermometerDisplay(Visual):
    """A class implementing a thermometer display."""

    def __init__(self, state, goal=None, size=(128,340), nr_segments=10,
                 gap=3, frame_line_width=20, active_colour=(150,150,150),
                 inactive_colour=(0,0,0), frame_colour=(100,100,100),
                 goal_colour=(0,255,0), gap_colour=(255,255,255),
                 position=None):
        """Initializing a thermometer display.

        Parameters:
        -----------
        state : int
            The state of the thermometer in percent.
        goal : int, optional
            The goal state indication in percent.
        size : (int, int), optional
            The size of the thermometer display (default=(128,340)).
        nr_segments : int, optional
            The number of segments to use (default=10).
        gap : int, optional
            The visual gap between the individual segments (default=3).
        frame_line_width : int, optional
            The line width of the frame around the thermometer display
            (default=20).
        active_colour : (int, int, int), optional
            The colour of the active segments (default=(150,150,150)).
        inactive_colour : (int, int, int), optional
            The colour of the inactive segments (default=(0,0,0)).
        frame_colour : (int, int, int), optional
            The colour of the frame around the thermometer display
            (default=(100,100,100)).
        goal_colour : (int, int, int), optional
            The colour of the goal indicator (default=(0,255,0)).
        gap_colour : (int, int, int), optional
            The gap colour of the thermometer stimulus
            (default=(255,255,255)).
        position : (int, int), optional
            The position of the thermometer display).
        """

        self._state = state
        self._goal = goal
        self._size = size
        self._nr_segments = nr_segments
        self._gap = gap
        self._frame_line_width = frame_line_width
        self._active_colour = active_colour
        self._inactive_colour = inactive_colour
        self._frame_colour = frame_colour
        self._goal_colour = goal_colour
        self._gap_colour = gap_colour
        self._position = position

        Visual.__init__(self, position)

    _getter_exception_message = "Cannot set {0} if surface exists!"

    @property
    def state(self):
        """Getter for state."""
        return self._state

    @state.setter
    def state(self, value):
        """Setter for state."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format("state"))
        else:
            self._state = value

    @property
    def goal(self):
        """Getter for goal."""
        return self._goal

    @goal.setter
    def goal(self, value):
        """Setter for goal."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format("goal"))
        else:
            self._goal = value

    @property
    def size(self):
        """Getter for size."""
        return self._size

    @size.setter
    def size(self, value):
        """Setter for size."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format("size"))
        else:
            self._size = value

    @property
    def nr_segments(self):
        """Getter for nr_segments."""
        return self._nr_segments

    @nr_segments.setter
    def nr_segments(self, value):
        """Setter for nr_segments."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format(
                    "nr_segments"))
        else:
            self._nr_segments = value

    @property
    def gap(self):
        """Getter for gap."""
        return self._gap

    @gap.setter
    def gap(self, value):
        """Setter for gap."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format("gap"))
        else:
            self._gap = value

    @property
    def frame_line_width(self):
        """Getter for frame_line_width."""
        return self._frame_line_width

    @frame_line_width.setter
    def frame_line_width(self, value):
        """Setter for frame_line_width."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format(
                    "frame_line_width"))
        else:
            self._frame_line_width = value

    @property
    def active_colour(self):
        """Getter for active_colour."""
        return self._active_colour

    @active_colour.setter
    def active_colour(self, value):
        """Setter for active_colour."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format(
                    "active_colour"))
        else:
            self._active_colour = value

    @property
    def inactive_colour(self):
        """Getter for inactive_colour."""
        return self._inactive_colour

    @inactive_colour.setter
    def inactive_colour(self, value):
        """Setter for inactive_colour."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format(
                    "inactive_colour"))
        else:
            self._inactive_colour = value

    @property
    def frame_colour(self):
        """Getter for frame_colour."""
        return self._frame_colour

    @frame_colour.setter
    def frame_colour(self, value):
        """Setter for frame_colour."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format(
                    "frame_colour"))
        else:
            self._frame_colour = value

    @property
    def goal_colour(self):
        """Getter for goal_colour."""
        return self._goal_colour

    @goal_colour.setter
    def goal_colour(self, value):
        """Setter for goal_colour."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format(
                    "goal_colour"))
        else:
            self._goal_colour = value

    @property
    def gap_colour(self):
        """Getter for gap_colour."""
        return self._gap_colour

    @gap_colour.setter
    def gap_colour(self, value):
        """Setter for gap_colour."""

        if self.has_surface:
            raise AttributeError(
                ThermometerDisplay._getter_exception_message.format(
                    "gap_colour"))
        else:
            self._gap_colour = value

    def _create_surface(self):
        """Create the surface of the stimulus."""

        surface = pygame.surface.Surface((self._size[0] +
                                          self._frame_line_width,
                                          self._size[1] +
                                          self._frame_line_width),
                                         pygame.SRCALPHA).convert_alpha()

        if self._gap_colour is not None:
            surface.fill(self._gap_colour)

        parts = []
        width = self._size[0] - self._frame_line_width - \
                2 * self._frame_line_width % 2
        height = self._size[1] - self._frame_line_width - \
                 2 * self._frame_line_width % 2 + 1
        s_height = int(height - (self._nr_segments + 1) *
                       self._gap) // self._nr_segments
        for x in range(self._nr_segments):
            if x < self._state / 100.0 * self._nr_segments:
                colour = self._active_colour
            else:
                colour = self._inactive_colour
            s = Rectangle((width - self._gap * 2,
                           s_height), colour=colour,
                          position=(0, -height // 2 + s_height // 2 + x *
                                    height // self.nr_segments + self._gap))
            parts.append(s)
        parts.append(Rectangle(self._size, colour=self._frame_colour,
                               line_width=self._frame_line_width,
                               position=self._position))
        parts.append(Rectangle((width - self._gap, height - self._gap * 2),
                               colour=self._gap_colour,
                               line_width=self._gap,
                               position=self._position))

        if self._goal is not None:
            x = int(round(self._goal / 100.0 * self._nr_segments)) - 1
            current_y_pos = -height // 2 + s_height // 2 + \
                            x * height // self._nr_segments + self._gap
            above_y_pos = -height // 2 + s_height // 2 + \
                          (x + 1) * height // self._nr_segments + self._gap
            g1 = Rectangle((self._frame_line_width * 1.25,
                            self._frame_line_width * 1.25),
                           colour=self._goal_colour,
                           position=(
                               -self._size[0] // 2 - self._frame_line_width // 2,
                               (current_y_pos + above_y_pos) // 2))
            g2 = Rectangle((self._frame_line_width * 1.25,
                            self._frame_line_width * 1.25),
                           colour=self._goal_colour,
                           position=(
                               self._size[0] // 2 + self._frame_line_width // 2,
                               (current_y_pos + above_y_pos) // 2))
            g1.rotate(45)
            g2.rotate(45)
            parts.append(g1)
            parts.append(g2)

        for stim in parts:
            stim.rect = pygame.Rect((0, 0), stim.surface_size)
            surface_size = surface.get_size()
            stim.rect.center = [stim.position[0] + surface_size[0] // 2,
                                - stim.position[1] + surface_size[1] // 2]
            surface.blit(stim._get_surface(), stim.rect)
        return surface

if __name__ == "__main__":
    from ... import control
    control.set_develop_mode(True)
    control.defaults.event_logging = 0
    exp = control.initialize()
    thermometer_display = ThermometerDisplay(50, 50)
    thermometer_display.present()
    #exp.clock.wait(1000)
    exp.keyboard.wait()


from __future__ import annotations

from typing import TYPE_CHECKING, Any

from ..network import ID_TOKEN_AUTHENTICATOR
from .by_plugin import AuthByPlugin, AuthType
from .webbrowser import AuthByWebBrowser

if TYPE_CHECKING:
    from ..connection import SnowflakeConnection

class AuthByIdToken(AuthByPlugin):
    """Internal IdToken Based Authentication.

    Works by accepting an id_toke and use that to authenticate. Only be used when users are using EXTERNAL_BROWSER_AUTHENTICATOR
    """

    @property
    def type_(self) -> AuthType:
        return AuthType.ID_TOKEN

    @property
    def assertion_content(self) -> str:
        return self._id_token

    def __init__(
        self,
        id_token: str,
        application: str,
        protocol: str | None,
        host: str | None,
        port: str | None,
        **kwargs,
    ) -> None:
        """Initialized an instance with an IdToken."""
        super().__init__(**kwargs)
        self._id_token: str | None = id_token
        self._application = application
        self._protocol = protocol
        self._host = host
        self._port = port

    def reset_secrets(self) -> None:
        self._id_token = None

    def prepare(self, **kwargs: Any) -> None:
        pass

    def reauthenticate(
        self,
        *,
        conn: SnowflakeConnection,
        **kwargs: Any,
    ) -> dict[str, bool]:
        conn.auth_class = AuthByWebBrowser(
            application=self._application,
            protocol=self._protocol,
            host=self._host,
            port=self._port,
            timeout=conn.login_timeout,
            backoff_generator=conn._backoff_generator,
        )
        conn._authenticate(conn.auth_class)
        conn._auth_class.reset_secrets()
        return {"success": True}

    def update_body(self, body: dict[Any, Any]) -> None:
        """Idtoken needs the authenticator and token attributes set."""
        body["data"]["AUTHENTICATOR"] = ID_TOKEN_AUTHENTICATOR
        body["data"]["TOKEN"] = self._id_token


from __future__ import annotations

from contextlib import contextmanager
from datetime import datetime
from decimal import Decimal
from typing import ContextManager, List, Text
from uuid import UUID, uuid4

import sqlalchemy as sa
import sqlalchemy_utils as sa_utils
from injector import inject
from sqlalchemy import func
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.orm import Session

from application.bus import Event, EventBus
from application.db import Base, Transaction
from currency import BTCRate, Currency

from .interface import BuyOrder, Repository
from .types import BtcAmountColumn, FiatAmountColumn

class DBBuyOrder(Base, BuyOrder):
    def __init__(
            self,
            request_id: UUID,
            paid: Decimal,
            bought: Decimal,
            exchange_rate: BTCRate,
    ) -> None:
        self.id = uuid4()
        self.request_id = request_id
        self.paid = paid
        self.bought = bought
        self.exchange_rate = exchange_rate

    __tablename__ = "buy_orders"

    _db_id: int = sa.Column("id", sa.Integer, primary_key=True)
    id: UUID = sa.Column("order_id", sa_utils.UUIDType, unique=True)
    request_id: UUID = sa.Column(sa_utils.UUIDType, unique=True)

    paid: Decimal = sa.Column(FiatAmountColumn)
    bought: Decimal = sa.Column(BtcAmountColumn)

    _currency: Currency = sa.Column(
        "currency", sa.Enum(Currency, naive_enum=False)
    )
    _price: Decimal = sa.Column("price", FiatAmountColumn)
    _rate_date: datetime = sa.Column("rate_date", sa.DateTime)

    @hybrid_property
    def exchange_rate(self) -> BTCRate:
        return BTCRate(
            price=self._price,
            currency=self._currency,
            on_date=self._rate_date,
        )

    @exchange_rate.setter
    def exchange_rate(self, value: BTCRate) -> None:
        self._price = value.price
        self._currency = value.currency
        self._rate_date = value.on_date

    when_created: datetime = sa.Column(
        "when_created", sa.DateTime, default=datetime.utcnow, index=True
    )
    _when_updated: datetime = sa.Column(
        "when_updated", sa.DateTime, onupdate=datetime.utcnow,
    )

    def __repr__(self) -> Text:
        return (
            "<"
            f"{self.__class__.__name__}"
            f" db_id={self._db_id}"
            f" order_id={self.id}"
            f" request_id={self.request_id}"
            f" amount={self.bought!s}BTC"
            ">"
        )

@inject
class ORMRepository(Repository):
    _session: Session
    _pending_events: List[Event]

    def __init__(self, transaction: Transaction, bus: EventBus) -> None:
        self._transaction = transaction
        self._bus = bus

    @contextmanager
    def lock(self) -> ContextManager[Decimal]:
        events = []
        with self._transaction() as session:
            self._session = session
            self._pending_events = events

            lock = f"LOCK TABLE {DBBuyOrder.__tablename__} IN EXCLUSIVE MODE;"
            session.execute(lock)

            yield self._calculate_balance()

            del self._session
            del self._pending_events

        for event in events:
            self._bus.emit(event)

    def _calculate_balance(self) -> Decimal:
        balance, = self._session.query(func.sum(DBBuyOrder.bought)).one()
        return Decimal(balance or 0)

    def create(
            self,
            request_id: UUID,
            paid: Decimal,
            bought: Decimal,
            with_rate: BTCRate,
    ) -> BuyOrder:
        entry = DBBuyOrder(
            request_id=request_id,
            paid=paid,
            bought=bought,
            exchange_rate=with_rate,
        )

        self._session.add(entry)
        self._session.flush()
        return entry

    def emit(self, event: Event) -> None:
        self._pending_events.append(event)

    def get_order_id(self, for_request_id: UUID) -> BuyOrder | None:
        with self._transaction() as session:
            result = (
                session.query(DBBuyOrder.id)
                .filter_by(request_id=for_request_id)
                .one_or_none()
            )
        return result and result[0]