import sys
import time

import grequests

with open('tests/pub.key') as f:
    PUBLIC_KEY = f.read()

def main():
    if len(sys.argv) != 3:
        print('Usage: python registration_load_test.py [start_id] [end_id]')
        sys.exit(1)

    start_id = int(sys.argv[1])
    end_id = int(sys.argv[2])

    with open('data/hospitals.txt') as f:
        hospitals = [line.strip() for line in f.readlines() if line.strip()]
        num_hospitals = len(hospitals)

    reqs = (
        grequests.post(f'http://localhost:{port}/patient_reg', json={
            'id': str(i),
            'name': f'Patient {i}',
            'pub_key': PUBLIC_KEY,
        }) for port in range(8100, 8100 + num_hospitals) for i in range(start_id, end_id)
    )

    start_time = time.time()
    responses = grequests.map(reqs)
    end_time = time.time()

    # print([res.json() for res in responses])

    # Ensure that the number of successful requests is correct
    num_success = len([1 for res in responses if res.json()['success']])
    if num_success != end_id - start_id:
        print(f'Expected {end_id-start_id} successes, got {num_success} successes')
    else:
        print(f'Correctly got {num_success} successes')

    print(f'Time taken: {end_time-start_time} for {end_id-start_id} requests each at {num_hospitals} hospitals')

if __name__ == "__main__":
    main()


import tensorflow as tf
for i in tf.config.experimental.list_physical_devices('GPU'):
    tf.config.experimental.set_memory_growth(i,True)
from tensorflow import keras
import keras.backend as K
from tensorflow.keras import layers
from tensorflow.python.profiler.model_analyzer import profile
from tensorflow.python.profiler.option_builder import ProfileOptionBuilder

def get_flops(model, input_sig=[tf.TensorSpec([1, 450, 8, 8, 3])]):
    forward_pass = tf.function(
        model.call,
        input_signature=input_sig)
    graph_info = profile(forward_pass.get_concrete_function().graph,
                            options=ProfileOptionBuilder.float_operation())
    return graph_info.total_float_ops

def to_tf(datatape, dtype=tf.float16):
    return tf.data.Dataset.from_generator(lambda :datatape, output_types=(dtype, dtype), output_shapes=(datatape.shape, datatape.shape[:1]))

"""
ours ↓
"""

class SpectralTransform(layers.Layer):

    def __init__(self, c, size=3):
        super().__init__()
        self.conv1 = keras.Sequential([
            layers.Conv1D(c*2, size, 1, padding='same'),
            layers.BatchNormalization(),
            layers.Activation('relu'),
        ])

    def call(self, y):
        x = tf.transpose(y, (0, 2, 1))
        x = tf.signal.rfft(x)
        r, i = tf.math.real(x), tf.math.imag(x)
        x = tf.concat([r, i], axis=-2)
        x = tf.transpose(x, (0, 2, 1))
        x = self.conv1(x)
        x = tf.transpose(x, (0, 2, 1))
        r, i = tf.split(x, 2, axis=-2)
        x = tf.complex(r, i)
        x = tf.signal.irfft(x)
        x = tf.transpose(x, (0, 2, 1))
        return layers.add([x,y])

class M_1(keras.Model):

    def __init__(self):
        super().__init__()
        self.a = keras.Sequential([
            layers.Reshape((1350, -1)),
            layers.Conv1D(64, 3, 3),
        ])
        self.ST1 = SpectralTransform(64, 5)
        self.conv1_1 = layers.Conv1D(64, 10, 1, padding='same')
        self.atv1 = keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')])
        self.ST2 = SpectralTransform(64, 3)
        self.conv2 = layers.Conv1D(32, 5, 1, padding='same')
        self.atv2 = keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')])
        self.z = keras.Sequential([
            layers.Conv1D(1, 1, 1),
            layers.Reshape((-1,)),
        ])
        self.w = {}

    def call(self, x):
        x = (x - tf.reshape(tf.reduce_mean(x, axis=(1, 2, 3)), (-1, 1, 1, 1, 3)))/tf.reshape(tf.math.reduce_std(x, axis=(1, 2, 3)), (-1, 1, 1, 1, 3))

        x = tf.transpose(x, (0, 1, 4, 2, 3))
        x = self.a(x)
        x = self.ST1(x)
        x = self.conv1_1(x)
        x = self.atv1(x)
        x = self.ST2(x)
        x = self.conv2(x)
        x = self.atv2(x)
        return self.z(x)

    def cross(self):
        if 'bn' not in self.w:
            self.w['bn'] = self.atv1, self.atv2
            self._ = keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')]), keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')])
            self._[0].build(input_shape=(None, None, 64))
            self._[1].build(input_shape=(None, None, 32))
            self._[0].set_weights(self.atv1.get_weights())
            self._[1].set_weights(self.atv2.get_weights())
            self.atv1, self.atv2 = lambda x:self._[0](x, training=True), lambda x:self._[1](x, training=True)

    def intra(self):
        if 'bn' in self.w:
            self.atv1, self.atv2 = self.w['bn']
            self.w.clear()

class M_2(keras.Model):

    def __init__(self):
        super().__init__()
        self.a = keras.Sequential([
            layers.Reshape((1350, -1)),
            layers.Conv1D(64, 3, 3),
        ])
        #self.ST1 = SpectralTransform(64, 5)
        self.conv1_1 = layers.Conv1D(64, 10, 1, padding='same')
        #self.conv1_2 = layers.Conv1D(32, 5, 1, padding='same')
        self.atv1 = keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')])
        #self.ST2 = SpectralTransform(64, 3)
        self.conv2 = layers.Conv1D(32, 5, 1, padding='same')
        self.atv2 = keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')])
        self.z = keras.Sequential([
            layers.Conv1D(1, 1, 1),
            layers.Reshape((-1,)),
        ])
        self.w = {}

    def call(self, x):
        x = x - tf.reshape(tf.reduce_mean(x, axis=(1, 2, 3)), (-1, 1, 1, 1, 3))
        x = tf.transpose(x, (0, 1, 4, 2, 3))
        x = self.a(x)
        #x = self.ST1(x)
        x = self.conv1_1(x)
        x = self.atv1(x)
        #x = self.ST2(x)
        x = self.conv2(x)
        x = self.atv2(x)
        return self.z(x)

    def cross(self):
        if 'atv' not in self.w:
            self.w['atv'] = self.atv1, self.atv2
            self._ = keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')]), keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')])
            self._[0].build(input_shape=(None, None, 64))
            self._[1].build(input_shape=(None, None, 32))
            self._[0].set_weights(self.atv1.get_weights())
            self._[1].set_weights(self.atv2.get_weights())
            self.atv1, self.atv2 = lambda x:self._[0](x, training=True), lambda x:self._[1](x, training=True)

    def intra(self):
        if 'atv' in self.w:
            self.atv1, self.atv2 = self.w['atv']
            self.w.clear()

class M_3(keras.Model):

    def __init__(self):
        super().__init__()
        self.a = keras.Sequential([
            layers.Reshape((1350, -1)),
            layers.Conv1D(64, 3, 3),
        ])
        self.ST1 = SpectralTransform(64, 5)
        #self.conv1_1 = layers.Conv1D(64, 10, 1, padding='same')
        #self.atv1 = keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')])
        self.ST2 = SpectralTransform(64, 3)
        #self.conv2 = layers.Conv1D(32, 5, 1, padding='same')
        #self.atv2 = keras.Sequential([layers.BatchNormalization(), layers.Activation('relu')])
        self.z = keras.Sequential([
            layers.Conv1D(1, 1, 1),
            layers.Reshape((-1,)),
        ])
        self.w = {}

    def call(self, x):
        x = x - tf.reshape(tf.reduce_mean(x, axis=(1, 2, 3)), (-1, 1, 1, 1, 3))

        x = tf.transpose(x, (0, 1, 4, 2, 3))
        x = self.a(x)
        x = self.ST1(x)
        #x = self.conv1_1(x)
        #x = self.atv1(x)
        x = self.ST2(x)
        #x = self.conv2(x)
        #x = self.atv2(x)
        return self.z(x)

"""
TS-CAN & DeepPhys & EfficientPhys ↓
"""

class Attention_mask(tf.keras.layers.Layer):
    def call(self, x):
        xsum = K.sum(x, axis=1, keepdims=True)
        xsum = K.sum(xsum, axis=2, keepdims=True)
        xshape = K.int_shape(x)
        return x / xsum * xshape[1] * xshape[2] * 0.5

    def get_config(self):
        config = super(Attention_mask, self).get_config()
        return config

class TSM(tf.keras.layers.Layer):
    def call(self, x, n_frame, fold_div=3):
        nt, h, w, c = x.shape
        x = K.reshape(x, (-1, n_frame, h, w, c))
        fold = c // fold_div
        last_fold = c - (fold_div - 1) * fold
        out1, out2, out3 = tf.split(x, [fold, fold, last_fold], axis=-1)

        # Shift left
        padding_1 = tf.zeros_like(out1)
        padding_1 = padding_1[:, -1, :, :, :]
        padding_1 = tf.expand_dims(padding_1, 1)
        _, out1 = tf.split(out1, [1, n_frame - 1], axis=1)
        out1 = tf.concat([out1, padding_1], axis=1)

        # Shift right
        padding_2 = tf.zeros_like(out2)
        padding_2 = padding_2[:, 0, :, :, :]
        padding_2 = tf.expand_dims(padding_2, 1)
        out2, _ = tf.split(out2, [n_frame - 1, 1], axis=1)
        out2 = tf.concat([padding_2, out2], axis=1)

        out = tf.concat([out1, out2, out3], axis=-1)
        out = K.reshape(out, (-1, h, w, c))

        return out

    def get_config(self):
        config = super(TSM, self).get_config()
        return config

