#!/usr/bin/env python3
from argparse import ArgumentParser
from datetime import timedelta
from importlib import import_module
import logging.config
import os
from signal import SIGINT, SIGTERM
import sys
import time

import json
import numpy as np
import tensorflow as tf
from tensorflow.contrib import slim

import common
import lbtoolbox as lb
import loss
from nets import NET_CHOICES
from heads import HEAD_CHOICES
from utils import available_gpu_num, get_available_gpus, make_parallel

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True

parser = ArgumentParser(description='Train a ReID network.')

# Required.

parser.add_argument(
    '--experiment_root', required=True, type=common.writeable_directory,
    help='Location used to store checkpoints and dumped data.')

parser.add_argument(
    '--train_set',
    help='Path to the train_set csv file.')

parser.add_argument(
    '--image_root', type=common.readable_directory,
    help='Path that will be pre-pended to the filenames in the train_set csv.')

# Optional with sane defaults.

parser.add_argument(
    '--resume', action='store_true', default=False,
    help='When this flag is provided, all other arguments apart from the '
         'experiment_root are ignored and a previously saved set of arguments '
         'is loaded.')

parser.add_argument(
    '--model_name', default='resnet_v1_50', choices=NET_CHOICES,
    help='Name of the model to use.')

parser.add_argument(
    '--head_name', default='fc1024', choices=HEAD_CHOICES,
    help='Name of the head to use.')

parser.add_argument(
    '--embedding_dim', default=128, type=common.positive_int,
    help='Dimensionality of the embedding space.')

parser.add_argument(
    '--initial_checkpoint', default=None,
    help='Path to the checkpoint file of the pretrained network.')

# TODO move these defaults to the .sh script?
parser.add_argument(
    '--batch_p', default=32, type=common.positive_int,
    help='The number P used in the PK-batches')

parser.add_argument(
    '--batch_k', default=4, type=common.positive_int,
    help='The numberK used in the PK-batches')

parser.add_argument(
    '--net_input_height', default=256, type=common.positive_int,
    help='Height of the input directly fed into the network.')

parser.add_argument(
    '--net_input_width', default=128, type=common.positive_int,
    help='Width of the input directly fed into the network.')

parser.add_argument(
    '--pre_crop_height', default=288, type=common.positive_int,
    help='Height used to resize a loaded image. This is ignored when no crop '
         'augmentation is applied.')

parser.add_argument(
    '--pre_crop_width', default=144, type=common.positive_int,
    help='Width used to resize a loaded image. This is ignored when no crop '
         'augmentation is applied.')
# TODO end

parser.add_argument(
    '--loading_threads', default=8, type=common.positive_int,
    help='Number of threads used for parallel loading.')

parser.add_argument(
    '--margin', default='soft', type=common.float_or_string,
    help='What margin to use: a float value for hard-margin, "soft" for '
         'soft-margin, or no margin if "none".')

parser.add_argument(
    '--metric', default='euclidean', choices=loss.cdist.supported_metrics,
    help='Which metric to use for the distance between embeddings.')

parser.add_argument(
    '--loss', default='batch_hard', choices=loss.LOSS_CHOICES.keys(),
    help='Enable the super-mega-advanced top-secret sampling stabilizer.')

parser.add_argument(
    '--learning_rate', default=3e-4, type=common.positive_float,
    help='The initial value of the learning-rate, before it kicks in.')

parser.add_argument(
    '--lr_decay_factor', default=0.96, type=common.positive_float,
    help='Learning rate decay factor')

parser.add_argument(
    '--lr_decay_steps', default=4000, type=common.positive_int,
    help='Learning rate decay steps')

parser.add_argument(
    '--train_iterations', default=25000, type=common.positive_int,
    help='Number of training iterations.')

parser.add_argument(
    '--decay_start_iteration', default=15000, type=int,
    help='At which iteration the learning-rate decay should kick-in.'
         'Set to -1 to disable decay completely.')

parser.add_argument(
    '--weight_decay_factor', default=0.001, type=common.positive_float,
    help='Weight decay factor')

parser.add_argument(
    '--checkpoint_frequency', default=1000, type=common.nonnegative_int,
    help='After how many iterations a checkpoint is stored. Set this to 0 to '
         'disable intermediate storing. This will result in only one final '
         'checkpoint.')

parser.add_argument(
    '--flip_augment', action='store_true', default=False,
    help='When this flag is provided, flip augmentation is performed.')

parser.add_argument(
    '--crop_augment', action='store_true', default=False,
    help='When this flag is provided, crop augmentation is performed. Based on'
         'The `crop_height` and `crop_width` parameters. Changing this flag '
         'thus likely changes the network input size!')

parser.add_argument(
    '--detailed_logs', action='store_true', default=False,
    help='Store very detailed logs of the training in addition to TensorBoard'
         ' summaries. These are mem-mapped numpy files containing the'
         ' embeddings, losses and FIDs seen in each batch during training.'
         ' Everything can be re-constructed and analyzed that way.')

def sample_k_fids_for_pid(pid, all_fids, all_pids, batch_k):
    """ Given a PID, select K FIDs of that specific PID. """
    possible_fids = tf.boolean_mask(all_fids, tf.equal(all_pids, pid))

    # The following simply uses a subset of K of the possible FIDs
    # if more than, or exactly K are available. Otherwise, we first
    # create a padded list of indices which contain a multiple of the
    # original FID count such that all of them will be sampled equally likely.
    count = tf.shape(possible_fids)[0]
    padded_count = tf.cast(tf.ceil(batch_k / count), tf.int32) * count
    full_range = tf.mod(tf.range(padded_count), count)

    # Sampling is always performed by shuffling and taking the first k.
    shuffled = tf.random_shuffle(full_range)
    selected_fids = tf.gather(possible_fids, shuffled[:batch_k])

    return selected_fids, tf.fill([batch_k], pid)

def sample_k_fids_for_model(model, all_fids, all_pids, all_models, all_colors, batch_k, batch_p):
    """ Given a PID, select K FIDs of that specific PID. """
    # possible_fids = tf.boolean_mask(all_fids, tf.equal(all_models, model))
    # possible_pids = tf.boolean_mask(all_pids, tf.equal(all_models, tf.cast(model, tf.int32)))
    possible_pids = tf.boolean_mask(all_pids, tf.equal(all_models, model))
    unique_pids, _ = tf.unique(possible_pids)

    # The following simply uses a subset of K of the possible FIDs
    # if more than, or exactly K are available. Otherwise, we first
    # create a padded list of indices which contain a multiple of the
    # original FID count such that all of them will be sampled equally likely.
    count = tf.shape(unique_pids)[0]
    padded_count = tf.cast(tf.ceil(batch_p / count), tf.int32) * count
    full_range = tf.mod(tf.range(padded_count), count)

    # Sampling is always performed by shuffling and taking the first k.
    shuffled = tf.random_shuffle(full_range)
    selected_pids = tf.gather(unique_pids, shuffled[:batch_p])

    selected_fids, selected_pids = tf.map_fn(lambda pid:sample_k_fids_for_pid(pid,
        all_fids=all_fids, all_pids=all_pids, batch_k=batch_k),
        selected_pids, dtype=(tf.string, tf.string))

    selected_fids = tf.reshape(selected_fids, [-1])
    selected_pids = tf.reshape(selected_pids, [-1])
    # selected_models = tf.fill([tf.shape(selected_fids)[0]], model)
    selected_models = tf.fill([batch_k * batch_p], model)
    # selected_models = tf.map_fn(lambda fid:tf.boolean_mask(all_models, tf.equal(all_fids, fid)), selected_fids)
    selected_colors = tf.map_fn(lambda fid:tf.boolean_mask(all_colors, tf.equal(all_fids, fid)), selected_fids)
    # print(selected_fids, selected_pids, selected_models, selected_colors)

    # return selected_fids, tf.fill([batch_k], pid)
    return selected_models, selected_colors, selected_fids, selected_pids

