###################################################################################################
# this class checks the quality of the calculated SST according to several tests and returns an array with quality flags
# description              ¦  bit ¦ value ¦
#                          ¦      ¦       ¦
# GROSS IR/Valid VIS (day) ¦  01  ¦     1 ¦
# VIS cloud (day)          ¦  02  ¦     2 ¦
# NIR/VIS ratio (day)      ¦  03  ¦     4 ¦
# IR cloud      (night)    ¦  04  ¦     8 ¦
# Low Stratus   (night)    ¦  05  ¦    16 ¦
# SST LE -5°C OR GE 35°C   ¦  06  ¦    32 ¦
# land/sea mask            ¦  07  ¦    64 ¦
# cloud mask               ¦  08  ¦   128 ¦
# Exclude singel pixels    ¦  09  ¦   256 ¦
# Spatial STDV > 3°K       ¦  10  ¦   512 ¦
# Spatial STDV > 1.5°K     ¦  11  ¦  1024 ¦
# SZA > 55°                ¦  12  ¦  2048 ¦
# SZA > 45°                ¦  13  ¦  4096 ¦
# Glint Angle < 36°        ¦  14  ¦  8192 ¦
# Error in QFlag proc.     ¦  15  ¦ 32767 ¦
###################################################################################################

import numpy
from scipy import ndimage