def TSM_Cov2D(x, n_frame, nb_filters=128, kernel_size=(3, 3), activation='tanh', padding='same'):
    x = TSM()(x, n_frame)
    x = layers.Conv2D(nb_filters, kernel_size, padding=padding, activation=activation)(x)
    return x

def TS_CAN(n_frame, nb_filters1, nb_filters2, input_shape, kernel_size=(3, 3), dropout_rate1=0.25, dropout_rate2=0.5,
           pool_size=(2, 2), nb_dense=128):
    diff_input = layers.Input(shape=input_shape, name='input_1')
    rawf_input = layers.Input(shape=input_shape, name='input_2')

    d1 = TSM_Cov2D(diff_input, n_frame, nb_filters1, kernel_size, padding='same', activation='tanh')
    d2 = TSM_Cov2D(d1, n_frame, nb_filters1, kernel_size, padding='valid', activation='tanh')

    r1 = layers.Conv2D(nb_filters1, kernel_size, padding='same', activation='tanh')(rawf_input)
    r2 = layers.Conv2D(nb_filters1, kernel_size, activation='tanh')(r1)

    g1 = layers.Conv2D(1, (1, 1), padding='same', activation='sigmoid')(r2)
    g1 = Attention_mask()(g1)
    gated1 = layers.multiply([d2, g1])

    d3 = layers.AveragePooling2D(pool_size)(gated1)
    d4 = layers.Dropout(dropout_rate1)(d3)

    r3 = layers.AveragePooling2D(pool_size)(r2)
    r4 = layers.Dropout(dropout_rate1)(r3)

    d5 = TSM_Cov2D(d4, n_frame, nb_filters2, kernel_size, padding='same', activation='tanh')
    d6 = TSM_Cov2D(d5, n_frame, nb_filters2, kernel_size, padding='valid', activation='tanh')

    r5 = layers.Conv2D(nb_filters2, kernel_size, padding='same', activation='tanh')(r4)
    r6 = layers.Conv2D(nb_filters2, kernel_size, activation='tanh')(r5)

    g2 = layers.Conv2D(1, (1, 1), padding='same', activation='sigmoid')(r6)
    g2 = Attention_mask()(g2)
    gated2 = layers.multiply([d6, g2])

    d7 = layers.AveragePooling2D(pool_size)(gated2)
    d8 = layers.Dropout(dropout_rate1)(d7)

    d9 = layers.Flatten()(d8)
    d10 = layers.Dense(nb_dense, activation='tanh')(d9)
    d11 = layers.Dropout(dropout_rate2)(d10)
    out = layers.Dense(1)(d11)
    model = keras.models.Model(inputs=[diff_input, rawf_input], outputs=out)
    return model

def CAN(nb_filters1, nb_filters2, input_shape, kernel_size=(3, 3), dropout_rate1=0.25, dropout_rate2=0.5,
        pool_size=(2, 2), nb_dense=128):
    diff_input = layers.Input(shape=input_shape)
    rawf_input = layers.Input(shape=input_shape)

    d1 = layers.Conv2D(nb_filters1, kernel_size, padding='same', activation='tanh')(diff_input)
    d2 = layers.Conv2D(nb_filters1, kernel_size, activation='tanh')(d1)

    r1 = layers.Conv2D(nb_filters1, kernel_size, padding='same', activation='tanh')(rawf_input)
    r2 = layers.Conv2D(nb_filters1, kernel_size, activation='tanh')(r1)

    g1 = layers.Conv2D(1, (1, 1), padding='same', activation='sigmoid')(r2)
    g1 = Attention_mask()(g1)
    gated1 = layers.multiply([d2, g1])

    d3 = layers.AveragePooling2D(pool_size)(gated1)
    d4 = layers.Dropout(dropout_rate1)(d3)

    r3 = layers.AveragePooling2D(pool_size)(r2)
    r4 = layers.Dropout(dropout_rate1)(r3)

    d5 = layers.Conv2D(nb_filters2, kernel_size, padding='same', activation='tanh')(d4)
    d6 = layers.Conv2D(nb_filters2, kernel_size, activation='tanh')(d5)

    r5 = layers.Conv2D(nb_filters2, kernel_size, padding='same', activation='tanh')(r4)
    r6 = layers.Conv2D(nb_filters2, kernel_size, activation='tanh')(r5)

    g2 = layers.Conv2D(1, (1, 1), padding='same', activation='sigmoid')(r6)
    g2 = Attention_mask()(g2)
    gated2 = layers.multiply([d6, g2])

    d7 = layers.AveragePooling2D(pool_size)(gated2)
    d8 = layers.Dropout(dropout_rate1)(d7)

    d9 = layers.Flatten()(d8)
    d10 = layers.Dense(nb_dense, activation='tanh')(d9)
    d11 = layers.Dropout(dropout_rate2)(d10)
    out = layers.Dense(1)(d11)
    model = keras.models.Model(inputs=[diff_input, rawf_input], outputs=out)
    return model

class DeepPhys_end_to_end(keras.Model):
    def __init__(self, size=(36, 36)):
        super().__init__()
        self.size = size
        self.deepphys = CAN(32, 64, (*size, 3))

    def call(self, x):
        x_ = x[1:] - x[:-1]
        x_ = (x_ - tf.reshape(tf.reduce_mean(x_, axis=(1,2 )), (-1, 1, 1, 3)))/tf.reshape(tf.math.reduce_std(x_, axis=(1, 2))+1e-6, (-1, 1, 1, 3))
        return self.deepphys((tf.concat([x_, tf.zeros([1, *self.size, 3])], axis=0), tf.expand_dims(tf.reduce_mean(x, axis=(0, )), axis=0)))

class TS_CAN_end_to_end(keras.Model):
    def __init__(self, n=16, size=(36, 36)):
        super().__init__()
        self.size = size
        self.ts_can = TS_CAN(n, 32, 64, (*size, 3))

    def call(self, x):
        x_ = x[1:] - x[:-1]
        x_ = (x_ - tf.reshape(tf.reduce_mean(x_, axis=(1,2 )), (-1, 1, 1, 3)))/tf.reshape(tf.math.reduce_std(x_, axis=(1, 2))+1e-6, (-1, 1, 1, 3))
        return self.ts_can((tf.concat([x_, tf.zeros([1, *self.size, 3])], axis=0), tf.expand_dims(tf.reduce_mean(x, axis=(0, )), axis=0)))

class EP(keras.Model):
    def __init__(self, n=32):
        super().__init__()
        self.n_frame = n
        self.TSM = TSM()
        self.mc1 = layers.Conv2D(32, kernel_size=3, padding='same', activation='tanh')
        self.mc2 = layers.Conv2D(32, kernel_size=3, activation='tanh')
        self.mc3 = layers.Conv2D(64, kernel_size=3, padding='same', activation='tanh')
        self.mc4 = layers.Conv2D(64, kernel_size=3, activation='tanh')
        self.attc1 = layers.Conv2D(1, kernel_size=1, activation='sigmoid')
        self.msk = Attention_mask()
        self.attc2 = layers.Conv2D(1, kernel_size=1, activation='sigmoid')
        self.avgp = layers.AvgPool2D(2)
        self.dp1 = layers.Dropout(0.25)
        self.dp2 = layers.Dropout(0.5)
        self.dense1 = layers.Dense(128 ,activation='tanh')
        self.dense2 = layers.Dense(1)
        self.bn = layers.BatchNormalization()
        self.ft = layers.Flatten()

    def call(self, x):
        x = x[1:]-x[:-1]
        x = tf.concat([x, tf.zeros([1, 72, 72, 3])], axis=0)
        x = self.bn(x)
        x = self.TSM(x, self.n_frame)
        x = self.mc1(x)
        x = self.TSM(x, self.n_frame)
        x = self.mc2(x)
        msk = self.msk(self.attc1(x))
        x = layers.multiply([x, msk])
        x = self.avgp(x)
        x = self.dp1(x)
        x = self.TSM(x, self.n_frame)
        x = self.mc3(x)
        x = self.TSM(x, self.n_frame)
        x = self.mc4(x)
        msk = self.msk(self.attc2(x))
        x = layers.multiply([x, msk])
        x = self.avgp(x)
        x = self.dp1(x)
        #x = tf.reshape(x, (-1, tf.reduce_prod(x.get_shape()[1:])))
        x = self.ft(x)
        x = self.dense1(x)
        x = self.dp2(x)
        x = self.dense2(x)
        return x

"""
PhysNet ↓
"""