def main():
    args = parser.parse_args()

    # We store all arguments in a json file. This has two advantages:
    # 1. We can always get back and see what exactly that experiment was
    # 2. We can resume an experiment as-is without needing to remember all flags.
    args_file = os.path.join(args.experiment_root, 'args.json')
    if args.resume:
        if not os.path.isfile(args_file):
            raise IOError('`args.json` not found in {}'.format(args_file))

        print('Loading args from {}.'.format(args_file))
        with open(args_file, 'r') as f:
            args_resumed = json.load(f)
        args_resumed['resume'] = True  # This would be overwritten.

        # When resuming, we not only want to populate the args object with the
        # values from the file, but we also want to check for some possible
        # conflicts between loaded and given arguments.
        for key, value in args.__dict__.items():
            if key in args_resumed:
                resumed_value = args_resumed[key]
                if resumed_value != value:
                    print('Warning: For the argument `{}` we are using the'
                          ' loaded value `{}`. The provided value was `{}`'
                          '.'.format(key, resumed_value, value))
                    args.__dict__[key] = resumed_value
            else:
                print('Warning: A new argument was added since the last run:'
                      ' `{}`. Using the new value: `{}`.'.format(key, value))

    else:
        # If the experiment directory exists already, we bail in fear.
        if os.path.exists(args.experiment_root):
            if os.listdir(args.experiment_root):
                print('The directory {} already exists and is not empty.'
                      ' If you want to resume training, append --resume to'
                      ' your call.'.format(args.experiment_root))
                exit(1)
        else:
            os.makedirs(args.experiment_root)

        # Store the passed arguments for later resuming and grepping in a nice
        # and readable format.
        with open(args_file, 'w') as f:
            json.dump(vars(args), f, ensure_ascii=False, indent=2, sort_keys=True)

    log_file = os.path.join(args.experiment_root, "train")
    logging.config.dictConfig(common.get_logging_dict(log_file))
    log = logging.getLogger('train')

    # Also show all parameter values at the start, for ease of reading logs.
    log.info('Training using the following parameters:')
    for key, value in sorted(vars(args).items()):
        log.info('{}: {}'.format(key, value))

    # Check them here, so they are not required when --resume-ing.
    if not args.train_set:
        parser.print_help()
        log.error("You did not specify the `train_set` argument!")
        sys.exit(1)
    if not args.image_root:
        parser.print_help()
        log.error("You did not specify the required `image_root` argument!")
        sys.exit(1)

    # Load the data from the CSV file.
    pids, fids, models, colors = common.load_dataset_cls(args.train_set, args.image_root)
    max_fid_len = max(map(len, fids))  # We'll need this later for logfiles.

    # Setup a tf.Dataset where one "epoch" loops over all PIDS.
    # PIDS are shuffled after every epoch and continue indefinitely.
    '''
    train_block = {}
    for model_idx in range(len(fids)):
        if models[model_idx] not in train_block.keys():
            train_block[models[model_idx]] = {pids[model_idx]:fids[model_idx]}
        elif pids[model_idx] not in train_block[models[model_idx]].keys():
            train_block[models[model_idx]][pids[model_idx]] = [fids[model_idx]]
        else:
            train_block[models[model_idx]][pids[model_idx]].append(fids[model_idx])
    print(train_block)
    '''
    # unique_pids = np.unique(pids)
    unique_models = np.unique(models)
    dataset = tf.data.Dataset.from_tensor_slices(unique_models)
    dataset = dataset.shuffle(len(unique_models))

    # Constrain the dataset size to a multiple of the batch-size, so that
    # we don't get overlap at the end of each epoch.
    # dataset = dataset.take((len(unique_pids) // args.batch_p) * args.batch_p)
    dataset = dataset.repeat(None)  # Repeat forever. Funny way of stating it.

    # For every PID, get K images.
    # dataset = dataset.map(lambda pid: sample_k_fids_for_pid(
    #     pid, all_fids=fids, all_pids=pids, batch_k=args.batch_k))

    dataset = dataset.map(lambda model: sample_k_fids_for_model(
        model, all_fids=fids, all_pids=pids,
        all_models=models, all_colors=colors,
        batch_k=args.batch_k, batch_p=args.batch_p))
    print('flatten {} {}'.format(dataset.output_shapes, dataset.output_types))

    # Ungroup/flatten the batches for easy loading of the files.
    dataset = dataset.apply(tf.contrib.data.unbatch())
    print('flatten {} {}'.format(dataset.output_shapes, dataset.output_types))

    # Convert filenames to actual image tensors.
    net_input_size = (args.net_input_height, args.net_input_width)
    pre_crop_size = (args.pre_crop_height, args.pre_crop_width)
    dataset = dataset.map(
        lambda c_model, c_color, fid, pid: common.fid_to_image_cls(
            c_model, c_color, fid, pid, image_root=args.image_root,
            image_size=pre_crop_size if args.crop_augment else net_input_size),
        num_parallel_calls=args.loading_threads)

    # Augment the data if specified by the arguments.
    if args.flip_augment:
        dataset = dataset.map(
            lambda im, c_model, c_color, fid, pid: (tf.image.random_flip_left_right(im), c_model, c_color, fid, pid))
    if args.crop_augment:
        dataset = dataset.map(
                lambda im, car_model, car_color, fid, pid: (tf.random_crop(im, net_input_size + (3,)), car_model, car_color, fid, pid))

    # Group it back into PK batches.
    batch_size = args.batch_p * args.batch_k
    dataset = dataset.batch(batch_size)

    # Overlap producing and consuming for parallelism.
    dataset = dataset.prefetch(1)

    # Since we repeat the data infinitely, we only need a one-shot iterator.
    images, c_models, c_colors, fids, pids = dataset.make_one_shot_iterator().get_next()

    # Create the model and an embedding head.
    model = import_module('nets.' + args.model_name)
    head = import_module('heads.' + args.head_name)

    def get_model(images, fids, pids):
        global endpoints, body_prefix
        global dists, neg_dists, pos_dists
        global losses, train_top1, prec_at_k, num_active
        # Feed the image through the model. The returned `body_prefix` will be used
        # further down to load the pre-trained weights for all variables with this
        # prefix.
        endpoints, body_prefix = model.endpoints(images, is_training=True)
        with tf.name_scope('head'):
            endpoints = head.head(endpoints, args.embedding_dim, is_training=True)

        # Create the loss in two steps:
        # 1. Compute all pairwise distances according to the specified metric.
        # 2. For each anchor along the first dimension, compute its loss.
        dists = loss.cdist(endpoints['emb'], endpoints['emb'], metric=args.metric)
        dists_origin = loss.cdist(endpoints['emb_joint'], endpoints['emb_joint'], metric=args.metric)
        losses, train_top1, prec_at_k, _, neg_dists, pos_dists = loss.LOSS_CHOICES[args.loss](
            dists, dists_origin, pids, args.margin, batch_precision_at_k=args.batch_k-1)

        # Count the number of active entries, and compute the total batch loss.
        num_active = tf.reduce_sum(tf.cast(tf.greater(losses, 1e-5), tf.float32))
        loss_mean = tf.reduce_mean(losses, keep_dims=True)

        return loss_mean

    loss_mean = make_parallel(get_model,
                            available_gpu_num(),
                            images=images,
                            fids=fids,
                            pids=pids)

    # Some logging for tensorboard.
    tf.summary.histogram('loss_distribution', losses)
    tf.summary.scalar('loss', tf.reduce_mean(loss_mean))
    tf.summary.scalar('batch_top1', train_top1)
    tf.summary.scalar('batch_prec_at_{}'.format(args.batch_k-1), prec_at_k)
    tf.summary.scalar('active_count', num_active)
    tf.summary.histogram('embedding_dists', dists)
    tf.summary.histogram('embedding_pos_dists', pos_dists)
    tf.summary.histogram('embedding_neg_dists', neg_dists)
    tf.summary.histogram('embedding_lengths', tf.norm(endpoints['emb_raw'], axis=1))

    # Create the mem-mapped arrays in which we'll log all training detail in
    # addition to tensorboard, because tensorboard is annoying for detailed
    # inspection and actually discards data in histogram summaries.
    if args.detailed_logs:
        log_embs = lb.create_or_resize_dat(
            os.path.join(args.experiment_root, 'embeddings'),
            dtype=np.float32, shape=(args.train_iterations, batch_size, args.embedding_dim))
        log_loss = lb.create_or_resize_dat(
            os.path.join(args.experiment_root, 'losses'),
            dtype=np.float32, shape=(args.train_iterations, batch_size))
        log_fids = lb.create_or_resize_dat(
            os.path.join(args.experiment_root, 'fids'),
            dtype='S' + str(max_fid_len), shape=(args.train_iterations, batch_size))

    # These are collected here before we add the optimizer, because depending
    # on the optimizer, it might add extra slots, which are also global
    # variables, with the exact same prefix.
    model_variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, body_prefix)

    print('variables {}'.format(tf.global_variables()))

    # Define the optimizer and the learning-rate schedule.
    # Unfortunately, we get NaNs if we don't handle no-decay separately.
    global_step = tf.Variable(0, name='global_step', trainable=False)
    if 0 <= args.decay_start_iteration < args.train_iterations:
        learning_rate = tf.train.exponential_decay(
            args.learning_rate,
            tf.maximum(0, global_step - args.decay_start_iteration),
            # args.train_iterations - args.decay_start_iteration, args.weight_decay_factor)
            args.lr_decay_steps, args.lr_decay_factor, staircase=True)
    else:
        learning_rate = args.learning_rate
    tf.summary.scalar('learning_rate', learning_rate)
    optimizer = tf.train.AdamOptimizer(learning_rate)
    # Feel free to try others!
    # optimizer = tf.train.AdadeltaOptimizer(learning_rate)

    # Update_ops are used to update batchnorm stats.
    trainable_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, '^(?!.*({}|{})).*$'.format(body_prefix, 'emb'))
    # trainable_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, '^(?!{}).*$'.format(body_prefix))
    print('trainable variables {}'.format(trainable_variables))
    with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
        train_op = optimizer.minimize(loss_mean,
                                    global_step=global_step,
                                    colocate_gradients_with_ops=True,
                                    var_list=trainable_variables)

    # Define a saver for the complete model.
    checkpoint_saver = tf.train.Saver(max_to_keep=0)

    with tf.Session(config=config) as sess:
        if args.resume:
            # In case we're resuming, simply load the full checkpoint to init.
            last_checkpoint = tf.train.latest_checkpoint(args.experiment_root)
            log.info('Restoring from checkpoint: {}'.format(last_checkpoint))
            checkpoint_saver.restore(sess, last_checkpoint)
        else:
            # But if we're starting from scratch, we may need to load some
            # variables from the pre-trained weights, and random init others.
            sess.run(tf.global_variables_initializer())
            if args.initial_checkpoint is not None:
                saver = tf.train.Saver(model_variables)
                # saver = tf.train.Saver()
                saver.restore(sess, args.initial_checkpoint)

            # In any case, we also store this initialization as a checkpoint,
            # such that we could run exactly reproduceable experiments.
            checkpoint_saver.save(sess, os.path.join(
                args.experiment_root, 'checkpoint'), global_step=0)

        merged_summary = tf.summary.merge_all()
        summary_writer = tf.summary.FileWriter(args.experiment_root, sess.graph)

        start_step = sess.run(global_step)
        log.info('Starting training from iteration {}.'.format(start_step))

        # Finally, here comes the main-loop. This `Uninterrupt` is a handy
        # utility such that an iteration still finishes on Ctrl+C and we can
        # stop the training cleanly.
        with lb.Uninterrupt(sigs=[SIGINT, SIGTERM], verbose=True) as u:
            for i in range(start_step, args.train_iterations):

                # Compute gradients, update weights, store logs!
                start_time = time.time()
                _, summary, step, b_prec_at_k, b_embs, b_loss, b_fids, b_models = \
                    sess.run([train_op, merged_summary, global_step,
                              prec_at_k, endpoints['emb'], losses, fids, c_models])
                elapsed_time = time.time() - start_time
                # log.info('models in a batch {}'.format(b_models))

                # Compute the iteration speed and add it to the summary.
                # We did observe some weird spikes that we couldn't track down.
                summary2 = tf.Summary()
                summary2.value.add(tag='secs_per_iter', simple_value=elapsed_time)
                summary_writer.add_summary(summary2, step)
                summary_writer.add_summary(summary, step)

                if args.detailed_logs:
                    log_embs[i], log_loss[i], log_fids[i] = b_embs, b_loss, b_fids

                # Do a huge print out of the current progress.
                seconds_todo = (args.train_iterations - step) * elapsed_time
                log.info('iter:{:6d}, loss min|avg|max: {:.3f}|{:.3f}|{:6.3f}, '
                         'batch-p@{}: {:.2%}, ETA: {} ({:.2f}s/it)'.format(
                             step,
                             float(np.min(b_loss)),
                             float(np.mean(b_loss)),
                             float(np.max(b_loss)),
                             args.batch_k-1, float(b_prec_at_k),
                             timedelta(seconds=int(seconds_todo)),
                             elapsed_time))
                sys.stdout.flush()
                sys.stderr.flush()

                # Save a checkpoint of training every so often.
                if (args.checkpoint_frequency > 0 and
                        step % args.checkpoint_frequency == 0):
                    checkpoint_saver.save(sess, os.path.join(
                        args.experiment_root, 'checkpoint'), global_step=step)

                # Stop the main-loop at the end of the step, if requested.
                if u.interrupted:
                    log.info("Interrupted on request!")
                    break

        # Store one final checkpoint. This might be redundant, but it is crucial
        # in case intermediate storing was disabled and it saves a checkpoint
        # when the process was interrupted.
        checkpoint_saver.save(sess, os.path.join(
            args.experiment_root, 'checkpoint'), global_step=step)

if __name__ == '__main__':
    main()


'''
   Copyright 2020 Matthew Ralston

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

'''

import sys
import os
import copy
import gzip
import yaml
import numpy as np
#import jsonschema
sys.path.append('..')

from kmerdb import fileutil

import logging
logger = logging.getLogger(__file__)