class QualityCheck:
    def __init__(self, lswt, rf1, rf2, bt3, bt4, bt5, sat_za, sun_za, rel_az, lwm, cmsk, height, width, sattype):
        self.qfilter = {'cmask': 7, 'r1lim': 0.005, 'r2lim': 0.1, 'r21lim': 1.0, 'T4lim': 263.15,
                        'lswtLim': [268.15, 308.15], 'vzaLim': 55., 'georefRel': 1.8, 'georefRel2': 1.8, 'std1': 3.,
                        'std2': 1.5, 'VZA2': 45., 'sunglint': 36., 'ircldtshld': 1.0, 'strattshld': -0.6,
                        'cloudTEST': 0}
        self.org_dim = lswt.shape
        self.lswt = numpy.reshape(lswt, (height, width))
        self.rf1 = numpy.reshape(rf1, (height, width))
        self.rf2 = numpy.reshape(rf2, (height, width))
        self.bt3 = numpy.reshape(bt3, (height, width))
        self.bt4 = numpy.reshape(bt4, (height, width))
        self.bt5 = numpy.reshape(bt5, (height, width))

        self.sat_za = numpy.reshape(sat_za, (height, width))
        self.sun_za = numpy.reshape(sun_za, (height, width))
        # Relative azimuth angle at the point for the view angle between the ray to the satellite and the ray to the Sun
        self.rel_az = numpy.reshape(rel_az, (height, width))
        if lwm is not None:
            self.lwm = numpy.reshape(lwm, (height, width))
        else:
            self.lwm = None
        if cmsk is not None:
            self.cmsk = numpy.reshape(cmsk, (height, width))
        else:
            self.cmsk = None

        self.height = height
        self.width = width

        self.sattype = sattype

    def check_quality(self):
        quality_flag = numpy.zeros(self.lswt.shape)

        ###############################################################
        # GROSS IR TEST / VALID VIS TEST # if the channel 4 temperature is less than  #10° then
        # do not compute a SST (BOM)
        ###############################################################
        quality_flag += numpy.logical_and(
            numpy.logical_or(self.bt4 <= self.qfilter['T4lim'], self.rf1 < self.qfilter['r1lim']), self.sun_za < 90) * 1

        ###############################################################
        # VISIBLE CLOUD THRESHOLD TEST # reflectance higher than 0.1 then
        # do not compute a SST (BoM)
        ###############################################################
        quality_flag += numpy.logical_and(self.rf2 >= self.qfilter['r2lim'], self.sun_za < 90) * 2

        ###############################################################
        # NIR/VISIBLE RATIO TEST # ration higher than 1.0 then
        # do not compute a SST (BoM)
        ###############################################################
        if numpy.nanmax(self.sun_za) < 90:
            r21 = self.rf2 / self.rf1
            quality_flag += (r21 > self.qfilter['r21lim']) * 4

        ###############################################################
        # NIGHTTIME IR CLOUD TEST # if a calculated channel 4 temperature based on the channel 5 value
        # (channel 5 temp * 1.0439 # 11.49) differs from the actual channel 4 temperature by more than
        # 1.0° C then do not compute a SST (BoM)
        ###############################################################

        # TODO, GL, this is masking out the LAKES !!!
        syn_ch4 = (self.bt5 * 1.0439) - 11.49
        quality_flag += numpy.logical_and(numpy.absolute(syn_ch4 - self.bt4) >= self.qfilter['ircldtshld'],
                                          self.sun_za > 85) * 8

        ###############################################################
        # NIGHTTIME LOW STRATUS TEST # the difference obtained when subtracting channel 3
        # temperature from channel 5 temperature must be LE #0.6° C. (BoM)
        ###############################################################
        quality_flag += numpy.logical_and((self.bt5 - self.bt3) > self.qfilter['strattshld'], self.bt3 > 100) * 16

        ###############################################################
        # Creating final product, limiting data to #5°C to +35°C
        ###############################################################
        quality_flag += numpy.logical_or(self.lswt <= self.qfilter['lswtLim'][0],
                                         self.lswt >= self.qfilter['lswtLim'][1]) * 32

        ###############################################################
        # Applying Sea and Lake Mask
        ###############################################################
        if self.lwm is not None:
            quality_flag += (self.lwm == 0) * 64
        else:
            quality_flag += 64

        ###############################################################
        # Applying Cloud Mask
        ###############################################################

        # if self.qfilter['cloudTEST'] == 0:
        if self.cmsk is not None:
            if self.sattype == 'AVHRR':
                quality_flag += numpy.logical_or(self.cmsk >= self.qfilter['cmask'], self.cmsk == 0) * 128
            else:
                quality_flag += (self.cmsk != 0) * 128
        else:
            quality_flag += 128

        #####################################
        # EXPERIMENTEL Cloud Mask alternative TEST
        #####################################
        # if self.qfilter['cloudTEST'] == 1:
        #     if numpy.nanmax(self.sun_za) < 80:  # daytime check
        #         cmsk_alt = numpy.logical_not(numpy.logical_or(
        #             numpy.logical_and(numpy.logical_and((self.rf1 + self.rf2) <= 1.2, self.bt5 >= 265),
        #                               (self.rf1 + self.rf2) <= 0.8), self.bt5 >= 285))
        #         cmsk_alt = numpy.logical_or(numpy.logical_or(cmsk_alt, self.cmsk >= self.qfilter['cmask']),
        #                                     self.cmsk == 0)
        #         quality_flag += cmsk_alt * 128
        #     else:
        #         cmsk_alt = numpy.logical_or(
        #             numpy.logical_or(self.bt4 < 243.15, (self.bt5 - self.bt3) > self.qfilter['strattshld']),
        #             numpy.absolute(((self.bt5 * 1.0439) - 11.49) - self.bt4) <= self.qfilter['ircldtshld'])
        #         cmsk_alt = numpy.logical_or(numpy.logical_or(cmsk_alt, self.cmsk >= self.qfilter['cmask']),
        #                                     self.cmsk == 0)
        #         quality_flag += cmsk_alt * 128
        #
        ###############################################################
        # Valid pixels only those which are not completely surrounded by nonvaliddata
        # according to Schwab et al 1999
        ###############################################################
        # the nxn window total not including central pixel
        n = 3
        if quality_flag.ndim == 1:
            raise ArrayDimensionException(quality_flag.ndim)
        else:
            kernel = numpy.ones((n, n))
        # check which pix are good so far
        good_pix = (quality_flag == 0).astype(int)
        # count good neighbors
        neighbors = ndimage.convolve(good_pix, kernel)
        # use only those pix with two or more neighbors
        quality_flag += (neighbors < 2) * 256

        ###############################################################
        # Valid pixels only those which sourrounding pixels have a SDEV from lower 3°C
        # according to Schwab et al 1999
        ###############################################################
        tmp_lswt = self.lswt
        bad_pix = (good_pix == 0).nonzero()
        if bad_pix[0].size != 0:
            tmp_lswt[bad_pix] = 0
        lswt_stdv = self.spatial_stddev(tmp_lswt, 3)

        quality_flag += numpy.logical_and(numpy.isfinite(lswt_stdv), lswt_stdv > self.qfilter['std1']) * 512
        quality_flag += numpy.logical_and(numpy.isfinite(lswt_stdv), lswt_stdv > self.qfilter['std2']) * 1024

        ###############################################################
        # Use only pixels for VZA LE 45 deg.
        ###############################################################
        quality_flag += (self.sat_za > self.qfilter['vzaLim']) * 2048
        quality_flag += (self.sat_za > self.qfilter['VZA2']) * 4096

        ###############################################################
        # Calculate glint angle
        ###############################################################
        glint = numpy.arcsin(numpy.sin(numpy.radians(self.sat_za)) * numpy.sin(numpy.radians(self.sun_za)) * numpy.cos(
            numpy.radians(self.rel_az)) + numpy.cos(numpy.radians(self.sat_za)) * numpy.cos(numpy.radians(self.sun_za)))
        glint = numpy.degrees(glint)
        quality_flag += (glint < self.qfilter['sunglint']) * 8192
        return quality_flag

    def get_quality_flags(self, flag):
        quality_flags = numpy.ones(flag.shape)
        flag = numpy.uint16(flag)
        init_id = numpy.logical_not(numpy.bitwise_and(flag, 511))
        if numpy.sum(init_id) > 0:
            quality_flags[init_id.nonzero()] = 2  # Q1
            quality_flags += init_id * 2  # georef1 Q2
            q3_id = numpy.bitwise_and(init_id, numpy.logical_not(self.bit(flag, 11)))  # VZA1
            if numpy.sum(q3_id) > 0:
                quality_flags[q3_id.nonzero()] = 8  # Q3
            q4_id = numpy.bitwise_and(q3_id, numpy.logical_not(self.bit(flag, 12)))  # VZA2
            if numpy.sum(q4_id) > 0:
                quality_flags[q4_id.nonzero()] = 16  # Q4
            q5_id = numpy.bitwise_and(q4_id, numpy.logical_not(self.bit(flag, 13)))  # glint
            if numpy.sum(q5_id) > 0:
                quality_flags[q5_id.nonzero()] = 32  # Q5
            q6_id = numpy.bitwise_and(q5_id, numpy.logical_not(self.bit(flag, 10)))  # std2
            if numpy.sum(q6_id) > 0:
                quality_flags[q6_id.nonzero()] = 64  # Q6
        return quality_flags

    # Different version: quality levels have value 1 to 6, like in old processor
    # def get_quality_flags(self, flag):
    #     quality_flags = numpy.zeros(flag.shape)
    #     flag = numpy.uint16(flag)
    #     init_id = numpy.logical_not(numpy.bitwise_and(flag, 511))
    #     if numpy.sum(init_id) > 0:
    #         quality_flags[init_id.nonzero()] = 1  # Q1
    #         quality_flags += init_id  # georef1 Q2
    #         q3_id = numpy.bitwise_and(init_id, numpy.logical_not(self.bit(flag, 11)))  # VZA1
    #         if numpy.sum(q3_id) > 0:
    #             quality_flags[q3_id.nonzero()] += 1  # Q3
    #         q4_id = numpy.bitwise_and(q3_id, numpy.logical_not(self.bit(flag, 12)))  # VZA2
    #         if numpy.sum(q4_id) > 0:
    #             quality_flags[q4_id.nonzero()] += 1  # Q4
    #         q5_id = numpy.bitwise_and(q4_id, numpy.logical_not(self.bit(flag, 13)))  # glint
    #         if numpy.sum(q5_id) > 0:
    #             quality_flags[q5_id.nonzero()] += 1  # Q5
    #         q6_id = numpy.bitwise_and(q5_id, numpy.logical_not(self.bit(flag, 10)))  # std2
    #         if numpy.sum(q6_id) > 0:
    #             quality_flags[q6_id.nonzero()] += 1  # Q6
    #     return quality_flags

    @staticmethod
    def bit(value, shift):
        return numpy.bitwise_and(7, numpy.right_shift(value, shift))

    @staticmethod
    def spatial_stddev(in_data, n):
        numpy.set_printoptions(threshold=numpy.nan)
        # Initialize kernel for convolution operation
        if in_data.ndim == 1:
            raise ArrayDimensionException(in_data.ndim)
        else:
            kernel = numpy.ones((n, n))
        # Compute the square of each data point
        data_sq = ndimage.convolve(numpy.power(in_data, 2), kernel)  # square of each data point
        sub_sum = ndimage.convolve(in_data, kernel)  # sum of each n x n pixel window
        n_valid = ndimage.convolve((in_data != 0).astype(int), kernel)
        # number of valid points in n x n pixel window
        data_me = sub_sum / n_valid  # mean of n x n pixel around each pixel, invalid pixels ignored

        # Compute variance
        data_var = data_sq - 2 * data_me * sub_sum + n_valid * numpy.power(data_me, 2)
        data_var /= (n_valid - 1)

        # compute standarddeviation
        return numpy.sqrt(numpy.absolute(data_var))