class PhysNet(keras.Model):

    def __init__(self, norm='batch'):
        self.norm = norm
        if norm == 'batch':
            norm = layers.BatchNormalization
        if norm == 'layer':
            norm = lambda :layers.LayerNormalization(axis=(1,))
        if norm == 'layer_frozen':
            norm = lambda :layers.LayerNormalization(axis=(1,), trainable=False)
        super().__init__()
        self.ConvBlock1 = keras.Sequential([
            layers.Conv3D(16, kernel_size=(1, 5, 5), strides=1, padding='same'),
            norm(),
            layers.Activation('relu')
        ])
        self.ConvBlock2 = keras.Sequential([
            layers.Conv3D(32, kernel_size=(3, 3, 3), strides=1, padding='same'),
            norm(),
            layers.Activation('relu')
        ])
        self.ConvBlock3 = keras.Sequential([
            layers.Conv3D(64, kernel_size=(3, 3, 3), strides=1, padding='same'),
            norm(),
            layers.Activation('relu')
        ])
        self.ConvBlock4 = keras.Sequential([
            layers.Conv3D(64, kernel_size=(3, 3, 3), strides=1, padding='same'),
            norm(),
            layers.Activation('relu')
        ])
        self.ConvBlock5 = keras.Sequential([
            layers.Conv3D(64, kernel_size=(3, 3, 3), strides=1, padding='same'),
            norm(),
            layers.Activation('relu')
        ])
        self.ConvBlock6 = keras.Sequential([
            layers.Conv3D(64, kernel_size=(3, 3, 3), strides=1, padding='same'),
            norm(),
            layers.Activation('relu')
        ])
        self.ConvBlock7 = keras.Sequential([
            layers.Conv3D(64, kernel_size=(3, 3, 3), strides=1, padding='same'),
            norm(),
            layers.Activation('relu')
        ])
        self.ConvBlock8 = keras.Sequential([
            layers.Conv3D(64, kernel_size=(3, 3, 3), strides=1, padding='same'),
            norm(),
            layers.Activation('relu')
        ])
        self.ConvBlock9 = keras.Sequential([
            layers.Conv3D(64, kernel_size=(3, 3, 3), strides=1, padding='same'),
            norm(),
            layers.Activation('relu')
        ])
        self.upsample = keras.Sequential([
            layers.Conv3DTranspose(64, kernel_size=(4, 1, 1), strides=(2, 1, 1), padding='same'),
            norm(),
            layers.Activation('elu')
        ])
        self.upsample2 = keras.Sequential([
            layers.Conv3DTranspose(64, kernel_size=(4, 1, 1), strides=(2, 1, 1), padding='same'),
            norm(),
            layers.Activation('elu')
        ])
        self.convBlock10 = layers.Conv3D(1, kernel_size=(1, 1, 1), strides=1)
        self.MaxpoolSpa = layers.MaxPool3D((1, 2, 2), strides=(1, 2, 2))
        self.MaxpoolSpaTem = layers.MaxPool3D((2, 2, 2), strides=2)
        self.poolspa = layers.AvgPool3D((1, 2, 2))
        self.flatten = layers.Reshape((-1,))

    def call(self, x):
        if self.norm == 'batch':
            training=True
        else:
            training=False
        x = self.ConvBlock1(x, training=training)
        x = self.MaxpoolSpa(x)
        x = self.ConvBlock2(x, training=training)
        x = self.ConvBlock3(x, training=training)
        x = self.MaxpoolSpaTem(x)
        x = self.ConvBlock4(x, training=training)
        x = self.ConvBlock5(x, training=training)
        x = self.MaxpoolSpaTem(x)
        x = self.ConvBlock6(x, training=training)
        x = self.ConvBlock7(x, training=training)
        x = self.MaxpoolSpa(x)
        x = self.ConvBlock8(x, training=training)
        x = self.ConvBlock9(x, training=training)
        x = self.upsample(x, training=training)
        x = self.upsample2(x, training=training)
        x = self.poolspa(x)
        x = self.convBlock10(x, training=training)
        x = self.flatten(x)
        x = x-tf.expand_dims(tf.reduce_mean(x, axis=-1), -1)
        return x

"""
ResNet ↓

resnet18 = ResNet([2, 2, 2, 2])
"""
class BasicBlock(layers.Layer):

    def __init__(self, filter_num, stride=1):
        super().__init__()
        self.conv1 = layers.Conv2D(filter_num, 3, strides=stride, padding='same')
        self.bn1 = layers.BatchNormalization()
        self.relu = layers.ReLU()
        self.conv2 = layers.Conv2D(filter_num, 3, strides=1, padding='same')
        self.bn2 = layers.BatchNormalization()
        if stride != 1:
            self.downsample = keras.Sequential()
            self.downsample.add(layers.Conv2D(filter_num, (1, 1), strides=stride))
        else:
            self.downsample = lambda x:x

    def call(self, inputs, training=False):
        out = self.conv1(inputs, training=training)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out, training=training)
        out = self.bn2(out)
        identity = self.downsample(inputs)
        output = layers.add([out, identity])
        output = self.relu(output)
        return output

class ResNet(keras.Model):
    def __init__(self, layer_dims, out=1):
        super(ResNet, self).__init__()
        # 第一层
        self.stem = keras.Sequential([
                                  layers.Conv2D(64, 7, strides=(1, 1)),
                                  layers.BatchNormalization(),
                                  layers.Activation('relu'),
                                  layers.AvgPool2D(pool_size=(2, 2), strides=(2, 2), padding='same')
                                  ])
        # 中间层的四个残差块：conv2_x，conv3_x，conv4_x，conv5_x
        self.layer1 = self.build_resblock(64,  layer_dims[0])
        self.layer2 = self.build_resblock(128, layer_dims[1], stride=2)
        self.layer3 = self.build_resblock(256, layer_dims[2], stride=2)
        self.layer4 = self.build_resblock(512, layer_dims[3], stride=2)
        # 池化
        self.pool = layers.AvgPool2D(7, padding='same')

        # 全连接层
        self.fc = keras.Sequential([layers.Flatten(), layers.Dense(out)])

    def call(self, inputs, training=None, **kwargs):
        x = inputs
        x = self.stem(x,training=training)
        x = self.layer1(x,training=training)
        x = self.layer2(x,training=training)
        x = self.layer3(x,training=training)
        x = self.layer4(x,training=training)
        x = self.pool(x,training=training)
        x1 = self.fc(x,training=training)
        return x1

    # 构建残差块（将几个相同的残差模块堆叠在一起）
    def build_resblock(self, filter_num, blocks, stride=1):
        res_blocks = keras.Sequential()
        res_blocks.add(BasicBlock(filter_num, stride))
        for _ in range(1, blocks):
            res_blocks.add(BasicBlock(filter_num, stride=1))
        return res_blocks

"""
loss ↓
"""

"""
在TS-CAN和DeepPhys中使用以下loss都需要将输出和标签转置
def np_loss_tscan(x, y):
    return np_loss(tf.transpose(x), tf.transpose(y))
"""
def np_loss(x, y):
    x_, y_ = tf.expand_dims(tf.reduce_mean(x, axis=-1), -1), tf.expand_dims(tf.reduce_mean(y, axis=-1), -1)
    return 1-tf.reduce_sum((x-x_)*(y-y_), axis=-1)/(tf.reduce_sum((x-x_)**2, axis=-1)*tf.reduce_sum((y-y_)**2, axis=-1))**0.5

def SNR(x, y):
    x, y = (x-tf.expand_dims(tf.reduce_mean(x, axis=-1), -1))/(tf.expand_dims(tf.math.reduce_std(x, axis=-1), -1)+1e-6), (y-tf.expand_dims(tf.reduce_mean(y, axis=-1), -1))/(tf.expand_dims(tf.math.reduce_std(y, axis=-1), -1)+1e-6)
    A_s = tf.reduce_mean(tf.abs(tf.signal.rfft(x)), axis=-1)
    A_n = tf.reduce_mean(tf.abs(tf.signal.rfft(y-x)), axis=-1)
    #A_n = tf.reduce_mean(tf.abs(tf.signal.rfft(y)), axis=-1) - A_s
    return 8.685889*tf.math.log((A_s/A_n)**2)
    #return (A_s/A_n)**2

def nSNR_loss(x, y):
    return -SNR(x, y)


import sys
import time
import numpy as np

def read_doc(filepath):

    with open(filepath, 'rt', encoding='utf-8') as fp:
        doc = [line for line in fp.read().splitlines() if line.strip() != '']

    return doc

def load_stopwords(filepath):
    """Loads stopwords from text file."""

    with open(filepath, 'rt') as fp:
        stopwords = set(fp.read().split())

    return stopwords

def get_sentences(doc):
    assert isinstance(doc, list)

    from re import split

    sentences = []
    for line in doc:
        sent = [s.strip() for s in split(r"[\?!.]", line)]
        sentences += (([s for s in sent if s != '']))

    return sentences

def clean(s):
    from re import finditer
    pattern = r"[a-z]+('[a-z])?[a-z]*"

    return [match.group(0) for match in finditer(pattern, s.lower())]

# Bag of Words
def gen_bag_of_words(sents):
    assert isinstance(sents, list)

    # load stopwords
    STOPWORDS = load_stopwords('data/english_stopwords.txt')

    ls = [None] * len(sents)
    for i, s in enumerate(sents):
        ls[i] = set([w for w in (set(clean(s))) if w not in STOPWORDS])

    return ls

def create_ids(bags):
    """
    Given a bag of words, creates a dictionary word index and a
    reverse lookup from ID to word.

    Returns:
        Tuple of two dictionaries
        `word_to_id` and `id_to_word`

    """
    # Generate Word IDs
    from random import choices

    all_words = set()
    for b in bags:
        all_words |= b

    word_to_id = {w: k for k, w in enumerate(all_words)}
    id_to_word = {k: w for k, w in enumerate(all_words)}

    # check ID creation
    for w in choices(list(all_words), k=5):
        assert id_to_word[word_to_id[w]] == w

    return word_to_id, id_to_word

# Sparse Matrix Coordinates and Values
def gen_coords(bags, word_to_id):
    # Some code to help you get started:
    m, n = len(word_to_id), len(bags)
    rows, cols, vals = [], [], []

    # Construct rows, cols, and vals:
    # rows = word ID
    # cols = sentence ID

    import numpy as np

    # essentially, this is the total count of each unique word in the entire document
    # num of sentence containing `i` word
    # iterate over each sentence/bag/document, then over each word
    ni = np.zeros(m)
    for sent in bags:
        for i in sent:
            # each word ID corresponds to the index of `ni`
            ni[word_to_id[i]] += 1

    for idx, sent in enumerate(bags):
        for w in sent:
            word_id = word_to_id[w]
            # word ID from dict
            rows.append(word_id)
            # sentence index
            cols.append(idx)
            # values
            a_ij = 1 / np.log((n + 1) / ni[word_id])
            vals.append(a_ij)

    # Returns your arrays:
    return rows, cols, vals