def open(filepath, mode="r", k=None, idx=None):
    if type(filepath) is not str:
        raise TypeError("kmerdb.index.open expects a str as its first positional argument")
    elif type(mode) is not str:
        raise TypeError("kmerdb.index.open expects a str as its second positional argument")
    elif "w" in mode and type(idx) is not np.ndarray:
        if "x" in mode and type(k) is not int: # Specifically, if its not a
            raise TypeError("kmerdb.index.open in create mode expects an int as its first *args argument")

    elif "w" in mode and type(k) is not int:
        raise TypeError("kmerdb.index.open in write mode expects an int as its second *args argument")
    modes = set(list(mode.lower()))
    if modes - set("xrwbt") or len(mode) > len(modes):
        raise ValueError("invalid mode: {}".format(mode))

    creating = "x" in modes
    reading  = "r" in modes
    writing  = "w" in modes
    binary   = "b" in modes
    text     = "t" in modes

    if text and binary:
        raise ValueError("can't have text and binary mode simultaneously")
    elif not (creating or reading or writing):
        raise ValueError("must have exactly one or read/write")

    if "r" in mode.lower():
        return IndexReader(filepath) # Will attempt to read this as an index file
    elif "x" in mode.lower() or set(list("xw")).intersection(modes):
        logger.info("Building index from file and constructing (mode='x') new .kdbi file")
        return IndexBuilder(filepath, k) # Creates a new .kdbi exactly named for the kdb file
    elif "w" in mode.lower():
        logger.info("Writing index to file instead of building index")
        write_index(idx, indexfile, k) # index, indexfile, k

class IndexReader:
    def __init__(self, indexfile:str):
        if type(indexfile) is not str:
            raise TypeError("kmerdb.index.IndexReader expects a str as its first positional argument")
        elif not os.path.exists(indexfile):
            raise TypeError("kmerdb.index.IndexReader expects an existing .kdbi index file as its first positional argument")
        elif not is_gz_file(indexfile):
            raise IOError("kmerdb.index.IndexReader expects a gzip compressed index file as its first positional argument")

        idx = np.array([], dtype="int64")
        #i = 0
        with gzip.open(indexfile, 'rt') as ifile:
            first_line = ifile.readline().rstrip()
            try:
                logger.debug("First line of index:\n{0}".format(first_line))
                k = int(first_line)
                idx = np.zeros(4**k, dtype="int64")
            except ValueError as e:
                logger.debug(e)
                raise IOError("Could not determine k by casting the first line of the index as an integer. Improperly formatted index")
            i = 0
            for line in ifile:
                # if i == 0:
                #     logger.warning("i: {0}, line:\n{1}".format(i, line))

                #     sys.exit(1)
                #     pass
                # else:
                kmer_id, line_offset = [int(i) for i in line.rstrip().split("\t")]
                idx[kmer_id] = line_offset
                i += 1
        self.index = idx

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        del self.index
        return

class IndexBuilder:
    def __init__(self, kdbfile:str, k:int):
        if type(kdbfile) is not str:
            raise TypeError("kmerdb.index.IndexBuilder expects a str as its first positional argument")
        elif not os.path.exists(kdbfile):
            logger.error("Could not find {0}".format(kdbfile))
            raise IOError("kmderdb.index.IndexBuilder expects an existing .kdb file as its first positional argument")
        elif os.path.splitext(kdbfile)[-1] != ".kdb":
            raise IOError("kmerdb.index.IndexBuilder was passed a non-kdb filetype as its first positional argument, should be .kdb")
        elif type(k) is not int:
            raise TypeError("kmerdb.index.IndexBuilder expects an int as its second positional argument")
        N = 4**k
        line_index = np.zeros(N, dtype="int64")
        #line_index = np.array([], dtype="int64")
        #line_index = set() # Can only use set transiently as it might sort
        indexfile = kdbfile + "i"

        with fileutil.open(kdbfile, mode='r') as kdbrdr:
            pos = kdbrdr.tell() # This could be the byte encoding, since we're reading it as plain text

            logger.info("Reading the .kdb file '{0}' in mode 'r' to generate the index...".format(kdbrdr._filepath))
            logger.info("Top position in the file after metadata parsing is said to be 'with this as kdbrdr:' {0}".format(pos))
            logger.info("KDBReader logs the position in mode 'r' as 'header_offset', after appending blocks: {0}".format(kdbrdr.metadata["header_offset"])) # This could be a compressed offset
            logger.info("Asked for the position one more time in mode 'r' in 'with this' and received: {0}".format(kdbrdr.tell()))
            line_index[0] = kdbrdr.tell()

            # SOmething should be checked here
            #kdbrdr.seek(kdbrdr.metadata["header_offset"])
            #kdbrdr._load_block()
            i = 1
            this_line = line_index[0]
            logger.debug("0th index, points to the first row? : {0}".format(line_index[0]))
            for line in kdbrdr:
                try:
                    next_line = kdbrdr.tell()
                    kmer_id, count = [int(i) for i in line.rstrip().split("\t")[:-1]]
                    line_index[kmer_id] = this_line
                    kdbrdr.seek(this_line)
                    new_line = kdbrdr.readline()
                    if new_line == line:
                        logger.debug("i: {0}, K-mer id: {1}, count: {2}, index line: {3}, offset: {4}".format(i, kmer_id, count, i, this_line))
                    elif next_line == 0 or this_line == 0:
                        logger.warning("i: 0, k-mer id: {1}, this_line: {2}, next_line{3}".format(i, kmer_id, this_line, next_line))
                        raise ValueError("Encountered a line where kdbrdr.tell() is producing 0")
                    else:
                        logger.warning("This line was not reproducibly generated...")
                        logger.warning("i: {0}, K-mer id: {1}, offset: {2}".format(i, kmer_id, this_line))
                        sys.exit(1)

                    this_line = next_line
                    if kdbrdr.tell() == 0:
                        raise IOError("kmerdb.index.IndexBuilder expects the kdbrdr to increment with each line")
                    # elif i > 10:
                    #     logger.error("Position in file: {0}".format(pos))
                    #     sys.exit(1)
                except IndexError as e: # The final
                    logger.warning("{0}:\t{1}".format(i, line.rstrip()))
                    raise e # STILL DONT KNOW WHATS HAPPENING HERE
                i += 1
        write_index(line_index, indexfile, k)

def has_index(kdbfile):
    if type(kdbfile) is not str:
        raise TypeError("kmerdb.index.has_index expects a str as its first positional argument")
    elif not os.path.exists(kdbfile):
        raise IOError("kmerdb.index.has_index received a filename that does not exist")
    elif os.path.splitext(kdbfile)[-1] == ".kdb" and not os.path.exists(kdbfile + "i"):
        return False
    elif os.path.splitext(kdbfile)[-1] == ".kdb" and os.path.exists(kdbfile + "i"):
        return True
    elif os.path.splitext(kdbfile)[-1] != ".kdb":
        raise ValueError("unexpected file extension, not a .kdb file.")
    else:
        raise RuntimeError("kmerdb.index.has_index encountered unexpected user inputs:\nkdbfile: '{0}', exists: {1}".format(kdbfile, os.path.exists(kdbfile)))

def read_line(kdbrdr:fileutil.KDBReader, kdbidx:IndexReader, kmer_id):
    if not isinstance(kdbrdr, fileutil.KDBReader):
        raise TypeError("kmerdb.index.read_line expects a kmerdb.fileutil.KDBReader as its first positional argument")
    elif not isinstance(kdbidx, IndexReader):
        raise TypeError("kmerdb.index.read_line expects a kmerdb.index.IndexReader as its second positional argument")
    elif type(kmer_id) is not int:
        raise TypeError("kmerdb.index.read_line expects an int as its third positional argument")
    if kdbidx.index[kmer_id] == 0:
        logger.warning("Found an index value without an offset, in fact, it's in the header")
        return None, None, None
        #raise ValueError("Invalid index value 0")
    else:
        kdbrdr.seek(kdbidx.index[kmer_id])
        kmer_id, count, metadata = kdbrdr.read_line()
        return kmer_id, count, metadata

def is_gz_file(filepath):
    try:
        gzip

        with gzip.open(filepath, 'rb') as f:
            f.readline()
        return True
    except OSError as e:
        return False

def write_index(index:np.array, indexfile:str, k:int):
    if not isinstance(index, np.ndarray):
        raise TypeError("kmerdb.index.write_index expects a str as its first positional argument")
    elif type(indexfile) is not str:
        raise TypeError("kmerdb.index.write_index expects a str as its second positional argument")
    elif os.path.exists(indexfile):
        logger.warning("kmerdb.index.write_index is overwriting an existing index '{0}'...".format(indexfile))
        mode = "wt"
    elif type(k) is not int:
        raise TypeError("kmerdb.index.write_index expects an int as its third positional argument")
    else:
        mode = "xt"
    with gzip.open(indexfile, mode) as ofile:
        ofile.write("{0}\n".format(k))
        for i, line_offset in enumerate(index):
            if int(line_offset) == 0:
                logger.debug("K-mer {0} was not observed, not writing to file".format(i))
                pass
            else:
                ofile.write("{0}\t{1}\n".format(i, line_offset))


import numpy as np

# [samplestore] is a sub-project of [graph-learn](https://github.com/alibaba/graph-learn),
# which is for preprocessing graph data, e.g., the user-item bipartite click graph of a
# recommender system. However, the samplestore API is now deprecated, and its functionality
# has been integrated into graph-learn. Please consider remove the samplestore related code
# and use your own way to prepare the data at your convenience.
import samplestore as ss
import tensorflow as tf