class ArrayDimensionException(Exception):
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)


from absl.testing import parameterized
import numpy as np
import tensorflow as tf

from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
from official.nlp.modeling.networks import bert_encoder

# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
# guarantees forward compatibility of this code for the V2 switchover.
@keras_parameterized.run_all_keras_modes
class BertEncoderTest(keras_parameterized.TestCase):

  def tearDown(self):
    super(BertEncoderTest, self).tearDown()
    tf.keras.mixed_precision.experimental.set_policy("float32")

  def test_network_creation(self):
    hidden_size = 32
    sequence_length = 21
    # Create a small BertEncoder for testing.
    test_network = bert_encoder.BertEncoder(
        vocab_size=100,
        hidden_size=hidden_size,
        num_attention_heads=2,
        num_layers=3)
    # Create the inputs (note that the first dimension is implicit).
    word_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    mask = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    type_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    data, pooled = test_network([word_ids, mask, type_ids])

    self.assertIsInstance(test_network.transformer_layers, list)
    self.assertLen(test_network.transformer_layers, 3)
    self.assertIsInstance(test_network.pooler_layer, tf.keras.layers.Dense)

    expected_data_shape = [None, sequence_length, hidden_size]
    expected_pooled_shape = [None, hidden_size]
    self.assertAllEqual(expected_data_shape, data.shape.as_list())
    self.assertAllEqual(expected_pooled_shape, pooled.shape.as_list())

    # The default output dtype is float32.
    self.assertAllEqual(tf.float32, data.dtype)
    self.assertAllEqual(tf.float32, pooled.dtype)

    test_network_dict = bert_encoder.BertEncoder(
        vocab_size=100,
        hidden_size=hidden_size,
        num_attention_heads=2,
        num_layers=3,
        dict_outputs=True)
    # Create the inputs (note that the first dimension is implicit).
    inputs = dict(
        input_word_ids=word_ids, input_mask=mask, input_type_ids=type_ids)
    _ = test_network_dict(inputs)

    test_network_dict.set_weights(test_network.get_weights())
    batch_size = 2
    vocab_size = 100
    num_types = 2
    word_id_data = np.random.randint(
        vocab_size, size=(batch_size, sequence_length))
    mask_data = np.random.randint(2, size=(batch_size, sequence_length))
    type_id_data = np.random.randint(
        num_types, size=(batch_size, sequence_length))
    list_outputs = test_network([word_id_data, mask_data, type_id_data])
    dict_outputs = test_network_dict(
        dict(
            input_word_ids=word_id_data,
            input_mask=mask_data,
            input_type_ids=type_id_data))
    self.assertAllEqual(list_outputs[0], dict_outputs["sequence_output"])
    self.assertAllEqual(list_outputs[1], dict_outputs["pooled_output"])

  def test_all_encoder_outputs_network_creation(self):
    hidden_size = 32
    sequence_length = 21
    # Create a small BertEncoder for testing.
    test_network = bert_encoder.BertEncoder(
        vocab_size=100,
        hidden_size=hidden_size,
        num_attention_heads=2,
        num_layers=3,
        return_all_encoder_outputs=True)
    # Create the inputs (note that the first dimension is implicit).
    word_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    mask = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    type_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    all_encoder_outputs, pooled = test_network([word_ids, mask, type_ids])

    expected_data_shape = [None, sequence_length, hidden_size]
    expected_pooled_shape = [None, hidden_size]
    self.assertLen(all_encoder_outputs, 3)
    for data in all_encoder_outputs:
      self.assertAllEqual(expected_data_shape, data.shape.as_list())
    self.assertAllEqual(expected_pooled_shape, pooled.shape.as_list())

    # The default output dtype is float32.
    self.assertAllEqual(tf.float32, all_encoder_outputs[-1].dtype)
    self.assertAllEqual(tf.float32, pooled.dtype)

  def test_network_creation_with_float16_dtype(self):
    hidden_size = 32
    sequence_length = 21
    tf.keras.mixed_precision.experimental.set_policy("mixed_float16")
    # Create a small BertEncoder for testing.
    test_network = bert_encoder.BertEncoder(
        vocab_size=100,
        hidden_size=hidden_size,
        num_attention_heads=2,
        num_layers=3)
    # Create the inputs (note that the first dimension is implicit).
    word_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    mask = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    type_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    data, pooled = test_network([word_ids, mask, type_ids])

    expected_data_shape = [None, sequence_length, hidden_size]
    expected_pooled_shape = [None, hidden_size]
    self.assertAllEqual(expected_data_shape, data.shape.as_list())
    self.assertAllEqual(expected_pooled_shape, pooled.shape.as_list())

    # If float_dtype is set to float16, the data output is float32 (from a layer
    # norm) and pool output should be float16.
    self.assertAllEqual(tf.float32, data.dtype)
    self.assertAllEqual(tf.float16, pooled.dtype)

  @parameterized.named_parameters(
      ("all_sequence", None, 21),
      ("output_range", 1, 1),
  )
  def test_network_invocation(self, output_range, out_seq_len):
    hidden_size = 32
    sequence_length = 21
    vocab_size = 57
    num_types = 7
    # Create a small BertEncoder for testing.
    test_network = bert_encoder.BertEncoder(
        vocab_size=vocab_size,
        hidden_size=hidden_size,
        num_attention_heads=2,
        num_layers=3,
        type_vocab_size=num_types,
        output_range=output_range)
    # Create the inputs (note that the first dimension is implicit).
    word_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    mask = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    type_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
    data, pooled = test_network([word_ids, mask, type_ids])

    # Create a model based off of this network:
    model = tf.keras.Model([word_ids, mask, type_ids], [data, pooled])

    # Invoke the model. We can't validate the output data here (the model is too
    # complex) but this will catch structural runtime errors.
    batch_size = 3
    word_id_data = np.random.randint(
        vocab_size, size=(batch_size, sequence_length))
    mask_data = np.random.randint(2, size=(batch_size, sequence_length))
    type_id_data = np.random.randint(
        num_types, size=(batch_size, sequence_length))
    outputs = model.predict([word_id_data, mask_data, type_id_data])
    self.assertEqual(outputs[0].shape[1], out_seq_len)

    # Creates a BertEncoder with max_sequence_length != sequence_length
    max_sequence_length = 128
    test_network = bert_encoder.BertEncoder(
        vocab_size=vocab_size,
        hidden_size=hidden_size,
        max_sequence_length=max_sequence_length,
        num_attention_heads=2,
        num_layers=3,
        type_vocab_size=num_types)
    data, pooled = test_network([word_ids, mask, type_ids])
    model = tf.keras.Model([word_ids, mask, type_ids], [data, pooled])
    outputs = model.predict([word_id_data, mask_data, type_id_data])
    self.assertEqual(outputs[0].shape[1], sequence_length)

    # Creates a BertEncoder with embedding_width != hidden_size
    test_network = bert_encoder.BertEncoder(
        vocab_size=vocab_size,
        hidden_size=hidden_size,
        max_sequence_length=max_sequence_length,
        num_attention_heads=2,
        num_layers=3,
        type_vocab_size=num_types,
        embedding_width=16)
    data, pooled = test_network([word_ids, mask, type_ids])
    model = tf.keras.Model([word_ids, mask, type_ids], [data, pooled])
    outputs = model.predict([word_id_data, mask_data, type_id_data])
    self.assertEqual(outputs[0].shape[-1], hidden_size)
    self.assertTrue(hasattr(test_network, "_embedding_projection"))

  def test_serialize_deserialize(self):
    # Create a network object that sets all of its config options.
    kwargs = dict(
        vocab_size=100,
        hidden_size=32,
        num_layers=3,
        num_attention_heads=2,
        max_sequence_length=21,
        type_vocab_size=12,
        intermediate_size=1223,
        activation="relu",
        dropout_rate=0.05,
        attention_dropout_rate=0.22,
        initializer="glorot_uniform",
        return_all_encoder_outputs=False,
        output_range=-1,
        embedding_width=16,
        dict_outputs=True,
        embedding_layer=None)
    network = bert_encoder.BertEncoder(**kwargs)
    expected_config = dict(kwargs)
    expected_config["activation"] = tf.keras.activations.serialize(
        tf.keras.activations.get(expected_config["activation"]))
    expected_config["initializer"] = tf.keras.initializers.serialize(
        tf.keras.initializers.get(expected_config["initializer"]))

    self.assertEqual(network.get_config(), expected_config)
    # Create another network object from the first object's config.
    new_network = bert_encoder.BertEncoder.from_config(network.get_config())

    # Validate that the config can be forced to JSON.
    _ = network.to_json()

    # If the serialization was successful, the new config should match the old.
    self.assertAllEqual(network.get_config(), new_network.get_config())

    # Tests model saving/loading.
    model_path = self.get_temp_dir() + "/model"
    network.save(model_path)
    _ = tf.keras.models.load_model(model_path)