def get_svd_largest(A):
    from scipy.sparse.linalg import svds
    from numpy import abs
    u, s, v = svds(A, k=1, which='LM', return_singular_vectors=True)
    return s, abs(u.reshape(A.shape[0])), abs(v.reshape(A.shape[1]))

def rank_words(u0, v0):
    """
    Computes rank and returns the IDs of words in descending order
    of importance.
    """

    # NOTE:
    #    reshape Sigma and U to be a row vector (610,1)
    #    reshape V to be a column vector (1,110)

    # multiply Sigma and U, then DOT PROD with V.transpose
    rank = (sigma0 * u0).reshape(-1, 1).dot(v0.reshape(-1, 1).T)

    # sort rankings (largest to smallest)
    # `axis=0` sorts row-wise (rows are words)
    idx_rank = np.argsort(rank, axis=0)[::-1]

    # return ranked words
    return idx_rank[:, 0]

def rank_sentences(u0, v0):
    rank = (sigma0 * u0).reshape(-1, 1).dot(v0.reshape(-1, 1).T)

    # sort rankings (largest to smallest)
    # `axis=1` sorts column-wise (columns are sentence IDs)
    idx_rank = np.argsort(rank, axis=1)

    # return ranked words
    return idx_rank[0, :][::-1]

def get_matrix(bags, word_to_id):
    # create sparse matrix of words and sentences

    from scipy.sparse import csr_matrix

    rows, cols, vals = gen_coords(bags, word_to_id)

    return csr_matrix((vals, (rows, cols)), shape=(len(word_to_id), len(bags)))

def get_topn_words(u0, v0, n=10):
    # compute work rank
    word_ranking = rank_words(u0, v0)

    return [id_to_word[k] for k in word_ranking[:n]]

def get_topn_sentences(u0, v0, n=5):
    # compute sentence rank
    sentence_ranking = rank_sentences(u0, v0)
    raw_sents = get_sentences(raw_doc)
    return [raw_sents[k] for k in sentence_ranking[:n]]

if __name__ == '__main__':

    start_time = time.time()

    # provide filepath
    filepath = sys.argv[1]

    # read file
    raw_doc = read_doc(filepath)

    print('Document type:', type(raw_doc))
    print('Document length:', len(raw_doc))

    # get bag of words
    bags = gen_bag_of_words(get_sentences(raw_doc))

    # get word ids
    word_to_id, id_to_word = create_ids(bags)

    # compute SVD
    A = get_matrix(bags, word_to_id)
    sigma0, u0, v0 = get_svd_largest(A)

    print('-'*100)
    # compute work rank
    top_ten_words = get_topn_words(u0, v0, n=10)
    print("\nTop 10 words:\n", top_ten_words)
    print()

    # compute sentence rank
    top_five_sentences = get_topn_sentences(u0, v0, n=5)

    print("Top 5 sentences")
    print('-' * 100)
    for k, s in enumerate(top_five_sentences):
        print(f"\n{k}:", s)

    print('\n\nExecution time:', round(time.time() - start_time, 4), 'sec')
    print('-'*100)


import unittest
import subprocess
from time import sleep
import utils.file_utils as fu
import crawler.ffmod as ffm
import sttest
SYS_CALLS = ["open", "read", "write"]

class Test(sttest.STTest):
        def test_simple_trace(self):
            sys_log_file = open(self.vi.sys_log, "w")
            subprocess.Popen("strace ls",
                             stderr=sys_log_file, stdout=sys_log_file,
                             shell=True)
            sys_log_file.close()
            sleep(5)  # wait for strace output to be available
            sys_logs = fu.read_file(self.vi.sys_log)
            for sys_call in SYS_CALLS:
                self.assertIn(sys_call, sys_logs)

        def test_attach_trace(self):
            sys_log_file = open(self.vi.sys_log, "w")
            vdisplay = ffm.start_xvfb()
            driver, _, proc = ffm.get_browser(self.vi.ff_log)
            subprocess.Popen("strace -p %d" % proc.pid,
                             stderr=sys_log_file, stdout=sys_log_file,
                             shell=True)
            driver.get("https://google.com")
            sleep(1)
            driver.close()
            proc.terminate()
            driver.quit()
            ffm.stop_xvfb(vdisplay)
            sys_log_file.close()
            sleep(5)  # wait for strace output to be available
            sys_logs = fu.read_file(self.vi.sys_log)
            for sys_call in SYS_CALLS:
                self.assertIn(sys_call, sys_logs,
                              """Cannot trace process, make sure you set up
                    ptrace permissions in /proc/sys/kernel/yama/ptrace_scope as
                    described in etc/setup.sh.""")

if __name__ == "__main__":
    unittest.main()


import math
import torch
import numpy as np
import torch.nn as nn
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module
from torch_geometric.nn import MessagePassing

class Graphsn_GCN_MM(Module):
    def __init__(self, in_features, out_features, bias=True):
        super(Graphsn_GCN_MM, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
        self.eps = nn.Parameter(torch.FloatTensor(1))
        if bias:
            self.bias = Parameter(torch.FloatTensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 0.9 / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)

        stdv_eps = 0.21 / math.sqrt(self.eps.size(0))
        nn.init.constant_(self.eps, stdv_eps)

        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, input, sc):
        v = (self.eps) * torch.diag(sc)
        mask = torch.diag(torch.ones_like(v))
        # print(f"v = {v}")
        # print(f"\tmask = {mask}")
        sc = mask * torch.diag(v) + (1. - mask) * sc
        support = torch.mm(input, self.weight)
        output = torch.spmm(sc, support)
        if self.bias is not None:
            return output + self.bias
        else:
            return output

    def __repr__(self):
        return self.__class__.__name__ + ' (' \
               + str(self.in_features) + ' -> ' \
               + str(self.out_features) + ')'

class Graphsn_GCN_MP(MessagePassing):
    def __init__(self, in_features, out_features, bias=True):
        super(Graphsn_GCN_MP, self).__init__(aggr='add')
        self.in_features = in_features
        self.out_features = out_features

        self.weight = nn.Parameter(torch.FloatTensor(in_features, out_features))
        self.eps = nn.Parameter(torch.FloatTensor(1))

        if bias:
            self.bias = nn.Parameter(torch.FloatTensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 0.9 / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)

        stdv_eps = 0.21 / math.sqrt(self.eps.size(0))
        nn.init.constant_(self.eps, stdv_eps)

        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, x, edge_index, edge_weight):
        # print(f"\tedge_index.shape = {edge_index.shape}")
        v = (self.eps) * torch.diag(edge_weight)
        mask = torch.diag(torch.ones_like(v))
        edge_weight = mask * torch.diag(v) + (1. - mask) * edge_weight

        support = torch.matmul(x, self.weight)
        output = self.propagate(edge_index, x=support, edge_weight=edge_weight)

        if self.bias is not None:
            output = output + self.bias

        return output

    def message(self, x_j, edge_weight):
        return edge_weight.view(-1, 1) * x_j

    # def update(self, aggr_out):
    #     print(f"\taggr_out.shape = {aggr_out.shape}")
    #     return aggr_out

    def __repr__(self):
        return self.__class__.__name__ + ' (' \
               + str(self.in_features) + ' -> ' \
               + str(self.out_features) + ')'


"""
Utilities to help with regging. This shouldn't completely replace filling out the form, since a small effort
barrier should hopefully keep people a little more honest.
At the same time, people shouldn't need to keep calling !user_info
"""

from discord.ext import commands
from .utils import checks, messages
import discord
import json
import glob
import re
import logging

# Note: Might need to pip install PyOpenSSL
import gspread
import gspread.utils
from oauth2client.service_account import ServiceAccountCredentials
import dateutil.tz
import dateutil.parser

log = logging.getLogger()

report_output = """
User: {0.name}
Joined at: {0.joined_at}
Acct created at: {0.created_at}
Application filed: {1}
"""

# Sheet is "https://docs.google.com/spreadsheets/d/1w0TOLj-SWXMFf0Hf8Zv7oJxX5oiDsGoEOfJFyt1vtkk/edit#gid=0"