class GraphInput(object):
    def __init__(self, FLAGS=None):
        print('u-i graph input')
        self.FLAGS = FLAGS

        time_bucket = self.FLAGS.time_buckets
        self.time_bucket = np.array(time_bucket.split(","), dtype=np.int32)

        self.features = self.init_input(FLAGS)

    def init_server(self):
        batch_size = self.FLAGS.batch_size

        edge_path = self.FLAGS.tables.split(',')[0]

        # read the user-item click graph (ui_src = User-Item data SouRCe)
        self.ui_src = ss.GraphSource(path=edge_path, alias="u-i")

        worker_count = len(self.FLAGS.worker_hosts.split(","))

        # allow a non-chief worker of the distributed cluster to restart if it fails
        ss.set_enable_failover(True)

        # padding with zero when the number of neighbors is less than the specified number
        ss.set_default_neighbor_id(0)

        # ss.set_batch_num_for_shuffle(num)
        # ss.set_datainit_threadnum(32)
        # ss.set_datainit_batchsize(2048*100)

        # start the samplestore service
        self.server = ss.Server(server_id=self.FLAGS.task_index, server_count=worker_count)
        print('server.start')
        self.server.start()

        # init data from source
        print('server.init')
        self.server.init(graph_source=[self.ui_src])

        print('finish server.init')

        self.ui_graph = ss.Graph(self.ui_src)

        # iterator for sampling a random set of users (u_node_iter = user node iterator)
        self.u_node_iter = ss.IterateNodeSampler(batch_size=batch_size)

        # iterator for sampling a random set of user-item edges (user-item clicks)
        self.edge_sampler = ss.IterateEdgeSampler(batch_size=batch_size, shuffle=True)

        if self.FLAGS.mode == 'train':
            # This samplers will return a list of neighbors of a given user, i.e., the items clicked by the user.
            # The list will be of size self.FLASG.hist_max, self.FLASG.hist_max = 20 by default.
            # The sampled clicks in the list are the consecutive clicks made by the user, i.e., they are in a
            # sliding window of size self.FLASG.hist_max.
            self.hop_u_sampler = ss.WindowNeighborSampler(high_size=0, low_size=self.FLAGS.hist_max)
        else:
            # This samplers will return a list of neighbors of a given user, i.e., the items clicked by the user.
            # The clicks in the list are the latest self.FLASG.hist_max (20 by default) clicks made by the user.
            self.hop_u_sampler = ss.WindowNeighborSampler(high_size=self.FLAGS.hist_max + 1, low_size=0)

    def init_input(self, FLAGS):
        self.u_ids = tf.placeholder(shape=[None], dtype=tf.string)  # [B], batch of user_ids
        self.i_ids = tf.placeholder(shape=[None], dtype=tf.string)  # [B], batch of positive item_ids
        self.item = tf.placeholder(shape=[None, FLAGS.hist_max], dtype=tf.string)  # [B,T], items clicked by the users
        self.nbr_mask = tf.placeholder(shape=[None], dtype=tf.int64)  # [B], the number of clicks made by each user
        self.samples_mask = tf.placeholder(shape=[None], dtype=tf.int64)  # [B], =0 means to ignore this sample

        # [B,T], the time when the click happens
        self.tb_feats = tf.placeholder(shape=[None, FLAGS.hist_max], dtype=tf.int64)

        print('finish placeholder init.')
        return {
            'uid': self.u_ids,
            'iid': self.i_ids,
            'item': self.item,
            'nbr_mask': self.nbr_mask,
            'samples_mask': self.samples_mask,
            'tb_feats': self.tb_feats
        }

    def get_tb_feats(self, target_weights, neights_weights):
        tb_feats = []
        # compuate time bucket
        for idx in range(np.shape(target_weights)[0]):
            cur_t = target_weights[idx]

            low_idx = idx * self.FLAGS.hist_max
            high_idx = (idx + 1) * self.FLAGS.hist_max

            hist_t = [(cur_t - i) // 60 for i in neights_weights[low_idx:high_idx]]

            hist_t = [np.sum(i >= self.time_bucket) for i in hist_t]

            tb_feats.append(hist_t)

        tb_feats = np.array(tb_feats, dtype=np.int64)

        return tb_feats

    def _next_sample(self, graph):
        # sample edges, each edge is of the form <user_id (src_id), clicked_item_id (dst_id)>
        edges = self.edge_sampler.get(graph, with_attr=False)

        # sample the latest self.FLAGS.hist_max clicks made by the user (src_id) before the specified click (dst_id)
        src_nbrs = self.hop_u_sampler.get([graph], edges.src_ids, edges.dst_ids, with_attr=False)

        # the corresponding timestamp of the click (<src_id, dst_id>)
        target_weights = np.array(graph.get_weights(edges.src_ids, edges.dst_ids, default_weight=0)).flatten()

        # the corresponding timestamp of the latest self.FLAGS.hist_max clicks
        neights_weights = np.array(
            graph.get_weights(np.tile(np.reshape(edges.src_ids, [-1, 1]), [1, self.FLAGS.hist_max]).flatten(),
                              src_nbrs.hop(1).ids.flatten(), default_weight=0)).flatten()

        # print("target_weights", target_weights)
        # print("neights_weights", neights_weights)
        #
        # print("time bucket", self.time_bucket)

        # time bucket features
        tb_feats = self.get_tb_feats(target_weights, neights_weights)

        return src_nbrs, edges, tb_feats

    def _next_user(self, graph):
        nodes = self.u_node_iter.get(graph, with_attr=False)
        src_ids = nodes.ids

        # sample the latest self.FLAGS.hist_max clicks made by the user (src_id)
        src_nbrs = self.hop_u_sampler.get([graph], src_ids, None, with_attr=False)
        nbrs_ids = np.reshape(src_nbrs.hop(1).ids, [-1, self.FLAGS.hist_max + 1])

        lst_ids = np.reshape(nbrs_ids[:, -1], -1)
        front_ids = nbrs_ids[:, :-1]

        nbr_mask = np.reshape(src_nbrs.hop(1).real_nbr_count, [-1]) - 1

        # the corresponding timestamp of the click (<src_id, dst_id>)
        target_weights = np.array(graph.get_weights(src_ids, lst_ids, default_weight=0)).flatten()

        # the corresponding timestamp of the latest self.FLAGS.hist_max clicks
        neights_weights = np.array(
            graph.get_weights(np.tile(np.reshape(src_ids, [-1, 1]), [1, self.FLAGS.hist_max]).flatten(),
                              front_ids.flatten(), default_weight=0)).flatten()

        # time bucket features
        tb_feats = self.get_tb_feats(target_weights, neights_weights)

        return src_ids, front_ids, nbr_mask, tb_feats

    def _next_item(self, graph):
        nodes = self.u_node_iter.get(graph, with_attr=False)
        src_ids = nodes.ids

        return src_ids

    def _feed_next_sample(self, pos_sample_table):
        src_nbrs, edges, tb_feats = self._next_sample(pos_sample_table)

        res = dict()

        src_ids = np.reshape(edges.src_ids, [-1])
        dst_ids = np.reshape(edges.dst_ids, [-1])

        nbr_mask = np.reshape(src_nbrs.hop(1).real_nbr_count, [-1])
        nbr_ids = np.reshape(src_nbrs.hop(1).ids, [-1, self.FLAGS.hist_max])
        samples_mask = np.where(nbr_mask >= 3, 1, 0)
        nbr_mask = nbr_mask * samples_mask
        tb_feats = np.reshape(tb_feats, [-1, self.FLAGS.hist_max])

        # print("---------------input------------------")
        # print("u_ids", src_ids)
        # print("i_ids", dst_ids)
        # print("item", nbr_ids)
        # print("nbr_mask", nbr_mask)
        # print("samples_mask", samples_mask)
        # print("tb_feats", tb_feats)
        # print("--------------------------------------")

        res[self.u_ids] = src_ids.astype(np.str)
        res[self.i_ids] = dst_ids.astype(np.str)
        res[self.item] = nbr_ids.astype(np.str)
        res[self.nbr_mask] = nbr_mask
        res[self.samples_mask] = samples_mask
        res[self.tb_feats] = tb_feats

        return res

    def stop(self):
        self.server.stop()

    def feed_next_item(self):
        src_ids = self._next_item(self.ui_graph)

        res = dict()
        res[self.i_ids] = src_ids.astype(np.str)

        return res

    def feed_next_user(self):
        src_ids, nbr_ids, nbr_mask, tb_feats = self._next_user(self.ui_graph)

        # print(tb_feats)
        res = dict()
        res[self.u_ids] = src_ids.astype(np.str)
        res[self.nbr_mask] = nbr_mask
        res[self.item] = nbr_ids.astype(np.str)
        res[self.tb_feats] = tb_feats

        return res

    def feed_next_sample_train(self):
        return self._feed_next_sample(self.ui_graph)

    def feed_next_sample_eval(self):
        return self._feed_next_sample(self.ui_graph_eval)


import numpy as np
import itertools
import json
import os
import time  # sleep() sekund pocakat, #clear klera
import random

# izbolsave ##########################
# dodaj edit team data
# dodaj se en dict z clasi da pol mas moznost optimize for dps  recimo pol gleda da majo dps mirage

######################################
# progam data

RoleList = ["solokite", "BT", "lamp", "hfb", "healscg", "cTank", "druid", "epi", "alaren", "bs", "plyon", "hk", "offchrono"]
clear = lambda: os.system('cls')

###################################

def remove_none(list):
    '''removes all None objects from a list'''

    newlist = []
    for i in list:
        if i:
            newlist.append(i)
    return newlist

def get_team_data():
    '''saves data about the static team including player names and roles'''

    print("input player names\n")
    players = []

    # names
    for i in range(10):
        x = input("player #{} name: ".format(i + 1))
        players.append(x)

    # roles
    print(
        "input player roles, separated by a single space, for each role. Note that code assumes everyone has a dps role available\n")
    print("choose from: solokite, BT, lamp, hfb, healscg, cTank, druid, epi, alaren, bs, plyon, hk, offchrono")
    RoleDict = {}
    for i in players:
        x = input("{}: ".format(i))
        x = x.split(" ")
        RoleDict[i] = x

    with open("TeamData.json", "w") as file:
        json.dump((players, RoleDict), file)
    print("Team data successfully saved to TeamData.json")
    print("Press Enter to return to Menu")
    return 0

def PickRandom(compositions, RelevantRoles, NumberOfRoles):
    RandomComp = compositions[random.randint(0, len(compositions) - 1)]  # picks random composition from list

    roles = [role for role in RelevantRoles] + ["DPS" for i in range(10 - NumberOfRoles)]
    Dict = {}
    for index, role in enumerate(roles):
        try:
            Dict[role].append(RandomComp[index])
        except:
            Dict[role] = [RandomComp[index]]

    return Dict

def MakeComp(perm, RelevantRoles, pylon=False):
    '''Create all possible compositions from list "RelevantRoles".
    Perm is calculated when the program starts and it is list of all permutations made from TeamData

    RelevantRoles: list of all required roles for the encounter. Spelling must match those in RoleList in program data.
    DO NOT INCLUDE DPS in the RelevantRoles'''

    clear()
    print("Calculating...")
    NumberOfRoles = len(RelevantRoles)

    for index, possiblecomp in enumerate(perm):
        for jdex, player in enumerate(possiblecomp):
            if jdex >= NumberOfRoles:
                continue  # try break ?
            elif RelevantRoles[jdex] not in RoleDict[player]:
                perm[index] = None  # isto k false ce klices bool
                continue

    # get rid of Nones
    perm = remove_none(perm.copy())

    # get rid of permutations of dps players
    for index, x in enumerate(perm):
        if x:
            for jdex, y in enumerate(perm[index + 1:]):
                if y:
                    if y[:NumberOfRoles] == x[:NumberOfRoles]:
                        perm[index + 1 + jdex] = None

    # get rid of permutations of pylon players
    if pylon:
        for index, x in enumerate(perm):
            if x:
                for jdex, y in enumerate(perm[index + 1:]):
                    if y:
                        if y[NumberOfRoles:] in list(itertools.permutations(x[NumberOfRoles:])):
                            perm[index + 1 + jdex] = None

    # get rid of Nones
    compositions = remove_none(perm)

    # printing
    clear()

    print("Found", len(compositions), "possible compositions. Save?")
    x = input("y/n: ")
    if x == "y":
        filename = input("enter file name: ")
        with open(str(filename) + ".txt", "w") as file:
            for role in RelevantRoles:
                file.write(str(role)+"\t")
            for i in range(10-NumberOfRoles):
                file.write("DPS" + "\t")
            file.write("\n______________________________________________________________________________\n\n")
            for i in compositions:
                for j in i:
                    file.write(str(j) + "\t")
                file.write("\n")
        print("Data successfully saved to {}.txt".format(filename))
        print("\n")
        time.sleep(3)

    # print out results
    clear()
    print("printing compositions:")
    print("ROLE ODRER:", [role for role in RelevantRoles] + ["DPS" for i in range(10 - NumberOfRoles)])
    for i in compositions:
        print(i)
    print("ROLE ORDER:", [role for role in RelevantRoles] + ["DPS" for i in range(10 - NumberOfRoles)])
    print("\n")

    # make radom comp pick
    pick = "y"
    while (pick == "y"):
        print("Randomly selected composition: ")
        Dict = PickRandom(compositions, RelevantRoles, NumberOfRoles)
        d1 = dict(list(Dict.items())[len(Dict) // 2:])
        d2 = dict(list(Dict.items())[:len(Dict) // 2])
        print("\t".join("{} > {}".format(k, v) for k, v in d2.items()))
        print("\t".join("{} > {}".format(k, v) for k, v in d1.items()))
        print("\nPick again?")
        pick = input("y/n: ")
    print("\n")
    input("Press Enter to return to Menu")

################################################################################
############################################################################
############# main UI

# check if there is already player data and load it
try:
    with open("TeamData.json", "r") as file:
        players, RoleDict = json.load(file)
except:
    input("It looks like this is your first time using this program. Press Enter to input team data")
    get_team_data()
    clear()

###### init
print("Initializing...")
permutacije = list(itertools.permutations(players))

clear()

while (1):  # i dont really know how to make a "menu" like loop so i improvise
    ###### Menu
    print("Menu:")
    print("Current possible commands are:\n")
    print("*premade compositions: standard, boonthief, w4, w6", "w7")  # will add dhuum when i add rr and kite and epi
    print("*Add your own composition: add")
    print("*Close the program: exit")
    print("*edit current team: edit")
    print("*Make a new team: new")
    print("\n")

    x = input("input: ")  # user select an option here
    clear()
    ###########################################possible options:
    if x == "exit":
        break

    elif x == "standard":
        MakeComp(permutacije.copy(), ["cTank", "druid", "bs", "alaren", "hfb"])
    elif x == "boonthief":
        MakeComp(permutacije.copy(), ["BT", "druid", "bs",  "alaren", "healscg"])
    elif x == "w4":
        MakeComp(permutacije.copy(), ["cTank", "druid", "bs", "alaren", "hfb", "hk"])
    elif x == "w6":
        print("there is multiple compositions available for this raid")
        print("press 1 for hfb variant OR press 2 for healscg + offchrono variant")
        x = input("1, 2: ")
        if x == "1":
            MakeComp(permutacije.copy(), ["solokite", "druid", "alaren", "hfb", "cTank", "lamp"])
        elif x == "2":
            MakeComp(permutacije.copy(), ["solokite", "druid", "alaren", "healscg", "cTank", "lamp", "offchrono"])
    elif x == "w7":
        MakeComp(permutacije.copy(), ["cTank", "druid", "bs", "alaren", "hfb", "pylon", "pylon", "pylon"], pylon=True)

    elif x == "new":
        get_team_data()
        print("Initializing...")
        permutacije = list(itertools.permutations(players))

        clear()


from styx_msgs.msg import TrafficLight
import rospy
import tensorflow as tf
import numpy as np
import cv2
import glob
from functools import partial

class TLClassifier(object):
    def __init__(self):
        #TODO load classifier
        self.DETECTION_THRESHOLD = 0.9
        self.COCO_TL_NUM = 10
        MODELCHUNK_BASE_DIR = rospy.get_param('~model_path')
        self.readsize = 1024

        #rospy.logwarn(MODEL_BASE_DIR)
        PATH_TO_MODEL = MODELCHUNK_BASE_DIR + '/rfcn_resnet101_coco_2018_01_28/frozen_inference_graph.pb'
        # reassemble the model from chunks. credit goes to team vulture for this idea
        output = open(PATH_TO_MODEL, 'wb')
        frozen_model_path = (MODELCHUNK_BASE_DIR+'/rfcn_resnet101_coco_2018_01_28/')
        chunks = glob.glob(frozen_model_path + 'chunk*')
        chunks.sort()
        for filename in chunks:
            with open(filename, 'rb') as fileobj:
                for chunk in iter(partial(fileobj.read, self.readsize), ''):
                    output.write(chunk)
        output.close()
        rospy.loginfo("Model recreated to run for TL detection")

        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(PATH_TO_MODEL,'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def,name='')
            self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
            self.d_boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
            self.d_scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
            self.d_classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
            self.num_d = self.detection_graph.get_tensor_by_name('num_detections:0')
        self.sess = tf.Session(graph=self.detection_graph)

    def get_classification(self, image):
        """Determines the color of the traffic light in the image

        Args:
            image (cv::Mat): image containing the traffic light

        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """
        #TODO implement light color prediction
        with self.detection_graph.as_default():
            image_expanded = np.expand_dims(image, axis=0)
            (boxes, scores, classes, num) = self.sess.run([self.d_boxes, self.d_scores, self.d_classes, self.num_d],feed_dict={self.image_tensor: image_expanded})
            #rospy.logwarn(scores[0][0])

        if int(num)==0:
            return TrafficLight.UNKNOWN
        else:
            valid = []
            for i in range(int(num)):
                if int(classes[0][i]) == self.COCO_TL_NUM and scores[0][i] > self.DETECTION_THRESHOLD:
                    valid.append(i)
            if len(valid)==0:
                return TrafficLight.UNKNOWN

            h, w = image.shape[0], image.shape[1]
            #As the scores are already sorted, we take the first box from the valid boxes
            ymin = int(h*boxes[0][valid[0]][0]) #ymin
            xmin = int(w*boxes[0][valid[0]][1]) #xmin
            ymax = int(h*boxes[0][valid[0]][2]) #ymax
            xmax = int(w*boxes[0][valid[0]][3]) #xmax

            # To figure out which area of traffic light is brighter i.e. lighted ON
            cropped_image = image[ymin:ymax, xmin:xmax]

            return self.red_yellow_green(cropped_image)

        #return TrafficLight.UNKNOWN

    def red_yellow_green(self, rgb_image):
        h,w = rgb_image.shape[0], rgb_image.shape[1]
        rgb_image = rgb_image[int(1*h/10):int(9*h/10),int(1*w/10):int(9*w/10)] #as too outer non-traffic-light area affects classification
        h,w = rgb_image.shape[0], rgb_image.shape[1]

        if h>w: # if traffic lights are vertical (topmost red)
            one_half_rgb_image = rgb_image[0:int(h/3),:] # red area
            two_half_rgb_image = rgb_image[int(h/3):int(2*h/3),:] # yellow area
            three_half_rgb_image = rgb_image[int(2*h/3):h,:] # green area
        else: # if traffic lights are horizontal like in Japan (rightmost red)
            one_half_rgb_image = rgb_image[:,int(2*w/3):w] # red area
            two_half_rgb_image = rgb_image[:,int(w/3):int(2*w/3)] # yellow area
            three_half_rgb_image = rgb_image[:,0:int(w/3)] # green area

        # To just consider the centre area of each half for more robust especially red vs. yellow detection.
        h1,w1 = one_half_rgb_image.shape[0], one_half_rgb_image.shape[1]
        one_half_rgb_image = one_half_rgb_image[int(1*h1/4):int(3*h1/4), int(1*w1/4):int(3*w1/4)]
        two_half_rgb_image = two_half_rgb_image[int(1*h1/4):int(3*h1/4), int(1*w1/4):int(3*w1/4)]
        three_half_rgb_image = three_half_rgb_image[int(1*h1/4):int(3*h1/4), int(1*w1/4):int(3*w1/4)]

        #f, (ax1, ax2, ax3) = plt.subplots(3,1)
        #ax1.imshow(one_half_rgb_image)
        #ax2.imshow(two_half_rgb_image)
        #ax3.imshow(three_half_rgb_image)
        #plt.hold(True)
        #f.savefig(os.path.splitext(image_path)[0] + "_cropped.png",bbox_inches='tight')
        #plt.close(f)

        #hsv = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2HSV)
        #sum_bright_full = float(np.sum(hsv[:,:,2]))
        hsv_one_half = cv2.cvtColor(one_half_rgb_image, cv2.COLOR_RGB2HSV)
        hsv_two_half = cv2.cvtColor(two_half_rgb_image, cv2.COLOR_RGB2HSV)
        hsv_three_half = cv2.cvtColor(three_half_rgb_image, cv2.COLOR_RGB2HSV)
        sum_bright_full = float(np.sum(hsv_one_half[:,:,2])+np.sum(hsv_two_half[:,:,2])+np.sum(hsv_three_half[:,:,2]))
        ratio_red = np.sum(hsv_one_half[:,:,2])/sum_bright_full
        ratio_yellow = np.sum(hsv_two_half[:,:,2])/sum_bright_full
        ratio_green = np.sum(hsv_three_half[:,:,2])/sum_bright_full
        rospy.logwarn('0(Red): ' + str(ratio_red) + ', 1(Yellow):' + str(ratio_yellow) + ', 2(Green):' + str(ratio_green))
        # state --> TrafficLight.UNKNOWN=4 , GREEN=2, YELLOW=1, RED=0

        # If the lights are off then it is expected that brightness is similar in all three halfs.
        # But if one of them is highlighted and then experimentally I found that it is always >35% of brigthness vs. overall.
        if ratio_red >= ratio_yellow and ratio_red >= ratio_green and ratio_red >= 0.35:
            return TrafficLight.RED # Red
        elif ratio_green >= ratio_yellow and ratio_green >= 0.35:
            return TrafficLight.GREEN # Yellow
        elif ratio_yellow >= 0.35:
            return TrafficLight.YELLOW # Green
        else:
            return TrafficLight.UNKNOWN # traffic light switched off / could be considered green as well :)


from keras.models import model_from_json,load_model
import numpy as np
import os,pickle
# os.environ["CUDA_VISIBLE_DEVICES"]="-1"
word_index=pickle.load(open("word_index_emotion.pkl","rb"))
# word_index_intent=pickle.load(open("word_index_intent.pkl","rb"))

with open("emotion_model.json","r") as emojson:
	emotion_model_json=emojson.read()
	emotion_loaded_model=model_from_json(emotion_model_json)
	emotion_loaded_model.load_weights("emotion_check_weights_best.h5")

# with open("intent_model.json","r") as intentjson:
# 	intent_model_json=intentjson.read()
# 	intent_loaded_model=model_from_json(intent_model_json)
# 	intent_loaded_model.load_weights("intent_adam.h5")

def preprocess(sentence):
	sentence=sentence.strip()
	sentence=sentence.lower()
	sentence=sentence.replace("\x92","")
	sentence = sentence.replace("_comma_",",")
	sentence=sentence.replace("'s","")
	sentence=sentence.replace("n't"," not")
	sentence=sentence.replace("'t","ot")
	sentence=sentence.replace("'m"," am")
	sentence=sentence.replace("'ve"," have")
	sen = ""
	for i in range(len(sentence)):
		if not sentence[i].isalpha():
			sen+=" "
		else:
			sen+=sentence[i]
	return sen

def complete_preprocessing(text,intent=0):
	l=[]
	text=preprocess(text)
	for i in text.split():
		try:
			if intent:
				l.append(word_index_intent[i])
			else:
				l.append(word_index[i])
		except:
			l.append(0)
	if intent:
		max_length=200
	else:
		max_length=50
	l=[0]*(max_length-min(len(l),max_length))+l[:max_length]
	return np.array(l)

def get_emotion(input_text):
	list_of_common_emotions=[["excited", "surprised", "joyful"],
	["afraid", "terrified", "anxious", "apprehensive"],
	["disgusted", "embarrassed", "guilty","ashamed"],
	["angry", "annoyed","jealous", "furious"],
	[  "faithful", "trusting", "grateful", "caring", "hopeful"],
	["sad" ,"disappointed", "devastated", "lonely" ,"nostalgic", "sentimental"],\
	["proud", "impressed", "content"],
	["anticipating", "prepared","confident"]]

	user_input = complete_preprocessing(input_text)
	# user_input.shape
	print(user_input)
	user_input = user_input.reshape((1,50))
	predicted = emotion_loaded_model.predict(user_input)
	return list_of_common_emotions[np.argmax(predicted)][0]

# def get_intent(input_text):
# 	user_input= complete_preprocessing(input_text,1)
# 	user_input=user_input.reshape((1,200))
# 	output=intent_loaded_model.predict(user_input)
# 	print(output)
# 	intent=np.argmax(output)
# 	return intent

def generate_empathetic_response(u1,u2,u3,emotion):
	u1_preprocessed=preprocess(u1)
	u2_preprocessed=preprocess(u2)
	u3_preprocessed=preprocess(u3)

	text=emotion+" "+u3_preprocessed+" . "+u2_preprocessed+" . "+u1_preprocessed+" ."
	print(text)
	with open("tmp_predict_this.txt","w") as file:
		file.write(text)

def generate_medical_response(input_text):
	user_input= complete_preprocessing(input_text)
	user_input=user_input.reshape((1,50))
	# output=

while True:
	print(get_emotion(input()))


#!/usr/bin/env python

"""Module to deal with trees

Some more documentation

For now (and for some time), this module assumes that in the Trees folder,
there is only one set of tree_bricks (i.e. TreeMaker was run for only one set
of tree_bricks).

TODO:
    - everything
    - pep8

"""
import matplotlib
import numpy as np
#from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
import pandas as pd
from scipy.io import FortranFile as FF
import yt
from glob import glob
import os
import yt.utilities.fortran_utils as fpu
from tqdm import tqdm

from . import physics
from .utils import find_outputs

class Forest(object):
    """
    Read the outputs from TreeMaker and gather in a dataframe
    in self.trees
    units are:
    x,y,z -> Mpccm
    vx,vy,xz -> km/s (to check)
    m,mvir -> 1e11Msun
    r,rvir -> Mpc
    """

    def __init__(self, Galaxy=False, path='.', big_run=False):
        paths = find_outputs(path+'/Outputs')
        self.prefix = path
        yt.funcs.mylog.setLevel(40)
        ds = yt.load(paths[-1])

        _sim = {}
        _sim['h'] = float(ds.cosmology.hubble_constant)
        _sim['Om'] = ds.omega_matter
        _sim['Ol'] = ds.omega_lambda
        _sim['Lbox'] = float(ds.length_unit.in_units('Mpccm'))

        treefolder = os.path.join(self.prefix, 'Trees')
        self.galaxies = False
        if Galaxy:
            treefolder = os.path.join(self.prefix, 'TreeStars')
            self.galaxies = True

        self.treefolder = treefolder
        self.sim = _sim
        self.ds = ds
        self.folder = treefolder
        self.snap = self._get_timestep_number()

        self.read_tree(big_run)

    def read_tree(self, big_run=False):
        """
        """

        if self.galaxies:
            tree_file = os.path.join(self.folder, 'tree.dat')
            self = self._read_tree(tree_file, 'd', big_run)
        else:
            tree_file = glob('{}/tree_file_???.001'.format(self.folder))[0]
            self = self._read_tree(tree_file, 'f', big_run)

        return

    def map_halo_ts_to_output(self, timestep):
        timestep = timestep.astype('int64')
        bricks = sorted(glob(os.path.join(self.treefolder, 'tree_bricks???')))
        mapping = np.empty(timestep.max(), dtype=np.int64)
        for istep in sorted(timestep.unique()):
            #ioutput = int(brick.split('bricks')[1])
            ioutput = int(bricks[istep-1].split('bricks')[1])
            mapping[istep-1] = ioutput

        return mapping[timestep-1]

    def get_all_progenitors(self, hnum, timestep=None):
        """Return the reduced tree containing ONLY the progenitors of halo hid
        ==========
        * hid: ID of the halo at output timestep given by the halo/galaxy finder
        * timestep (None): timestep at which the ID must be taken, default is last timestep
        """

        if timestep is None:
            ts = self.halo_ts.max()
        print('Care, this can be very long for galaxies. The get_main_progenitor is much faster')
        print('It is OK for Haloes')
        progenitors = self._get_progenitors(hnum, timestep=timestep)

        return progenitors

    def get_main_progenitor(self, hnum, timestep=None):
        """Return the reduced tree containing ONLY the main progenitors of halo hid
        ==========
        * hid: ID of the halo at output timestep given by the halo/galaxy finder
        * timestep (None): timestep at which the ID must be taken, default is last timestep
        """
        id_main = {}
        if timestep is None:
            ts = self.trees.halo_ts.max()
        else:
            ts = timestep

        list_progenitors = [(hnum, ts)]
        fatherMass = self.trees.loc[hnum,timestep].fatherMass
        fathersID = self.trees.loc[hnum,timestep].fathersID
        while (fathersID.size > 1):
            #update ts
            ts -= 1
            #remove background
            fatherMass = fatherMass[fathersID > 0]
            fathersID = fathersID[fathersID > 0]
            #main father is the one that contains most mass
            main_father = fathersID[fatherMass.argmax()]
            list_progenitors += [(main_father, ts)]
            fatherMass = self.trees.loc[main_father, ts].fatherMass
            fathersID = self.trees.loc[main_father, ts].fathersID
        main_progs = self.trees.loc[list_progenitors]

        return main_progs

    def get_main_children(self, hnum, timestep=None):
        """Return the reduced tree containing ONLY the main children of halo hid
        ==========
        * hid: ID of the halo at output timestep given by the halo/galaxy finder
        * timestep (None): timestep at which the ID must be taken, default is last timestep
        """
        # Get current timestep
        max_ts = self.trees.halo_ts.max()
        if timestep is None:
            ts = max_ts
        else:
            ts = timestep

        list_sons = [(hnum, ts)]
        while (ts < max_ts):
            main_son = self.trees.loc[hnum,ts].mainSon
            #update ts
            ts += 1
            list_sons += [(main_son, ts)]
        children = self.trees.loc[list_sons]

        return children

    def get_family(self, hnum, timestep=None):
        """Return the reduced tree containing ONLY the main progenitors/children of halo hid
        ==========
        * hid: ID of the halo at output timestep given by the halo/galaxy finder
        * timestep (None): timestep at which the ID must be taken, default is last timestep
        """
        # Get current timestep
        if timestep is None:
            ts = self.trees.halo_ts.max()
        else:
            ts = timestep

        child = self.get_main_children(hnum, timestep)
        fathers = self.get_main_progenitor(hnum, timestep)
        fathers = fathers.loc[fathers.halo_ts != ts]
        family = pd.concat((child, fathers))
        family = family.sort_index(level=1)

        return family

    def plot_all_trees(self, minmass=-1, maxmass=1e99, radius=1.0,
                       output=None, loc='./'):
        """
        See plot_halo_tree.
        Compute the above function for all halos with virial mass
        between `minmass` and `maxmass`. The width in Mpccm for the
        dynamics of the BH is radius. `loc` is the location to save the
        PDF.
        """

        if output is None:
            output = str(minmass) + 'To' + str(maxmass) + '_r' + str(radius)

        # Normalize min/max mass to 1e11
        minmass /= 1.e11
        maxmass /= 1.e11

        # Select halos at the last timestep of the tree, with the right mass
        # Also selects halos with at least 1 progenitor...
        mask = ((self.trees.halo_ts == self.snap) &
                (self.trees.mvir > minmass) &
                (self.trees.mvir < maxmass) &
                (self.trees.first_prog != -1))

        tid = self.trees[mask].halo_id
        print('Total: {} halos'.format(len(tid)))

        pdf = PdfPages(loc + 'trees{}.pdf'.format(output))
        try:
            fig = self.fig
        except AttributeError:
            self.fig = plt.figure(figsize=(12, 12))
        self.fig.savefig(pdf, format='pdf', dpi=200)

        for ihalo in tid:
            self.plot_halo_tree(hid=int(ihalo), radius=radius, pdffile=pdf)
            plt.close()
        pdf.close()
        plt.close('all')

        return 'OK'

    def plot_halo_tree(self, hid=None, hnum=None, hts=None, radius=1.0,
                       pdffile=None, loc='./'):
        """
        Plot Mass/Merger history of a halo
        Parameters
        ----------
        * hid (None): halo_id of the halo you want to have the merger/mass history
        * hnum (None): halo_num of the halo you want to have the
          merger history, if used then the timestep of this halo must
          be given with hts
        * radius (1): width of the window to plot the dynamics of the halo
        * loc ('./'): where to save the plot
        """

        # Define the selected halo
        if hid is None:
            hid = (self.trees.loc[(self.trees.halo_num == hnum) & (
                self.trees.halo_ts == hts)].halo_id).astype(int)
        halo = self.trees.ix[hid]
        sim = self.sim

        xh = halo['x']  # - sim['Lbox']/2 #/sim['Lbox'] + .5
        yh = halo['y']  # - sim['Lbox']/2   #/sim['Lbox'] + .5
        zh = halo['z']  # - sim['Lbox']/2   #/sim['Lbox'] + .5

        # Get progenitors
        progs = self._get_progenitors(hid)
        progs.loc[hid].descendent_id = -1
        main_progs = self.get_main_progenitor(hid)
        main_progs.loc[hid].descendent_id = -1

        # Recenter on selected halo and correct for periodic boundaries
        progs['x'] -= xh
        progs['y'] -= yh
        progs['z'] -= zh
        progs['x'].where(progs.x <= sim['Lbox'] / 2.,
                         progs.x - sim['Lbox'], inplace=True)
        progs['y'].where(progs.y <= sim['Lbox'] / 2.,
                         progs.y - sim['Lbox'], inplace=True)
        progs['z'].where(progs.z <= sim['Lbox'] / 2.,
                         progs.z - sim['Lbox'], inplace=True)
        progs['x'].where(progs.x >= -sim['Lbox'] / 2.,
                         progs.x + sim['Lbox'], inplace=True)
        progs['y'].where(progs.y >= -sim['Lbox'] / 2.,
                         progs.y + sim['Lbox'], inplace=True)
        progs['z'].where(progs.z >= -sim['Lbox'] / 2.,
                         progs.z + sim['Lbox'], inplace=True)
        main_progs['x'] -= xh
        main_progs['y'] -= yh
        main_progs['z'] -= zh
        main_progs['x'].where(main_progs.x <= sim['Lbox'] / 2.,
                              main_progs.x - sim['Lbox'], inplace=True)
        main_progs['y'].where(main_progs.y <= sim['Lbox'] / 2.,
                              main_progs.y - sim['Lbox'], inplace=True)
        main_progs['z'].where(main_progs.z <= sim['Lbox'] / 2.,
                              main_progs.z - sim['Lbox'], inplace=True)
        main_progs['x'].where(main_progs.x >= -sim['Lbox'] / 2.,
                              main_progs.x + sim['Lbox'], inplace=True)
        main_progs['y'].where(main_progs.y >= -sim['Lbox'] / 2.,
                              main_progs.y + sim['Lbox'], inplace=True)
        main_progs['z'].where(main_progs.z >= -sim['Lbox'] / 2.,
                              main_progs.z + sim['Lbox'], inplace=True)

        # Define plot range
        xc = (main_progs['x'].min() + main_progs['x'].max()) * .5
        yc = (main_progs['y'].min() + main_progs['y'].max()) * .5
        zc = (main_progs['z'].min() + main_progs['z'].max()) * .5
        xmin = ymin = zmin = -radius
        xmax = ymax = zmax = radius
        progs[['x', 'y', 'z']] -= [xc, yc, zc]
        main_progs[['x', 'y', 'z']] -= [xc, yc, zc]

        # Recenter other halos, and correct for periodic boundaries
        x = self.trees.x - (xc + xh)
        y = self.trees.y - (yc + yh)
        z = self.trees.z - (zc + zh)
        x.where(x <= sim['Lbox'] / 2., x - sim['Lbox'], inplace=True)
        y.where(y <= sim['Lbox'] / 2., y - sim['Lbox'], inplace=True)
        z.where(z <= sim['Lbox'] / 2., z - sim['Lbox'], inplace=True)
        x.where(x >= -sim['Lbox'] / 2., x + sim['Lbox'], inplace=True)
        y.where(y >= -sim['Lbox'] / 2., y + sim['Lbox'], inplace=True)
        z.where(z >= -sim['Lbox'] / 2., z + sim['Lbox'], inplace=True)

        # Select other halos within plot range
        others = ((x < radius) & (-radius < x) &
                  (y < radius) & (-radius < y) &
                  (z < radius) & (-radius < z))

        # Scalings
        scc = np.log10(progs.m)
        mainscc = np.log10(main_progs.m)
        sc = (scc - scc.min()) / (scc.max() - scc.min()) * 500.
        mainsc = (mainscc - scc.min()) / (mainscc.max() - scc.min()) * 500.

        osc = np.log10(self.trees[others].m)
        osc = (osc - scc.min()) / (scc.max() - scc.min()) * 500.
        ocol = self.trees[others].halo_ts
        ocol = .15 + .7 * (ocol - ocol.min()) / (ocol.max() - ocol.min())
        edg = np.where(self.trees[others].bush_id ==
                       halo['bush_id'], 'orange', 'k')

        # Plot halos dynamics in the comoving space
        # try:
        #    fig = self.fig
        # except AttributeError:
        self.fig = plt.figure(figsize=(12, 12))
        fig = self.fig
        # try:
        #    ax = self.ax
        # except AttributeError:
        ax10 = self.fig.add_axes([0.07, 0.05, 0.4, 0.4])
        ax00 = self.fig.add_axes([0.07, 0.5, 0.4, 0.4])
        ax01 = self.fig.add_axes([0.55, 0.5, 0.4, 0.4])
        ax11 = self.fig.add_axes([0.55, 0.05, 0.4, 0.4])
        self.ax = [ax00, ax01, ax10, ax11]
        ax = self.ax

        self.ax[0].scatter(x[others], y[others], c=ocol, s=osc, cmap='Greys',
                           vmin=0, vmax=1., edgecolor=edg,
                           rasterized=True)
        self.ax[0].scatter(progs.x, progs.y, c=progs.halo_ts, s=sc,
                           cmap='summer', rasterized=True)
        self.ax[0].scatter(main_progs.x, main_progs.y,
                           c=main_progs.halo_ts, s=mainsc,
                           cmap='magma', rasterized=True)
        self.ax[0].set_xlabel(r'$x$ (cMpc)')
        self.ax[0].set_ylabel(r'$y$ (cMpc)')
        self.ax[0].set_xlim(xmin, xmax)
        self.ax[0].set_ylim(ymin, ymax)

        self.ax[1].scatter(z[others], y[others], c=ocol, s=osc, cmap='Greys',
                           vmin=0, vmax=1., edgecolor=edg,
                           rasterized=True)
        self.ax[1].scatter(progs.z, progs.y, c=progs.halo_ts, s=sc,
                           cmap='summer', rasterized=True)
        self.ax[1].scatter(main_progs.z, main_progs.y,
                           c=main_progs.halo_ts, s=mainsc,
                           cmap='magma', rasterized=True)
        self.ax[1].set_xlabel(r'$z$ (cMpc)')
        self.ax[1].set_ylabel(r'$y$ (cMpc)')
        self.ax[1].set_xlim(zmin, zmax)
        self.ax[1].set_ylim(ymin, ymax)

        self.ax[2].scatter(x[others], z[others], c=ocol, s=osc, cmap='Greys',
                           vmin=0, vmax=1., edgecolor=edg,
                           rasterized=True)
        self.ax[2].scatter(progs.x, progs.z, c=progs.halo_ts, s=sc,
                           cmap='summer', rasterized=True)
        self.ax[2].scatter(main_progs.x, main_progs.z,
                           c=main_progs.halo_ts, s=mainsc,
                           cmap='magma', rasterized=True)
        self.ax[2].set_xlabel(r'$x$ (cMpc)')
        self.ax[2].set_ylabel(r'$z$ (cMpc)')
        self.ax[2].set_xlim(xmin, xmax)
        self.ax[2].set_ylim(zmin, zmax)

        # Draw some lines
        for progid in progs.halo_id:
            descid = progs.ix[progid].descendent_id
            if descid > 0:
                progenitor = progs.ix[progid]
                descendent = progs.ix[descid]
                px, py, pz = progenitor[['x', 'y', 'z']]
                dx, dy, dz = descendent[['x', 'y', 'z']]
                l0 = self.ax[0].plot([px, dx], [py, dy],
                                     lw=1, c='k', zorder=-1)
                l1 = self.ax[1].plot([pz, dz], [py, dy], lw=1, c='k',
                                     zorder=-1)
                l2 = self.ax[2].plot([px, dx], [pz, dz], lw=1, c='k',
                                     zorder=-1)

        title = "Halo #{t:d}\n(x, y, z) = ({x:3.3f}, {y:3.3f}, {z:3.3f})"
        title = title.format(t=int(halo.halo_num),
                             x=xh / sim['Lbox'] + .5,
                             y=yh / sim['Lbox'] + .5,
                             z=zh / sim['Lbox'] + .5)
        suptitle = self.fig.suptitle(title, fontsize=18)

        if not self.fig.texts:
            self.fig.texts.append(suptitle)
        if pdffile:
            self.fig.savefig(pdffile, format='pdf', dpi=200)
        else:
            plt.savefig(loc + 'halo_{:d}dynamics.png'.format(int(halo.halo_num)), dpi=100, format='png')

        # A bit of cleaning
        # for axx in self.ax:
        #    axx.clear()
        #self.fig.texts = []
        fig.clf()
        # try:
        #    fig = self.fig
        # except AttributeError:
        #    self.fig = plt.figure(figsize=(12, 12))
        # try:
        #    ax = self.ax
        # except AttributeError:
        ax10 = self.fig.add_axes([0.07, 0.05, 0.4, 0.4])
        ax00 = self.fig.add_axes([0.07, 0.5, 0.4, 0.4])
        ax01 = self.fig.add_axes([0.55, 0.5, 0.4, 0.4])
        ax11 = self.fig.add_axes([0.55, 0.05, 0.4, 0.4])
        self.ax = [ax00, ax01, ax10, ax11]
        ax = self.ax

        # Plot halo mass/merger history

        time = np.array([self.timestep['age'][int(hts) - 1]
                         for hts in progs.halo_ts])
        maintime = np.array([self.timestep['age'][int(hts) - 1]
                             for hts in main_progs.halo_ts])
        self.ax[0].scatter(time, progs.m * 1e11, c=progs.halo_ts,
                           s=sc, cmap='summer', rasterized=True)
        self.ax[0].semilogy(
            [self.ds.cosmology.t_from_z(z) / (3600 * 24 * 365 * 1e9)
             for z in np.logspace(-10, 2, 1000)],
            1e12 * physics.MofZ(halo.m * 1e11 / 1e12,
                                np.logspace(-10, 2, 1000),
                                1 / self.timestep['aexp'][progs.loc[hid].halo_ts - 1] - 1),
            'k')
        self.ax[0].scatter(maintime, main_progs.m * 1e11,
                           c=main_progs.halo_ts, s=mainsc,
                           cmap='magma', rasterized=True)
        self.ax[0].set_xlabel(r'$t$ (Gyr)')
        self.ax[0].set_ylabel(r'$M$ ($\mathrm{M}_{\odot}$)')
        self.ax[0].set_ylim(1e11 * progs.m.min() / 3.,
                            1e11 * progs.m.max() * 3.)
        self.ax[0].set_xlim(time.min(), time.max())
        self.ax[0].set_yscale('log')
        z = np.unique([int((1 / self.timestep['aexp'][int(hts) - 1] - 1) * 2) / 2.
                       for hts in main_progs.halo_ts])
        z = z[1:-1]
        ax_z = self.ax[0].twiny()
        ax_z.set_xlim(self.ax[0].get_xlim())
        ax_z.set_xticks([float(self.ds.cosmology.t_from_z(
            zz) / (1e9 * 365 * 24 * 3600)) for zz in z])
        ax_z.set_xticklabels(z)
        ax_z.set_xlabel("redshift")

        main_prog = self.get_main_progenitor(hid)
        minor_mergers = []
        major_mergers = []
        for current_id in main_prog.halo_id:
            if main_prog.halo_ts[current_id] != main_prog.halo_ts.min():
                halo_ts = main_prog.halo_ts[current_id]
                progs = self.get_all_progenitors(
                    current_id, timestep=halo_ts - 1)
                mask_minor = ((progs.m / main_prog.m[current_id + 1] > 1. / 20) &
                              (progs.m / main_prog.m[current_id + 1] < 1. / 4))
                minor_mergers += [len(progs.m[mask_minor])]
                mask_major = (progs.m / main_prog.m[current_id + 1] >= 1. / 4)
                major_mergers += [len(progs.m[mask_major]) - 1]

        scc = np.log10(main_prog.m)
        sc = (scc - scc.min()) / (scc.max() - scc.min()) * 500.
        time = np.array([self.timestep['age'][hts - 1]
                         for hts in main_prog[main_prog.halo_ts !=
                                              self.trees.halo_ts[hid]].halo_ts])
        self.ax[2].plot(time, # np.cumsum(major_mergers)[-1] -
                        np.cumsum(major_mergers),
                        color='b', label='1:1 > mass ratio > 1:4 (sim)')

        self.ax[2].plot(time,
                        [physics.N_merger_until_z(
                            0.25,
                            halo.m * 1e11 / 1e12,
                            self.ds.cosmology.z_from_t(t * 1e9 * 365 * 24 * 3600),
                            self.ds.cosmology.z_from_t(time.max() * 1e9 * 365 * 24 * 3600))
                         for t in time],
                        color='b', linestyle='--',
                        label='1:1 > mass ratio > 1:4 (theory)')

        self.ax[2].plot(time, # np.cumsum(minor_mergers)[-1] -
                        np.cumsum(minor_mergers), color='r',
                        label='1:4 > mass ratio > 1:20 (sim)')

        self.ax[2].plot(time,
                        [physics.N_merger_until_z(
                            1. / 20,
                            halo.m * 1e11 / 1e12,
                            self.ds.cosmology.z_from_t(t * 1e9 * 365 * 24 * 3600),
                            self.ds.cosmology.z_from_t(time.max() * 1e9 * 365 * 24 * 3600))
                         - physics.N_merger_until_z(
                             1. / 4,
                             halo.m * 1e11 / 1e12,
                             self.ds.cosmology.z_from_t(t * 1e9 * 365 * 24 * 3600),
                             self.ds.cosmology.z_from_t(time.max() * 1e9 * 365 * 24 * 3600))
                         for t in time],
                        color='r', linestyle='--',
                        label='1:4 > mass ratio > 1:20 (theory)')
        # self.ax[2].scatter(time, [minor_mergers[m] for m in major_mergers],
        #    c=main_prog[main_prog.halo_ts != main_prog.halo_ts.max()].halo_ts,
        #    s=sc, cmap='summer', rasterized=True)
        # self.ax[2].scatter(time, [major_mergers[m] for m in major_mergers],
        #    c=main_prog[main_prog.halo_ts != main_prog.halo_ts.max()].halo_ts,
        #    s=sc, cmap='magma', rasterized=True, marker='*')
        self.ax[2].legend(loc='best')
        self.ax[2].set_ylabel('#merger left before t$_{max}$')
        self.ax[2].set_xlabel(r'$t$ (Gyr)')
        self.ax[2].set_xlim(time.min(), time.max())
        self.ax[2].set_ylim(ymin=0)
        z = np.unique([int(1 / self.timestep['aexp'][int(hts) - 1] - 1)
                       for hts in main_progs.halo_ts])
        z = z[1:-1]
        ax_z = self.ax[2].twiny()
        ax_z.set_xlim(self.ax[0].get_xlim())
        ax_z.set_xticks([np.copy(self.ds.cosmology.t_from_z(
            zz) / (1e9 * 365 * 24 * 3600)) for zz in z])
        ax_z.set_xticklabels(z)

        title = "Halo #{t:d}\nFinal virial mass = {x:.2e} M$_\odot$"
        title = title.format(t=int(halo.halo_num),
                             x=halo['mvir'] * 1e11)
        suptitle = self.fig.suptitle(title, fontsize=18)

        if not self.fig.texts:
            self.fig.texts.append(suptitle)
        if pdffile:
            self.fig.savefig(pdffile, format='pdf', dpi=200)
        else:
            plt.savefig('halo_{:d}.png'.format(
                int(halo.halo_num)), dpi=100, format='png')

        # A bit of cleaning
        # for axx in self.ax:
        #    axx.clear()
        #self.fig.texts = []
        fig.clf()

        return hid

    def plot_accretion_history(self, hid, scale='mass'):
        """Plot halo accretion history.

        Blah
        """
        # Define the selected halo
        halo = self.trees.ix[hid]
        sim = self.sim

        # Get progenitors
        progs = self._get_progenitors(hid)

        # Scalings
        scc = np.log10(progs.m)
        sc = (scc - scc.min()) / (scc.max() - scc.min()) * 500.

        try:
            fig = self.fig
        except AttributeError:
            self.fig = plt.figure(figsize=(12, 12))
        try:
            ax = self.ax
        except AttributeError:
            self.ax = self.fig.add_subplot(111)

        time = np.array([self.timestep['age'][hts - 1]
                         for hts in progs.halo_ts])
        self.ax.scatter(time, progs.m * 1e11, c=progs.halo_ts, s=sc,
                        cmap='summer', rasterized=True)
        self.ax.set_xlabel(r'$t$ (Gyr)')
        self.ax.set_ylabel(r'$M$ ($\mathrm{M}_{\odot}$)')
        self.ax.set_ylim(1e11 * progs.m.min() / 3., 1e11 * progs.m.max() * 3.)
        self.ax.set_xlim(time.min(), time.max())
        self.ax.set_yscale('log')

        self.fig.tight_layout()
        self.fig.savefig('halo_{:d}_history.pdf'.format(
            int(halo.halo_num)), dpi=200, format='pdf')

        return hid

    # CONVENIENCE FUNCTIONS
    def _get_timestep_number(self):
        l = len(glob(os.path.join(self.folder, 'halos_results.???')))
        return l

    def get_ioutlist(self, inputfile='input_TreeMaker.dat'):
        """Return the number of treebricks files from an input file

        The input_TreeMaker.dat file must be a list of:
            nsnaps some_number
            '/path/to/tree_bricks{n:03d}'
            '/path/to/tree_bricks{n+1:03d}'
            ...
            '/path/to/tree_bricks{n+nsnaps:03d}'
            ...

        What should be done:
            * Read all the /path/to/tree_bricks lines
            * For each line, get the tree_bricks number
            * Return the mapping between [1, nsnaps] and the tree_bricks list.
        """
        import os.path

        with open(self.folder + inputfile, 'r') as ifile:
            lines = ifile.readlines()

        outputs = []

        # We know that the name of the tree_bricks file should be
        # "tree_bricksXXX"
        tbbasename = 'tree_bricks'
        # We can skip the first line, which should be "nsnaps 1"
        for path in lines[1:]:
            # Remove those nasty "'" and \n
            tbname = os.path.basename(path[1:-2])
            assert tbname.startswith(tbbasename)
            outputs.append(int(tbname[len(tbbasename):]))

        return outputs

    def _get_progenitors(self, hid, timestep):
        if self.galaxies:
            progenitors = self.trees.loc[(self.trees.halo_num == hid) & (self.trees.halo_ts == timestep)]
            if len(progenitors) == 0:
                return
            else:
                if len(progenitors.fathersID.item()) == 1:
                    return progenitors
                else:
                    for fatherID in progenitors.fathersID.item():
                        progenitors = pd.concat([progenitors, self._get_progenitors(fatherID, timestep-1)])
                    return progenitors

        else:
            target = self.trees.loc[(self.trees.halo_num == hid) & (self.trees.halo_ts == timestep)].index
            print(target)
            mask = ((self.trees.halo_id >= self.trees.loc[target].halo_id.item()) &
                    (self.trees.halo_id <= self.trees.loc[target].last_prog.item()))

            progenitors = self.trees.loc[mask].copy()
            return progenitors

    def _read_halo(self, F, precision, big_run=False):
        [ID] = fpu.read_vector(F, 'i')
        [BushID] = fpu.read_vector(F, 'i')
        [timestep] = fpu.read_vector(F, 'i')
        [level, hosthaloID, hostsubID, nbsub, nextsub] = fpu.read_vector(F, 'i')
        [m] = fpu.read_vector(F, precision)
        [dmacc] = fpu.read_vector(F, 'd')
        [x, y, z] = fpu.read_vector(F, precision)
        [vx, vy, vz] = fpu.read_vector(F, precision)
        [Lx, Ly, Lz] = fpu.read_vector(F, precision)
        [r, a, b, c] = fpu.read_vector(F, precision)
        [ek, ep, et] = fpu.read_vector(F, precision)
        [spin] = fpu.read_vector(F, precision)
        [nbfather] = fpu.read_vector(F, 'i')
        if nbfather == 0:
            fatherID = []
            fatherMass = []
        else:
            fatherID = fpu.read_vector(F, 'i')
            fatherMass = fpu.read_vector(F, precision)
        [nbsons] = fpu.read_vector(F, 'i')
        if nbsons == 0:
            sonsID = []
            mainSon = []
        else:
            sonsID = fpu.read_vector(F, 'i')
        mainSon = fpu.read_vector(F, 'i')[0]
        if self.galaxies:
            [rvir, mvir, tvir, cvel,Reff] = fpu.read_vector(F, precision)
        else:
            [rvir, mvir, tvir, cvel] = fpu.read_vector(F, precision)
        [rho_0, r_c] = fpu.read_vector(F, precision)
        if not big_run:
            fpu.read_vector(F, 'i')  # ncont

        return (ID, timestep, level, hosthaloID, hostsubID, m,
                           dmacc, x, y, z, vx, vy, vz, Lx, Ly, Lz, r, a, b,
                           c, ek, ep, et, spin, rvir, mvir, tvir, cvel,
                           fatherID, fatherMass, sonsID, mainSon, Reff)

    def _read_tree(self, tree_file, precision, big_run=False):
        Key_tree = [
            'halo_num', 'tree_step', 'level', 'host_halo_id', 'host_sub_id', 'm',
            'dmacc', 'x', 'y', 'z', 'vx', 'vy', 'vz', 'Lx', 'Ly', 'Lz', 'r',
            'a', 'b', 'c', 'ek', 'ep', 'et', 'spin', 'rvir', 'mvir', 'tvir',
            'cvel', 'fathersID', 'fatherMass', 'sonsID', 'mainSon']
        if self.galaxies:
            Key_tree += ['Reff']

        with open(tree_file, 'rb') as F:
            self.timestep = {}
            [self.timestep['nsteps']] = fpu.read_vector(F, 'i')
            n_halo_tot = fpu.read_vector(F, 'i')
            self.timestep['nhalos'] = n_halo_tot[:self.timestep['nsteps']]
            self.timestep['aexp'] = fpu.read_vector(F, precision)
            fpu.read_vector(F, precision)
            self.timestep['age'] = fpu.read_vector(F, precision)

            data = np.empty(shape=(n_halo_tot.sum(), len(Key_tree)), dtype=object)
            if (os.path.exists('{}/tree.hdf'.format(self.folder))):
              self.trees = pd.read_hdf('{}/tree.hdf'.format(self.folder))
            else:
              print('Reading the tree for the first time... be patient!')
              j = 0
              for istep in tqdm(range(self.timestep['nsteps'])):
                for ihalo in range(n_halo_tot[istep]+n_halo_tot[istep+self.timestep['nsteps']]):
                    data[j] = self._read_halo(F, precision, big_run)
                    j += 1

              dd = {k: data[:, i]
                  for i, k in enumerate(Key_tree)}
              self.trees = pd.DataFrame(dd)
              for k in ['halo_num', 'tree_step', 'level', 'host_halo_id', 'host_sub_id']:
                    self.trees[k] = self.trees[k].astype(np.int32)
              for k in ['m', 'dmacc', 'x', 'y', 'z', 'vx', 'vy', 'vz', 'Lx', 'Ly',
                  'Lz', 'r', 'a', 'b', 'c', 'ek', 'ep', 'et', 'spin', 'rvir',
                  'mvir', 'tvir', 'cvel']:
                self.trees[k] = self.trees[k].astype(np.float32)

              #process a bit the data
              self.timestep['aexp'] = 1/(1+self.ds.cosmology.z_from_t(self.ds.arr(self.timestep['age'],'Gyr')))
              ####### Create halo_ts from the step in the tree
              self.trees['halo_ts'] = self.map_halo_ts_to_output(self.trees.tree_step)
              iout = self.trees.tree_step.values.astype(int)-1
              self.trees['aexp'] = self.timestep['aexp'][iout]
              aexp = self.trees['aexp']
              aexp_last = self.timestep['aexp'].max()

              factor = float(self.ds.length_unit.in_units('cm') / 3.08e24) * aexp / aexp_last
              self.trees['x'] = self.trees['x'] / factor + 0.5
              self.trees['y'] = self.trees['y'] / factor + 0.5
              self.trees['z'] = self.trees['z'] / factor + 0.5

              # Find outputs given halo_ts
              paths = find_outputs(self.prefix+'/Outputs')
              self.outputs = paths[-int(self.trees.halo_ts.max()):]
              step_first_gal = len(paths) - self.trees.halo_ts.max()
              self.trees.halo_ts += step_first_gal

              self.trees.m *= 1e11
              self.trees.dmacc *= 1e11

              self.trees.set_index(['halo_num','halo_ts'], drop=False, inplace=True)
              self.trees.to_hdf('{}/tree.hdf'.format(self.folder), 'hdf5')

        return self