if __name__ == "__main__":
  tf.test.main()


#!/Users/fynn/.pyenv/eth/bin/python

import sys
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

def get_moves_per_partition(latencies_path):
    moves_per_partition = {'1': {}}
    with open(latencies_path, 'r') as f:
        for line in f:
            line = line.split()
            if line[1] == 'moveTo' or line[1] == 'move2':
                partition = line[2]
                height = int(line[3])
                if partition not in moves_per_partition:
                    moves_per_partition[partition] = {}
                if height not in moves_per_partition[partition]:
                    moves_per_partition[partition][height] = 0
                moves_per_partition[partition][height]+=1

    return moves_per_partition

def multi_plot_absolute(tput_path):
    times = []
    txs_delta = []
    first = True
    txs = 0
    prev_time = 0

    partition_id = tput_path.split('-')[-1].split('.')[0]

    timestamps = []

    begin_path = '/'.join(tput_path.split('/')[:-1]) + '/begin-experiment.txt'
    with open(begin_path, 'r') as be:
        experiment_begin = int(be.readline())

    moves_per_partition = get_moves_per_partition('/'.join(tput_path.split('/')[:-1]) + '/latencies.txt')
    moves_in_partition = moves_per_partition[partition_id]

    with open(tput_path, 'r') as f:
        i = 0
        for line in f:
            line = line.split()
            txs = int(line[0])
            timestamp = int(line[1])
            height = int(line[2])
            if timestamp < experiment_begin:
                continue
            if first:
                initial_time = timestamp
                prev_txs = txs
                prev_time = timestamp
                first = False
                continue
            timestamps.append(timestamp/1e9)
            moves = 0
            for ts in moves_in_partition:
                if ts >= prev_time and ts <= timestamp:
                    moves+=1
            if height in moves_per_partition[partition_id]:
                txs -= moves_per_partition[partition_id][height]
            if txs - prev_txs < 0:
                print('Warning', txs - prev_txs, moves)
            txs_delta.append(txs - prev_txs)
            times.append((timestamp - prev_time)/1e9)
            prev_time = timestamp
            prev_txs = txs
            i += 1

    fig, ax = plt.subplots()
    times_from_zero = [0]
    for t in times[1:]:
        times_from_zero.append(times_from_zero[-1] + t)
    # ax.plot(times, txs_delta, '.')
    # print('Avg latency:', np.average(times))
    # print('Avg txs/s:', sum(txs_delta)/((timestamp - initial_time)/1e9))
    for i in range(len(txs_delta)):
        txs_delta[i] = txs_delta[i]/times[i]
    ax.plot(times_from_zero, txs_delta, '.')
    # ax.plot([i for i in range(len(txs_delta))], txs_delta, '.')

    # ax.plot(prev_time - initial_time, txs/float((prev_time - initial_time)/1e9), '.')
    ax.set(xlabel='time (s)', ylabel='tx/s')
    ylim1, ylim2 = plt.ylim()
    plt.ylim((0, ylim2))
    plt.title('Sending to multiple clients (1000 clients)')
    plt.show()
    return fig, ax