class RegularUtils:

    def __init__(self, bot):
        self.bot = bot
        self.config = bot.config
        self.gclient = None
        self.wks = None
        self._google_auth("1w0TOLj-SWXMFf0Hf8Zv7oJxX5oiDsGoEOfJFyt1vtkk")  # TODO: Update w/ actual form

    @property
    def guild(self):
        return self.bot.get_guild(111504456838819840)

    def _google_auth(self, sheet_key):
        """Authenticate with google"""
        scope = ['https://spreadsheets.google.com/feeds']
        with open("auth.json", "r", encoding="utf-8") as datafile:
            auth = json.load(datafile)["sheets"]
        self.creds = ServiceAccountCredentials.from_json_keyfile_dict(auth, scope)
        self.gclient = gspread.authorize(self.creds)
        # Open the spreadsheet.
        # Even if link sharing is enabled, the spreadsheet has to be explicitly shared with client.email in auth.json
        self.wks = self.gclient.open_by_key(sheet_key).sheet1

    @staticmethod
    def _timestamp_to_datetime(timestamp):
        """Stuff we have to do because strptime expects padding"""
        return dateutil.parser.parse(timestamp, tzinfo=dateutil.tz.gettz("America/New York"))

    @staticmethod
    def find_member_by_username(username, guild):
        """Process the username and discriminator and check to see if we can find a match"""
        return discord.utils.find(lambda m: str(m) == username, guild.members)

    async def process_reg_document(self):
        """
        Process the regular form document and check to see if there are any changes. If so, add them to the tracked
        config file.

        :return: List of member objects that have been added since the document was last checked.
        """

        # Store the last update position in the google doc itself
        last_checked_row = self.config.get("config:regulars:last_checked_row")
        last_checked_row = int(last_checked_row) if last_checked_row else 2
        last_valid_row = last_checked_row  # Keeping track of the last row that we know has a good value

        processed_members = []

        if last_checked_row == self.wks.row_count:
            return []

        # Starting at the last row we checked last time, scan through until we hit a blank cell.
        # If there is a new appeal, post it, and update the value for the last valid row so we don't scan through
        # the document multiple times.
        for cell in self.wks.range("A{}:A{}".format(last_checked_row + 1, self.wks.row_count)):  # Skip the first row
            if cell.value == "":  # End of list
                last_valid_row = cell.row - 1
                break

            # Timestamp is in est

            active_row = self.wks.row_values(cell.row)
            ts = active_row[0]

            parts = {
                "application_id": cell.row - 1,
                "timestamp_from_form": ts,
                "discord_username": active_row[1],
                "reddit_username": active_row[2],
                "joined_date": active_row[3],
                "favorite_pkmn": active_row[4],
                "birthday": active_row[5],
                "age": active_row[6]
            }

            # Check against the server to see if a member w/ their username exists on the server

            member = self.find_member_by_username(active_row[1], self.guild)

            if member:
                parts["discord_id"] = member.id
                parts["discord_joined_at"] = str(member.joined_at)
                parts["discord_created_at"] = str(member.created_at)

                self.wks.update_acell("M{}".format(cell.row), member.id)
                processed_members.append(member)

            else:
                parts["discord_id"] = None
                parts["discord_joined_at"] = None
                parts["discord_created_at"] = None

            if last_valid_row < cell.row:
                log.info("New reg app submitted by {}. Discord verified: {}".format(member.name if member else active_row[1],
                                                                                member is not None))
                self.config.hmset("config:regulars:apps:{}".format(member.id), parts)

        if last_checked_row != last_valid_row:
            self.config.set("config:regulars:last_checked_row", last_valid_row)

        return processed_members

    # @checks.is_pokemon_mod()
    @commands.command(hidden=True)
    async def generate_report(self, ctx, user_id: str):
        """
        Generate a report on a user, including the amount of messages that they have sent, what channels they talked in,
        how many messages per channel, etc.
        Pull that from the existing logs.
        Might take a while
        """

        # TODO: Only track messages that the user has filled out since a certain date

        if ctx.message.mentions:
            user_id = ctx.message.mentions[0].id

        member = discord.utils.get(self.guild.members, id=user_id)
        if member is None:
            await ctx.send("Member not found in the server.")
            return

        # Get the time that the member filed their confirmation
        elif not self.config.exists("config:regulars:apps:{}".format(user_id)):
            await ctx.send("Reg submission from user not found. Continuing without sub info.")

        else:
            reg_app = self.config.hmgetall("config:regulars:apps:{}".format(user_id))
            output = report_output.format(member, reg_app["timestamp_from_form"])
            await ctx.say(output)

        async with ctx.typing():

            # Generate activity report

            data = {}
            activity_output = ""
            channels = glob.glob("message_cache/channels/*")  # Find a neater solution to just find the filename

            # Get the number of messages the user has posted per channel

            for channel_path in channels:
                try:
                    channel_id = re.findall(r'(\d{17,18})', channel_path)[0]  # Find the ID from the path
                except IndexError:  # If the path has no ID in it
                    continue
                # Get the total number of messages the user has send among channels
                with open(channel_path, "r", encoding="utf-8") as channel_data:
                    data[channel_id] = 0
                    for line in channel_data.readlines():
                        if "({})".format(user_id) in line:  # User id is at the start of most new lines
                            data[channel_id] += 1  # count the total number of lines that they have posted per chan.

            # Get the names of those channels from the server, if they exist.

            for channel_id, message_count in data.items():
                channel_obj = self.bot.get_channel(channel_id)
                if channel_obj is None:
                    channel_name = "[DELETED CHANNEL] (ID: {})".format(channel_id)
                else:
                    channel_name = channel_obj.name

                if data[channel_id] == 0:
                    continue

                activity_output += "{}: **{}** messages\n".format(channel_name, data[channel_id])

                if len(activity_output) > 1900:  # If the message is too long, just send it and make a new one
                    await ctx.send(activity_output)
                    activity_output = ""

            if not activity_output:
                activity_output = "Nothing found."
        await ctx.send(activity_output)

    @checks.r_pokemon()
    @checks.is_pokemon_mod()
    @commands.command(hidden=True)
    async def reset_reg_requests(self):
        """Reset all active regular requests"""
        self.config.delete(*self.config.scan_iter("config:regulars:apps:*"))

    @commands.Cog.listener()
    async def on_timer_update(self, seconds):
        if seconds % 120 == 0:
            if self.creds.access_token_expired:
                self.gclient.login()  # Re-auth if our access tokens expire

            new_apps = await self.process_reg_document()
            for member in new_apps:
                await member.send(messages.reg_confirmation)

def setup(bot):
    bot.add_cog(RegularUtils(bot))


"""
Illustrates how to work with the kd_tree module.

"""

import kdtree

# Create an empty tree by specifying the number of
# dimensions its points will have
emptyTree = kdtree.create(dimensions=3)

# A kd-tree can contain different kinds of points, for example tuples
point1 = (2, 3, 4)

# Lists can also be used as points
point2 = [4, 5, 6]

# Other objects that support indexing can be used, too
import collections
Point = collections.namedtuple('Point', 'x y z')
point3 = Point(5, 3, 2)

# A tree is created from a list of points
tree = kdtree.create([point1, point2, point3])

# Each (sub)tree is represented by its root node
print(tree)

# Adds a tuple to the tree
tree.add( (5, 4, 3) )

# Removes the previously added point and returns the new root
tree = tree.remove( (5, 4, 3) )

# Retrieving the Tree in inorder
print(list(tree.inorder()))

# Retrieving the Tree in level order
print(list(kdtree.level_order(tree)))

# Find the nearest node to the location (1, 2, 3)
tree.search_nn( (1, 2, 3) )

# Add a point to make the tree more interesting
tree.add( (10, 2, 1) )

# Visualize the Tree
kdtree.visualize(tree)

# Take the right subtree of the root
subtree = tree.right

# and detatch it
tree.right = None
kdtree.visualize(tree)
kdtree.visualize(subtree)

# and re-attach it
tree.right = subtree
kdtree.visualize(tree)

# Add a node to make the tree unbalanced
print(tree.is_balanced)
tree.add( (6, 1, 5) )
tree.is_balanced
kdtree.visualize(tree)

# rebalance the tree
tree = tree.rebalance()
print(tree.is_balanced)
kdtree.visualize(tree)


import ForecastHandler, ExpenseForecast
import BudgetSet, AccountSet, MemoRuleSet
import MilestoneSet, MemoMilestone, AccountMilestone

import CompositeMilestone

import logging
import datetime

log_format = '%(asctime)s - %(levelname)-8s - %(message)s'
l_formatter = logging.Formatter(log_format)

l_stream = logging.StreamHandler()
l_stream.setFormatter(l_formatter)
l_stream.setLevel(logging.INFO)

l_file = logging.FileHandler('scratch__'+datetime.datetime.now().strftime('%Y%m%d_%H%M%S')+'.log')
l_file.setFormatter(l_formatter)
l_file.setLevel(logging.INFO)

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.propagate = False
logger.handlers.clear()
logger.addHandler(l_stream)
logger.addHandler(l_file)

if __name__ == '__main__':

    # account_set = AccountSet.AccountSet([])
    #
    # account_set.addAccount(name='Checking',
    #                        balance=1000,
    #                        min_balance=0,
    #                        max_balance=float('Inf'),
    #                        account_type="checking")
    #
    # budget_set = BudgetSet.BudgetSet([])
    # budget_set.addBudgetItem('20000101', '20000101', 1, 'once', '1', 'Core', deferrable=False, partial_payment_allowed=False)
    #
    # account_milestone__list = []
    # memo_milestone__list = []
    # composite_milestone__list = []
    # milestone_names_for_composite_milestone = ['Account Milestone 1','Memo Milestone 1']
    #
    # account_milestone__list.append(AccountMilestone.AccountMilestone('Account Milestone 1','Checking',0,10000))
    # memo_milestone__list.append(MemoMilestone.MemoMilestone('Memo Milestone 1', '.*'))
    #
    # composite_milestone__list.append( CompositeMilestone.CompositeMilestone('composite milestone name 1',account_milestone__list,memo_milestone__list,milestone_names_for_composite_milestone) )
    #
    #
    #
    # MS = MilestoneSet.MilestoneSet(account_set,budget_set,account_milestone__list,memo_milestone__list,composite_milestone__list)

    F = ForecastHandler.ForecastHandler()

    F.initialize_from_excel_file('expense_forecast__milestone_testing__input.ods')

    F.run_forecasts()

    #print(F.read_results_from_disk())

    # E1 = ExpenseForecast.initialize_from_json_file(path_to_json='Forecast__2023_04_07__14_06_02__031534.json')
    # E1.appendSummaryLines()
    #
    # E2 = ExpenseForecast.initialize_from_json_file(path_to_json='Forecast__2023_04_07__13_27_36__020658.json')
    # E2.appendSummaryLines()

    #print(E1.evaulateMemoMilestone('EMT class'))

    #F.generateCompareTwoForecastsHTMLReport(E1,E2)

    # core_budget_set = [['1']]
    #
    # CoreBudgetSet = BudgetSet.BudgetSet([])
    # CoreBudgetSet.addBudgetItem('20000101','20000101',1,'once','1','Core',deferrable=False,partial_payment_allowed=False)
    #
    # BudgetSetA2 = BudgetSet.BudgetSet([])
    # BudgetSetA2.addBudgetItem('20000101', '20000101', 1, 'once', '1', 'A2', deferrable=False, partial_payment_allowed=False)
    #
    # BudgetSetB2 = BudgetSet.BudgetSet([])
    # BudgetSetB2.addBudgetItem('20000101', '20000101', 1, 'once', '1', 'B2', deferrable=False, partial_payment_allowed=False)
    #
    # BudgetSetC3 = BudgetSet.BudgetSet([])
    # BudgetSetC3.addBudgetItem('20000101', '20000101', 1, 'once', '1', 'C3', deferrable=False, partial_payment_allowed=False)
    #
    # BudgetSetD3 = BudgetSet.BudgetSet([])
    # BudgetSetD3.addBudgetItem('20000101', '20000101', 1, 'once', '1', 'D3', deferrable=False, partial_payment_allowed=False)
    #
    # BudgetSetE3 = BudgetSet.BudgetSet([])
    # BudgetSetE3.addBudgetItem('20000101', '20000101', 1, 'once', '1', 'E3', deferrable=False, partial_payment_allowed=False)
    #
    # BudgetSetF4 = BudgetSet.BudgetSet([])
    # BudgetSetF4.addBudgetItem('20000101', '20000101', 1, 'once', '1', 'F4', deferrable=False, partial_payment_allowed=False)
    #
    # list_of_lists_of_budget_sets = [
    #     [ BudgetSetA2, BudgetSetB2 ],
    #     [ BudgetSetC3, BudgetSetD3, BudgetSetE3 ],
    #     [ BudgetSetF4 ]
    # ]
    #
    # F = ForecastHandler.ForecastHandler()
    #
    # account_set = AccountSet.AccountSet([])
    # memo_rule_set = MemoRuleSet.MemoRuleSet([])
    #
    # account_set.addAccount(name='Checking',
    #                        balance=1000,
    #                        min_balance=0,
    #                        max_balance=float('Inf'),
    #                        account_type="checking")
    #
    # memo_rule_set.addMemoRule(memo_regex='.*', account_from='Checking', account_to=None, transaction_priority=1)
    #
    # F.calculateMultipleChooseOne(account_set, CoreBudgetSet , memo_rule_set, '20000101', '20000103', list_of_lists_of_budget_sets)


"""
What is a mayor?  It is an abstraction of an binding map from attribute names
to editable widgets on screen.  It has two main incarnations.  One is a simple
binding map from attributes to widgets in a form layout.  The second is a grid
with rows and columns.
"""

import contextlib
from PySide6 import QtCore, QtWidgets
import valix
import apputils
from . import bindings
from . import utils

class SaveError(RuntimeError):
    pass

def commit_mayor(mayor):
    import apputils.widgets as widgets

    if isinstance(mayor, widgets.TableView):
        mayor.commit_editors()
    elif isinstance(mayor, bindings.Binder):
        if mayor.bound == None:
            raise NotImplementedError("unconnected singleton mayor ... mayday")
        mayor.save(mayor.bound)
    else:
        raise NotImplementedError("unsupported type of mayor")

def focus_to_edit(mayor, error):
    import apputils.widgets as widgets

    if isinstance(mayor, widgets.TableView):
        index = mayor.model().index_object_column(error[0], error[1])
        mayor.setCurrentIndex(index)
    elif isinstance(mayor, bindings.Binder):
        w = mayor.widgets[error[1]]
        w.setFocus()
    else:
        raise NotImplementedError("unsupported type of mayor")

def highlight_errors(mayor, errors):
    import apputils.widgets as widgets

    if isinstance(mayor, widgets.TableView):
        mayor.model().update_invalid_fields([], errors)

def containing_mayor(row, mayors):
    for mayor in mayors:
        import apputils.widgets as widgets

        if isinstance(mayor, widgets.TableView):
            if row in mayor.model()._main_rows:
                return mayor
        elif isinstance(mayor, bindings.Binder):
            if mayor.bound == row:
                return mayor
        else:
            raise NotImplementedError("unsupported type of mayor")
    return None

class ValidationSession:
    def __init__(self, frame):
        self.frame = frame

        self._errors = []

    def class_validate_rowset(self, klass, rows, mayor):
        x = valix.class_validate_rowset(klass, rows)
        if len(x):
            self._errors.append((mayor, x))

    def class_validate_row(self, klass, row, mayor):
        self.class_validate_rowset(klass, [row], mayor)

    def has_errors(self):
        return len(self._errors)

    def errors(self):
        for e in self._errors:
            mayor = containing_mayor(e[0], self._mayors)
            yield mayor, e

    def highlight_mayors(self):
        grouped = {id(m): (m, []) for m in self._mayors}
        for e in self._errors:
            mayor = containing_mayor(e[0], self._mayors)
            grouped[id(mayor)][1].append(e)
        for m, errors in grouped.items():
            if len(errors[1]) > 0:
                highlight_errors(*errors)

    def finalize(self):
        self.frame.finalize_session(self)

VAL_FRAME_CSS = """
QFrame {
    background: --;
}
""".replace(
    "--", valix.INVALID_RGBA
)

class ValidationFrame(QtWidgets.QFrame):
    focus_error = QtCore.Signal()

    def __init__(self, parent):
        super(ValidationFrame, self).__init__(parent)

        self.hide()

        self.obscured = parent
        self.obscured.installEventFilter(self)

        self.layout = QtWidgets.QHBoxLayout(self)
        self.status = QtWidgets.QLabel("no errors")
        self.gobtn = QtWidgets.QPushButton("&Fix Error")
        self.gobtn.clicked.connect(self.focus_error.emit)
        self.hidebtn = QtWidgets.QPushButton("&X")
        self.hidebtn.clicked.connect(self.hide)
        self.layout.addWidget(self.status)
        self.layout.addStretch(12)
        self.layout.addWidget(self.gobtn)
        self.layout.addWidget(self.hidebtn)

        self.setStyleSheet(VAL_FRAME_CSS)
        self._set_size()

        self.hide_timer = QtCore.QTimer(self)
        self.hide_timer.setInterval(1000 * 5)
        self.hide_timer.setSingleShot(True)
        self.hide_timer.timeout.connect(self.hide)

    def _set_size(self):
        s = self.obscured.size()
        s.setHeight(36)
        self.resize(s)

    def eventFilter(self, obj, ev):
        if obj == self.obscured and ev.type() == QtCore.QEvent.Resize:
            self._set_size()
            # self.show()
        return False

    def hideEvent(self, event):
        self.hide_timer.stop()
        return super(ValidationFrame, self).hideEvent(event)

    def session(self):
        return ValidationSession(self)

    def finalize_session(self, session):
        if session.has_errors():
            mayor, error = next(session.errors())
            self.status.setText(error[2])
            self.show()
            self.hide_timer.start()
            focus_to_edit(mayor, error)
            raise SaveError("validation failed")

class DocumentTracker(QtCore.QObject):
    dirty_change = QtCore.Signal(bool)  # is_dirty
    post_save = QtCore.Signal(str)  # mode

    def __init__(self, parent, save_callback):
        super(DocumentTracker, self).__init__(parent)

        self.widparent = parent
        self.save_callback = save_callback

        self.dirty = False
        self._dirty_fields = set()
        self.is_new_document = False
        self.load_lockout = False
        self._mayors = None

    @contextlib.contextmanager
    def loading(self, reset=True):
        self.load_lockout = True
        yield
        self.load_lockout = False
        if reset:
            self.reset_dirty()

    def connect_button(self, button):
        button.clicked.connect(lambda: self.command_save(asksave=False))
        self.dirty_change.connect(button.setEnabled)
        button.setEnabled(self.dirty)

    def connect_action(self, action):
        action.triggered.connect(lambda: self.command_save(asksave=False))
        self.dirty_change.connect(action.setEnabled)
        action.setEnabled(self.dirty)

    def is_dirty(self):
        return self.dirty

    def is_dirty_attr(self, row, attr):
        return (id(row), attr) in self._dirty_fields

    def dirty_rows(self, rows):
        dirties = set([x for x, _ in self._dirty_fields])
        return [row for row in rows if id(row) in dirties]

    def set_mayor_list(self, mayors):
        self._mayors = list(mayors)

    def reset_dirty(self):
        self.dirty = False
        self._dirty_fields = set()
        self.dirty_change.emit(self.dirty)

    def set_dirty(self, row, attr):
        if not self.load_lockout:
            self.dirty = True
            self._dirty_fields.add((id(row), attr))

            self.dirty_change.emit(self.dirty)

    def ask_save(self):
        """
        Returns one of 'Yes', 'No', 'Cancel'
        """
        if not self.dirty:
            return "No"
        return apputils.message(
            self.widparent,
            "Do you want to save changes?",
            buttons=["Yes", "No", "Cancel"],
        )

    def _commit_mayors(self):
        for mayor in self._mayors:
            commit_mayor(mayor)

    def _core_save(self, asksave):
        if self._mayors == None:
            return True

        self._commit_mayors()
        save = self.ask_save() if asksave else ("Yes" if self.dirty else "No")
        reset = False
        if save == "Yes":
            try:
                self.save_callback()
                reset = True
            except SaveError:
                return False
            except:
                utils.exception_message(
                    self.widparent, "Error Saving", logged="unknown"
                )
                return False
        if save == "Cancel":
            return False
        if reset:
            self.reset_dirty()

        return True

    def window_close(self, asksave=True):
        result = self._core_save(asksave=asksave)
        if result:
            self.post_save.emit("window_close")
        return result

    def window_new_document(self, asksave=True):
        result = self._core_save(asksave=asksave)
        if result:
            self.post_save.emit("window_close")
        return result

    def command_save(self, asksave=False, enforce_persist_new=False):
        if self.is_new_document and enforce_persist_new:
            self.set_dirty(None, "__new__")

        result = self._core_save(asksave)
        if result:
            self.post_save.emit("command_save")
        return result