def get_stopped_time_partition(stop_path):
    try:
        with open(stop_path, 'r') as f:
            return int(f.readline()[:-1])
    except:
        return None

def plot_aggregated(tput_paths, delta_time=5):
    raw_tputs = []
    for path in tput_paths:
        tput = []
        partition_id = path.split('-')[-1].split('.')[0]
        begin_path = '/'.join(path.split('/')[:-1]) + '/begin-experiment.txt'
        with open(begin_path, 'r') as be:
            experiment_begin = int(be.readline())
        moves_per_partition = get_moves_per_partition('/'.join(path.split('/')[:-1]) + '/latencies.txt')

        with open(path, 'r') as f:
            # f.readline()
            for line in f:
                timestamp = 0
                (total_txs, timestamp, height, validator) = line.split()
                total_txs = int(total_txs)
                timestamp = int(timestamp)
                height = int(height)
                if timestamp < experiment_begin:
                    continue

                if height in moves_per_partition[partition_id]:
                    total_txs -= moves_per_partition[partition_id][height]
                assert(total_txs >= 0)
                tput.append((timestamp, total_txs))

        raw_tputs.append(tput)

    # remove non-overlapping ends
    max_min = 0
    min_max = float('inf')
    for tput in raw_tputs:
        max_min = max(max_min, tput[0][0])
        min_max  = min(min_max, tput[-1][0])
    for tput in raw_tputs:
        while tput[0][0] < max_min:
            tput.remove(tput[0])
        while tput[-1][0] > min_max:
            tput.pop()

    aggregated_tputs = []

    delta = delta_time*1e9
    start_time = max_min

    data_points_i = [0 for i in range(len(raw_tputs))]
    running = True
    # Aggregate
    while True:
        aggregated_data = []
        for tput_i, tput in enumerate(raw_tputs):
            prev_idx = data_points_i[tput_i]
            while True:
                idx = data_points_i[tput_i]
                if idx >= len(tput):
                    running = False
                    break
                data_points_i[tput_i] += 1
                if tput[idx][0] > start_time + delta:
                    break

            if running == False:
                break

            time_passed = (tput[idx][0] - tput[prev_idx][0])/1e9
            if time_passed == 0:
                running = False
                break
            aggregated_data.append(((tput[idx][0] + tput[prev_idx][0])/2., (tput[idx][1] - tput[prev_idx][1])/time_passed ))
        if running == False:
            break
        avg_time = np.average(list(map(lambda i : i[0], aggregated_data)))
        sum_txs = sum(list(map(lambda i : i[1], aggregated_data)))

        aggregated_tputs.append((avg_time, sum_txs))
        start_time += delta

    # plot
    fig, ax = plt.subplots()
    avg_times = list(map(lambda i : (i[0] - max_min)/1e9, aggregated_tputs))

    ax.plot(avg_times, list(map(lambda i : i[1], aggregated_tputs)))
    ax.set(xlabel='time (s)', ylabel='tx/s', title='tput {} partitions'.format(len(tput_paths)))
    ylim1, ylim2 = plt.ylim()
    plt.ylim((0, ylim2))
    # plt.show()

    # for i in range(len(avg_times)):
    #     print('{} {}'.format(avg_times[i],aggregated_tputs[i][1]))

    total_tput = 0
    for p in raw_tputs:
        total_tput += p[-1][1]
    print(min_max - max_min, total_tput)

    return fig, ax

if __name__ == '__main__':
    partitions = int(sys.argv[1])
    tput_paths = ['{}/tput-partition-{}.txt'.format(sys.argv[2],i+1) for i in range(partitions)]
    fig, ax = plot_aggregated(tput_paths, delta_time=60)
    path = sys.argv[2] + '/' + 'tput-aggregated.pdf'
    fig.savefig(path)

    for tp in tput_paths:
        fig, ax = multi_plot_absolute(tp)
        p = tp[:-3] + 'pdf'
        fig.savefig(p)