#!/usr/bin/env python3
'''
|
|------------------------------------------------------------------------------
|
|  PARSE_ARGUMENTS.PY
|
|  UPDATED:    2023-10-04
|  AUTHOR:     kent campbell
|  CONTACT:    campbellkb@ornl.gov
|
|  DESCRIPTION
|
|     This module provides functions for processing command line arguments
|
|------------------------------------------------------------------------------
|
'''
# | Python 3 standard library |
import argparse

from os.path import isdir, isfile
from sys import stdout

# | Global Variable(s) |
__version__ = '3.3.2'

# |----------------------------------------------------------------------------
# | command line argument handling functions
# |----------------------------------------------------------------------------
def argument_handler():
    '''Primary argument handler.'''

    # Parser object.
    p = argparse.ArgumentParser(description = '----- Granulemeta v' + __version__ + ' -----')

    # Determine where to send output. Defaults to stdout.
    p.add_argument('-o', '--outfile', default = stdout, type = str,
                   help = 'the path to an output file (with ext)')

    # Specify a configuration file at the command line
    p.add_argument('-c', '--config', default = None, type = _file_and_dir_validator,
                   help = 'a configuration file')

    # Exclude all other types but: 'raster', 'vector', 'mdim', 'table', ...
    p.add_argument('-t', '--type', type = _type_validator,
                   help = "input data type: [raster, vector, mdim, table, audio, envi, archive]")

    # Only process a certain number of granules before stopping execution
    p.add_argument('-n', '--nfiles', default = -1, type = int,
                   help = 'number of granules to process before stopping execution (defaults to -1 for all files).')

    # Determine where to send output. Defaults to stdout.
    p.add_argument('--outtype', default = 'gmout', type = str,
                   help = 'type of output file to produce')

    # If the output file already exists, overwrite its contents instead of appending new content only.
    p.add_argument('--overwrite', default = False, action = 'store_true',
                   help = 'a flag for enabling the overwriting of an existing output file')

    # Opt to format output for Daymet Lower Latency granules.
    p.add_argument('--daymet', default = False, action = 'store_true',
                   help = 'a flag indicating format for daymet')

    # Opt to keep endtime as is instead of adjusting ones that are at midnight.
    p.add_argument('--no-adjust-midnight-endtime', default = False, action = 'store_true',
                   help = 'a flag for disabling the auto-adjustment of the endtime')

    # Opt to remove aux.xml files from the input directory.
    p.add_argument('--keep_aux_xml', default = False, action = 'store_true',
                   help = 'a flag to disable the automatic deletion of aux.xml files')

    # Opt to leave bounding box and time extents blank when not found.
    p.add_argument('--force_blanks', default = False, action = 'store_true',
                   help = 'a flag to leave bounding box and time extents blank when not found instead of exiting with an error message')

    p.add_argument('--no_stats', default = False, action = 'store_true',
                   help = 'a flag to prevent calculating statistics on granules during execution')

    p.add_argument('--no_checksum', default = False, action = 'store_true',
                   help = 'a flag to prevent calculating the md5 checksum on granules during execution')

    p.add_argument('--add_ext', default = False, action = 'store_true',
                   help = 'a flag to add \'.bin\' extensions to extensionless ENVI files')

    p.add_argument('--lat_name', default = None, type = str,
                   help = 'a string providing the name of the column (in the granule) from which the script extracts the latitude values.')

    p.add_argument('--lon_name', default = None, type = str,
                   help = 'a string providing the name of the column (in the granule) from which the script extracts the longitude values.')

    p.add_argument('--date_name', default = None, type = str,
                   help = 'a string providing the name of the column (in the granule) from which the script extracts the date values.')

    p.add_argument('--time_name', default = None, type = str,
                   help = 'a string providing the name of the column (in the granule) from which the script extracts the time values.')

    p.add_argument('--version', default = False, action = 'store_true',
                   help = 'a flag to allow for the version number of granulemeta to be printed')

    # Positional argument 1: Input directory or file.
    p.add_argument('path', type = _file_and_dir_validator, nargs = '+',
                   help = 'an input directory (OR file OR sequence of files OR glob pattern(s)) to traverse for metadata')

    # Return parsed arguments.
    return p.parse_args()

def _type_validator(input_type: str):
    '''An input 'type' validator.'''

    # Valid inputs:
    v = ['raster', 'vector', 'mdim', 'table', 'audio', 'envi', 'archive']

    # Bad input message:
    e = "Input data type ('-t', '--type') must be from: {}.".format(v)

    # If the input data type is not valid, raise ValueError. Else return it.
    if input_type not in v:
        raise ValueError(e)
    else:
        return input_type

def _file_and_dir_validator(input_path: str):
    '''A file or directory path validator.'''
    if isinstance(input_path, str):
        if not any([isfile(input_path), isdir(input_path)]):
            print('No file(s) found corresponding to supplied PATH argument')
            exit(1)
        else:
            return input_path
    elif isinstance(input_path, list):
        for x in input_path:
            if not any([isfile(x), isdir(x)]):
                print('No file(s) found corresponding to supplied PATH argument')
                exit(1)
        return(input_path)


"""
    OBJECTIVES:
    - add more nodes that are connected to at least 2 other nodes inside artist_data
    - to mine more neighbors of nodes from IndieNesia playlist (since there are not enough data/imbalanced)
    - get rid or avoid mining nodes with degree == 1

    TO-DO:
    - get rids nodes with deg==1:
        - for existings(nodes degree=1):
            - delete (preferred in top hits playlist)
            - find existing neighbors in dump (preferred in IndieNesia playlist)
        - for new:
            - check if nodes (artist_id) is connected with the in_traverse
              && at least the other one* inside artist_data
              *(listed in neighbors of all existing data (artist_data))
              [YES]--> add to queue

    STEPS
    [LOCAL]
    - read from local data, both in dump and artist_data.
    - get collaborators of indienesia (with no exclusion), only in dump.
    - filter out gotten collaborators ==> list_of_in_dump.
    - create intersection of set collaborators of each in list_of_in_dump...
    - ... with everyone in artist_data.
    - [one transaction with DB] Add to traverse queue if len intersection >= 2

    [DB CONNECTION]
    - after queue creation complete
    - move all collaborators from dump (id in queue) to artist_data
"""
import json

import time
from tqdm import tqdm
from google.cloud import firestore
import os

os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = './credentials/spotify-co-featuring2gra-34f8c-firebase-adminsdk-k0pe2-bfcd86823f.json'

db = firestore.Client(
    project="spotify-co-featuring2gra-34f8c"
)

DUMP_PATH = 'backup/data/export/staging-exceeds_dump_v2.json'
ARTIST_DATA_PATH = 'backup/data/export/staging-artist_data_v2.json'

def load_ids(filepath):
    excluded_ids = []
    with open(filepath) as json_file:
        data = json.load(json_file)
        for x in data['collections']:
            excluded_ids.append(x)
    return excluded_ids

def load_data(filepath):
    data = {}
    with open(filepath) as json_file:
        data = json.load(json_file)
    return data

def indienesia_only(data):
    collections = data['collections']
    out = {'collections':{}}
    for artist_id, payload in collections.items():
        if payload['playlist'] == "IndieNesia":
            out['collections'][artist_id] = payload
    return out

def get_excluded_collaborators_ids(data, artist_ids, dump_ids):
    collections = data['collections']
    merged = set()
    for artist_id, payload in collections.items():
        collaborators = payload['collaborators']
        merged = merged.union(set(collaborators).intersection(dump_ids))
    out = merged - set(artist_ids)
    return out

def filter_data_by_ids(data, artist_ids):
    collections = data['collections']
    out = {
        'collections': {a_id: collections[a_id] for a_id in artist_ids}
    }
    return out

def filter_data_if_intersection_mteq_2(data, artist_ids):
    collections = data['collections']
    out = {'collections':{}}
    for artist_id, payload in collections.items():
        collaborators = payload['collaborators']
        intersected = set(collaborators).intersection(set(artist_ids))
        if len(intersected) >= 2:
            out['collections'][artist_id] = payload
            artist_ids.append(artist_id)
    return out

def update_db(queues_ids, excluded_indie_collaborator_data):
    print('updating_db....')
    # transaction = db.transaction()
    queue_doc_ref = db.collection('staging-dashboard').document('queue')
    # dump_coll_ref = db.collection('staging-exceeds_dump')
    # artist_collection_ref = db.collection('staging-artist_data')
    dashboard_queue_ref = db.collection('staging-dashboard').document('total_in_queue')
    # dashboard_nodes_ref = db.collection('staging-dashboard').document('total_nodes')

    batch = db.batch()
    # for queues_id in queues_ids:
        # payload = excluded_indie_collaborator_data['collections'][queues_id]
        # doc_ref = artist_collection_ref.document(queues_id)
        # dump_ref = dump_coll_ref.document(queues_id)
        # batch.set(doc_ref, payload)
        # print(f'delete from dump: {excluded_indie_collaborator_data["collections"][queues_id]["name"]}')
        # batch.delete(dump_ref)

    # total_nodes = dashboard_nodes_ref.get().get('value')
    batch.update(queue_doc_ref, {'ids':queues_ids})
    batch.update(dashboard_queue_ref, {'value':len(queues_ids)})
    # batch.update(dashboard_nodes_ref, {'value': total_nodes + len(queues_ids)})
    batch.commit()

def run():
    # get data
    dump_ids = load_ids(DUMP_PATH)
    artist_data_ids = load_ids(ARTIST_DATA_PATH)
    artist_data = load_data(ARTIST_DATA_PATH)
    dump_data = load_data(DUMP_PATH)
    # filter to indinesia only
    indienesia_data = indienesia_only(artist_data)
    # get ids of excluded indinesia's collaborators (in dump)
    excluded_indie_collaborator_ids = get_excluded_collaborators_ids(indienesia_data, artist_data_ids, dump_ids)
    # query exclude_indie_collaborator_ids from dump
    excluded_indie_collaborator_data = filter_data_by_ids(dump_data, excluded_indie_collaborator_ids)
    # filter out again which is associated to at least 2 artist_ids
    filtered_excluded_indie_collaborator_data = filter_data_if_intersection_mteq_2(excluded_indie_collaborator_data, artist_data_ids)
    first_queue_ids = list(filtered_excluded_indie_collaborator_data['collections'].keys())
    # print(first_queue_ids)
    update_db(first_queue_ids, filtered_excluded_indie_collaborator_data)

    # ------------
    # add to queue
    # push data & delete docuemnt to artist_id (1 batched transaction)
    # start mining

    # [LOCAL] optional, while mining
    # create reference = collaborators of filtered_excluded_indie_collaborator_data which in artist_data
    # update reference's collaborators --> add artist id of related filtered_excluded_indie_collaborator_data

    # recompute degree for each artist_ids in indienesia
    # if len == 1:
    # check original collaborators outside

run()


from setuptools import setup

from veh import __version__

classifiers = [
    'Development Status :: 3 - Alpha',
    'Intended Audience :: Developers',
    'License :: OSI Approved :: GNU General Public License (GPL)',
    'Operating System :: OS Independent',
    'Programming Language :: Python',
    'Topic :: Utilities',
]

setup(
    name = "veh",
    version = __version__,
    description = "virtualenv for hg",
    long_description = """Tie virtualenvs to individual mercurial repositorys.""",
    license = "GNU GPL v3",
    author = "Nic Ferrier",
    author_email = "nic@ferrier.me.uk",
    url = "http://github.com/nicferrier/veh",
    download_url="http://github.com/nicferrier/veh/downloads",
    platforms = ["any"],
    packages=['veh'],
    install_requires=['pip', 'virtualenv', 'Mercurial'],
    entry_points = {
       'console_scripts': [
            'veh = veh:main'
        ],
    },
    classifiers =  classifiers,
    )


#!/usr/bin/env python3

from influxdb import InfluxDBClient
import paho.mqtt.client as mqtt
import time
import logging
import argparse

logFilename = '/home/pi/influxKairos-' + time.strftime('%Y%m%d-%H-%M-%S') + '.log'
logging.basicConfig(filename=logFilename,level=logging.INFO)

class Data:
    """Class for storing MQQT data with timestamp"""
    def __init__(self, database, cloud_broker):
        self.broker = cloud_broker
        self.database = database
        self.timestamp = 0
        self.dbdata = [
        {
            "measurement": "temp",
            "tags": {
                "region": "florianopolis"
            },
            "fields": {
                "Float_value": 0.0,
                "error": 0
            }
        }
    ]

    def add_value(self, data, error):
        self.dbdata[0]['fields']['Float_value'] = data
        self.dbdata[0]['fields']['error'] = error
        self.timestamp = time.time()

    def add_record(self, measurement, data, error):
        self.dbdata[0]['measurement'] = measurement
        self.dbdata[0]['fields']['Float_value'] = data
        self.dbdata[0]['fields']['error'] = error
        self.timestamp = time.time()

        self.database.write_points(self.dbdata)

        print(self.dbdata)

    def set_error(self, error):
        self.dbdata[0]['fields']['error'] = error

    def is_new(self, timestamp):
        if self.timestamp == timestamp:
            return False
        else:
            return True

def on_connect(client, userdata, flags, rc):

    logging.warning('Connected to localhost with result code: ' + str(rc) + ' ' + time.strftime('%Y-%m-%dT%H:%M:%S'))
    print('Connected to localhost with result code ' + str(rc))

    # Subscribing in on_connect() means that if we lose the connection and
    # reconnect then subscriptions will be renewed.
    # client.subscribe("$SYS/#")

    client.subscribe('tanque/#')

def on_connect_cloud(client, userdata, flags, rc):

    logging.warning('Connected to cloud with result code: ' + str(rc) + ' ' + time.strftime('%Y-%m-%dT%H:%M:%S'))
    print('Connected to cloud with result code ' + str(rc))

def on_disconnect(client, userdata, rc):
    if rc != 0:
        print('Unexpected disconnection: ' + str(rc) + ' ' + time.strftime('%Y-%m-%dT%H:%M:%S'))
        logging.warning('Unexpected disconnection: ' + str(rc) + ' ' + time.strftime('%Y-%m-%dT%H:%M:%S'))

    # client.reconnect()

def on_message(client, userdata, message):
    #global temperature# = Data()

    # Publish messages on the cloud
    ret = userdata.broker.publish(message.topic, message.payload)

    # Add values to the databse
    topic_split = str(message.topic).split('/')
    payload_split = list(message.payload.decode('ASCII'))
    data_value = (''.join(payload_split).split(';',1))

    measurement = ''.join(topic_split)
    userdata.add_record(measurement,float(data_value[0]), int(data_value[1]))

def main():

    myData = []

    timestamp = []
    timestamp.append(0)
    timestamp.append(0)
    timestamp.append(0)

    parser = argparse.ArgumentParser(description='Simple mqtt to influx gateWay')
    parser.add_argument('broker', help='Broker address')
    parser.add_argument('user', help='Broker user')
    parser.add_argument('password', help='Broker password')
    parser.add_argument('DBuser', help='InfluxDB user')
    parser.add_argument('DBpassword', help='InfluxDB password')

    options = parser.parse_args()

    print('Connecting to {} with user {}.'.format(options.broker, options.user))
    logging.info('Connecting to {} with user {}.'.format(options.broker, options.user))

    cloud_mqtt = mqtt.Client()
    cloud_mqtt.username_pw_set('juca',  'password')
    # cloud_mqtt.connect('m14.cloudmqtt.com')
    cloud_mqtt.connect('192.168.1.2')
    cloud_mqtt.on_connect = on_connect_cloud

    local_mqtt = mqtt.Client()
    local_mqtt.username_pw_set(options.user,  options.password)
    local_mqtt.connect(options.broker)
    local_mqtt.on_disconnect = on_disconnect
    local_mqtt.on_message = on_message
    local_mqtt.on_connect = on_connect

    client = InfluxDBClient('localhost', 8086, options.DBuser, options.DBpassword)
    client.switch_database('kairos')

    data = Data(client, cloud_mqtt)
    local_mqtt.user_data_set(data)

    local_mqtt.loop_start()
    cloud_mqtt.loop_start()

    try:
        while True:

            time.sleep(10)

    except KeyboardInterrupt:
        print('Ending...')
        local_mqtt.disconnect()
        client.close()

if __name__ == '__main__':
    main()


import re
import json
from operator import attrgetter
from utils import AthleteScore

class IFSCResultParser:

    rounds = ["qualification", "semis", "finals"]

    def __init__(self, json_results):
        self.json_results = json_results
        self.ranking = self.json_results["ranking"]
        self.score_regex = re.compile(r"\d+")

    def who_qualified(self, round="finals"):

        qualifiers = []

        # lookup a threshold which tells us how many rounds they got through
        # based on the name of the round we want
        # why doesn't python have built in enums
        threshold = self.rounds.index(round)

        for athlete in self.ranking:
            if len(athlete["rounds"]) > threshold:
                qualifiers.append(athlete)

        return qualifiers

    def get_rank_after_qualification(self, qualifiers):

        ranking_after_qualification = []

        for athlete in qualifiers:

            id, first_name, last_name = (
                athlete["athlete_id"],
                athlete["firstname"],
                athlete["lastname"],
            )

            qualification_score = athlete["rounds"][0]["score"]
            qualification_score = qualification_score.replace("\xa0", " ")

            # e.g. '4T5z 14 9',
            score_string, top_attempts, zone_attempts = qualification_score.split(" ")

            tops, zones = self.score_regex.findall(score_string)

            athlete_score = AthleteScore(
                id, first_name, last_name, tops, zones, top_attempts, zone_attempts
            )

            ranking_after_qualification.append(athlete_score)

        ranking_after_qualification = sorted(
            ranking_after_qualification,
            key=attrgetter("tops", "zones", "top_attempts", "zone_attempts"),
            reverse=True,
        )

        return [athlete.athlete_id for athlete in ranking_after_qualification]

    def get_rank_after_finals(self, qualifiers):

        ranking_after_finals = []

        for athlete in qualifiers:
            athlete_id, athlete_rank = athlete["athlete_id"], athlete["rank"]
            ranking_after_finals.append((athlete_id, athlete_rank))

        return [
            id
            for (id, _) in sorted(
                ranking_after_finals, key=lambda x: x[1], reverse=True
            )
        ]

    def save(self):

        with open(f"outputs/{self.json_results['event']}.json", "w+") as f:
            json.dump(self.json_results, f)