import os
from collections import defaultdict
from typing import NamedTuple, Sequence, Mapping

import tensorflow as tf
import numpy as np
import copy
import sys

# tensorflow inputs
from jtr.preprocess.vocab import Vocab
#from jtr.preprocess.batch import get_feed_dicts
from jtr.preprocess.batch import numpify, get_batches, GeneratorWithRestart
from jtr.util import tfutil
from projects.nerre.data import read_ann, convert_to_batchable_format, convert_batch_to_ann, reset_output_dir
from projects.nerre.eval import calculateMeasures
from numpy.random import choice

def create_placeholders():
    tokens = tf.placeholder(tf.int32, [None, None], name="sentences_as_ints")  # [batch_size, max_num_tokens]
    sentence_lengths = tf.placeholder(tf.int32, [None], name="sentence_length")
    tag_labels = tf.placeholder(tf.int32, [None, None], name="bio_labels_as_ints")  # [batch_size, max_num_tokens]
    type_labels = tf.placeholder(tf.int32, [None, None], name="type_labels_as_ints")  # [batch_size, max_num_tokens]
    target_relations = tf.placeholder(tf.int32, [None, None, None],
                                      "relation_matrices")  # [batch_size, max_num_tokens, max_num_tokens]]
    document_indices = tf.placeholder(tf.int32, [None], name="document_indices")  # [batch_size]
    token_char_offsets = tf.placeholder(tf.int32, [None, None, None], name="token_char_offsets")

    placeholders = {"sentences_as_ints": tokens, "sentence_length": sentence_lengths, "bio_labels_as_ints": tag_labels,
                "type_labels_as_ints": type_labels, "relation_matrices": target_relations,
                "document_indices": document_indices, "token_char_offsets": token_char_offsets}

    return placeholders

def create_model(placeholders, output_size, layers, dropout, num_words, emb_dim, max_sent_len, tag_size=3, type_size=4, rel_size=3, relations=True, tieOutputLayer=True, hierarchicalSoftm=True, a = 1, b = 1, c = 1):
    # sequence-to-sequence bi-LSTM with hierarchical softmax for keyphrase identification, classification and relation classification
    with tf.variable_scope("embeddings"):
        embeddings = tf.get_variable("embeddings", shape=[num_words, emb_dim], dtype=tf.float32)

    with tf.variable_scope("input"):
        embedded_input = tf.gather(embeddings, placeholders["sentences_as_ints"])  # embed sentence tokens

    with tf.variable_scope("model"):
        cell = tf.nn.rnn_cell.LSTMCell(
            output_size,
            state_is_tuple=True,
            initializer=tf.contrib.layers.xavier_initializer()
        )

        if layers > 1:
            cell = tf.nn.rnn_cell.MultiRNNCell([cell] * layers)

        if dropout != 0.0:
            cell_dropout = \
                tf.nn.rnn_cell.DropoutWrapper(cell, input_keep_prob=1.0 - dropout)
        else:
            cell_dropout = cell

        outputs, states = tf.nn.bidirectional_dynamic_rnn(
            cell_dropout,
            cell_dropout,
            embedded_input,
            sequence_length=placeholders["sentence_length"],
            dtype=tf.float32
        )

        # states = (states_fw, states_bw) = ( (c_fw, h_fw), (c_bw, h_bw) )  [batch_size, max_time, cell_fw.output_size] for fw and bw each
        output = tf.concat(2, outputs)  # concatenate along output_size dimension -> [batch_size, max_time, cell_fw.output_size*2]

        dim1, dim2, dim3 = tf.unstack(tf.shape(placeholders["relation_matrices"]))

        # masking output for sentence lengths  -- doesn't seem to help
        #output_mask = mask_for_lengths(sentence_lengths, dim1, max_length=max_sent_len, dim2=emb_dim*2, value=-1000)
        #output = output + output_mask

        output_mask_keys = mask_for_lengths(placeholders["sentence_length"], dim1, max_length=max_sent_len, value=0, mask_right=False)
        output_mask_rels = mask_for_lengths(placeholders["sentence_length"], dim1, max_length=max_sent_len, dim2=dim3, value=0, mask_right=False)

        output_with_tags = tf.contrib.layers.linear(output, tag_size)

        output_with_tags_softm = tf.nn.softmax(output_with_tags)
        predicted_tags = tf.arg_max(output_with_tags_softm, 2)
        l_tags = tf.nn.sparse_softmax_cross_entropy_with_logits(output_with_tags, placeholders["bio_labels_as_ints"])

        # applying sentence length mask
        #l_tags = l_tags + output_mask_keys
        l_tags = tf.where(output_mask_keys, l_tags, tf.zeros(tf.shape(l_tags)))

        loss_tags = tf.reduce_sum(l_tags)

        output_with_labels = tf.contrib.layers.linear(output, type_size)

        # set output_with_labels [batch_size, max_num_tokens, 0] to output_with_tags [batch_size, max_num_tokens, 0]  -- i.e. the params for O should be the same for tags and labels
        if tieOutputLayer == True:
            output_with_tags_slice = tf.slice(output_with_tags, [0, 0, 0], [dim1, dim2, 1])
            output_with_labels_slice = tf.slice(output_with_labels, [0, 0, 1], [dim1, dim2, type_size-1])
            output_with_labels = tf.concat(values=[output_with_tags_slice, output_with_labels_slice], axis=2)

        output_with_labels_softm = tf.nn.softmax(output_with_labels)

        predicted_labels = tf.arg_max(output_with_labels_softm, 2)
        l_labels = tf.nn.sparse_softmax_cross_entropy_with_logits(output_with_labels, placeholders["type_labels_as_ints"])

        # applying sentence length mask
        l_labels = tf.where(output_mask_keys, l_labels, tf.zeros(tf.shape(l_labels)))

        loss_labels = tf.reduce_sum(l_labels)

        if relations == True:
            output_with_rels = tf.contrib.layers.linear(output, max_sent_len * rel_size)
            output_with_rels = tf.reshape(output_with_rels, [dim1, dim2, dim3, -1])  # batch_size, seq_len, seq_len, num_rels

            # the params for [O, O] relation weights should be same as those for O tags and labels
            # do this separately for the rows and columns in the rel label matrix
            if tieOutputLayer == True:
                output_rels_unpacked1 = tf.unstack(output_with_rels, num=max_sent_len, axis=1)  # produces dim2 number of [dim1, dim3, dim4] slices
                #output_with_rels = tf.stack(output_rels_unpacked1, axis=1)

                outputs1 = tf.TensorArray(size=max_sent_len, dtype='float32', infer_shape=False)
                i = 0
                for rel_sl in output_rels_unpacked1:  # this is of len max_sent_len
                    output_with_rels_slice = tf.slice(rel_sl, [0, 0, 1], [dim1, dim3, rel_size - 1])  # these are all the non-O weights # bio_vocab("B")
                    output_with_labels_slice = tf.slice(output_with_tags, [0, 0, 0], [dim1, dim2, 1])  # these are the O weights

                    # multiply the weights for the B labels of keyphrase identification with the weights for both relations in the rel matrix
                    if hierarchicalSoftm == True:
                        output_with_labels_slice_B = tf.slice(rel_sl, [0, 0, 1], [dim1, dim3, 2])
                        output_with_rels_slice = tf.einsum('blr,blk->blr', output_with_rels_slice, output_with_labels_slice_B)

                    output_with_rels = tf.concat(values=[output_with_labels_slice, output_with_rels_slice], axis=2)
                    outputs1 = outputs1.write(i, output_with_rels)
                    i += 1

                # packs along axis 0, there doesn't seem to be a way to change that (?)
                output_with_rels = outputs1.pack()
                output_with_rels = tf.transpose(output_with_rels, perm=[1, 0, 2, 3])

                output_rels_unpacked2 = tf.unstack(output_with_rels, num=max_sent_len, axis=2)  # produces dim1 number of [dim1, dim2, dim4] slices

                outputs2 = tf.TensorArray(size=max_sent_len, dtype='float32', infer_shape=False)
                i = 0
                for rel_sl in output_rels_unpacked2:  # this is of len max_sent_len
                    output_with_rels_slice = tf.slice(rel_sl, [0, 0, 1], [dim1, dim3, rel_size - 1])  # these are all the non-O weights
                    output_with_labels_slice = tf.slice(output_with_tags, [0, 0, 0], [dim1, dim2, 1])  # these are the O weights

                    # multiply the weights for the B labels of keyphrase identification with the weights for both relations in the rel matrix
                    if hierarchicalSoftm == True:
                        output_with_labels_slice_B = tf.slice(rel_sl, [0, 0, 1], [dim1, dim3, 2])
                        output_with_rels_slice = tf.einsum('blr,blk->blr', output_with_rels_slice,
                                                           output_with_labels_slice_B)

                    output_with_rels = tf.concat(values=[output_with_labels_slice, output_with_rels_slice], axis=2)
                    outputs2 = outputs2.write(i, output_with_rels)
                    i += 1

                # packs along axis 0, there doesn't seem to be a way to change that (?)
                output_with_rels = outputs2.pack()
                output_with_rels = tf.transpose(output_with_rels, perm=[1, 0, 2, 3])

            output_with_labels_softm = tf.nn.softmax(output_with_rels)
            predicted_rels = tf.arg_max(output_with_labels_softm, 3)
            predicted_rels = tf.reshape(predicted_rels, [dim1, dim2, dim3])

            l_relations = tf.nn.sparse_softmax_cross_entropy_with_logits(output_with_rels, placeholders["relation_matrices"])

            # applying sentence length mask
            l_relations = tf.where(output_mask_rels, l_relations, tf.zeros(tf.shape(l_relations)))

            loss_rels = tf.reduce_sum(l_relations)

            loss_all = a * loss_tags + b * loss_labels + c * loss_rels

        else:
            predicted_rels = placeholders["relation_matrices"]
            loss_all = loss_tags + loss_labels

        return loss_all, predicted_tags, predicted_labels, predicted_rels

def train(placeholders, train_batches, dev_batches, vocab, max_sent_len, max_epochs=200, l2=0.0, learning_rate=0.001, emb_dim=10, output_size=10, layers=1, dropout=0.0, sess=None, clip=None, clip_op=tf.clip_by_value, pred_on_train=False, a=1, b=1, c=1, useGoldKeyphr=False, relations=True):

    loss, predicted_tags, predicted_labels, predicted_rels = create_model(placeholders, max_sent_len=max_sent_len, output_size=output_size, layers=layers, dropout=dropout, num_words=len(vocab), emb_dim=emb_dim, a=a, b=b, c=c, relations=relations)

    optim = tf.train.AdamOptimizer(learning_rate)

    if l2 != 0.0:
        loss = loss + tf.add_n(
            [tf.nn.l2_loss(v) for v in tf.trainable_variables()]) * l2

    if clip is not None:
        gradients = optim.compute_gradients(loss)
        if clip_op == tf.clip_by_value:
            capped_gradients = [(tf.clip_by_value(grad, clip[0], clip[1]), var)
                                for grad, var in gradients]
        elif clip_op == tf.clip_by_norm:
            capped_gradients = [(tf.clip_by_norm(grad, clip), var)
                                for grad, var in gradients]
        min_op = optim.apply_gradients(capped_gradients)
    else:
        min_op = optim.minimize(loss)

    # Do not take up all the GPU memory, all the time.
    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True

    if sess is None:
        sess = tf.Session(config=sess_config)

    tf.global_variables_initializer().run(session=sess)

    saver = tf.train.Saver()

    pred_batches = []

    i = 0
    j = 1
    last_loss = sys.maxsize
    fiveagoloss = sys.maxsize
    while (i < max_epochs-1):# or last_loss > 1.0):
        loss_all = []
        for j, batch in enumerate(train_batches):
            _, current_loss = sess.run([min_op, loss], feed_dict=batch)
            loss_all.append(current_loss)

            # in last epoch, get predictions on training data
            if i == max_epochs and pred_on_train == True:
                train_pred_batches_i = {}

                curr_predicted_tags, curr_predicted_labels, curr_predicted_rels \
                    = sess.run([predicted_tags, predicted_labels, predicted_rels], feed_dict=batch)

                train_pred_batches_i["bio_labels_as_ints"], train_pred_batches_i["type_labels_as_ints"], train_pred_batches_i["relation_matrices"] \
                    = curr_predicted_tags, curr_predicted_labels, curr_predicted_rels

                train_pred_batches_i["sentences_as_ints"], train_pred_batches_i["sentence_length"], train_pred_batches_i["document_indices"], train_pred_batches_i["token_char_offsets"]\
                    = batch[placeholders["sentences_as_ints"]], batch[placeholders["sentence_length"]], batch[placeholders["document_indices"]], batch[placeholders["token_char_offsets"]]

                pred_batches.append(train_pred_batches_i)

        average_loss = float(np.average(loss_all))
        if i%10 == 0:
            print("Epoch:", i, "\tAverage loss:", average_loss)
            saver.save(sess, "/tmp/model.ckpt")

        # early stopping
        if (average_loss > last_loss and average_loss > fiveagoloss):
            print("Stopping early", str(average_loss), str(last_loss), str(fiveagoloss))
            break

        last_loss = average_loss
        if j ==5:
            fiveagoloss = average_loss
            j = 0

        i += 1
        j -=1

    # Save the variables to disk.
    save_path = saver.save(sess, "/tmp/model.ckpt")
    print("Model saved in file: %s" % save_path)

    if pred_on_train == False:
        for j, dev_batch in enumerate(dev_batches):
            dev_pred_batches_i = {}

            curr_predicted_tags, curr_predicted_labels, curr_predicted_rels \
                = sess.run([predicted_tags, predicted_labels, predicted_rels], feed_dict=dev_batch)

            if useGoldKeyphr == False:
                dev_pred_batches_i["bio_labels_as_ints"], dev_pred_batches_i[
                    "type_labels_as_ints"] = curr_predicted_tags, curr_predicted_labels
                dev_pred_batches_i["relation_matrices"] = dev_batch[placeholders["relation_matrices"]]
            else:
                dev_pred_batches_i["bio_labels_as_ints"], dev_pred_batches_i["type_labels_as_ints"], \
                dev_pred_batches_i["relation_matrices"] \
                    = dev_batch[placeholders["bio_labels_as_ints"]], dev_batch[placeholders["type_labels_as_ints"]], \
                      dev_batch[placeholders["relation_matrices"]]

            dev_pred_batches_i["sentences_as_ints"], dev_pred_batches_i["sentence_length"], \
            dev_pred_batches_i["document_indices"], dev_pred_batches_i["token_char_offsets"] \
                = dev_batch[placeholders["sentences_as_ints"]], dev_batch[placeholders["sentence_length"]], dev_batch[
                placeholders["document_indices"]], dev_batch[placeholders["token_char_offsets"]]

            pred_batches.append(dev_pred_batches_i)

    return pred_batches

def test_only(placeholders, dev_batches, vocab, max_sent_len, emb_dim=10, output_size=10, layers=1, dropout=0.0, sess=None, pred_on_train=False, a=1, b=1, c=1, useGoldKeyphr=False, relations=True):

    loss, predicted_tags, predicted_labels, predicted_rels = create_model(placeholders, max_sent_len=max_sent_len, output_size=output_size, layers=layers, dropout=dropout, num_words=len(vocab), emb_dim=emb_dim, a=a, b=b, c=c, relations=relations)

    # Do not take up all the GPU memory, all the time.
    sess_config = tf.ConfigProto()
    sess_config.gpu_options.allow_growth = True

    if sess is None:
        sess = tf.Session(config=sess_config)

    tf.global_variables_initializer().run(session=sess)

    saver = tf.train.Saver()

    # Save the variables to disk.
    save_path = saver.restore(sess, "/tmp/model.ckpt")
    print("Model restored from %s" % save_path)

    pred_batches = []

    if pred_on_train == False:
        for j, dev_batch in enumerate(dev_batches):
            dev_pred_batches_i = {}

            curr_predicted_tags, curr_predicted_labels, curr_predicted_rels \
                = sess.run([predicted_tags, predicted_labels, predicted_rels], feed_dict=dev_batch)

            if useGoldKeyphr == False:
                dev_pred_batches_i["bio_labels_as_ints"], dev_pred_batches_i[
                    "type_labels_as_ints"] = curr_predicted_tags, curr_predicted_labels
                dev_pred_batches_i["relation_matrices"] = dev_batch[placeholders["relation_matrices"]]
            else:
                dev_pred_batches_i["bio_labels_as_ints"], dev_pred_batches_i["type_labels_as_ints"], \
                dev_pred_batches_i["relation_matrices"] \
                    = dev_batch[placeholders["bio_labels_as_ints"]], dev_batch[placeholders["type_labels_as_ints"]], \
                      dev_batch[placeholders["relation_matrices"]]

            dev_pred_batches_i["sentences_as_ints"], dev_pred_batches_i["sentence_length"], \
            dev_pred_batches_i["document_indices"], dev_pred_batches_i["token_char_offsets"] \
                = dev_batch[placeholders["sentences_as_ints"]], dev_batch[placeholders["sentence_length"]], dev_batch[
                placeholders["document_indices"]], dev_batch[placeholders["token_char_offsets"]]

            pred_batches.append(dev_pred_batches_i)

    return pred_batches

def get_feed_dicts(data, placeholders, batch_size=32, pad=0, bucket_order=None, bucket_structure=None,
                   exact_epoch=False):
    """Creates feed dicts for all batches with a given batch size.

    Args:
        `data` (dict): The input data for the feed dicts.
        `placeholders` (dict): The TensorFlow placeholders for the data
            (placeholders.keys() must form a subset of data.keys()).
        `batch_size` (int): The batch size for the data.
        `pad` (int): Padding symbol index to pad lists of different sizes.
        `bucket_order`: argument `order` in get_buckets (list with keys); `None` if no bucketing
        `bucket_structure`: argument `structure` in get_buckets; `None` if no bucketing
        `exact_epoch`: if set to `True`, final batch per bucket may be smaller, but each instance will be seen exactly
            once during training. Default: `False`, to be certain during training
            that each instance per batch gets same weight in the total loss.

    Returns:
        GeneratorWithRestart: Generator that yields a feed_dict for each
        iteration. A feed dict consists of '{ placeholder : data-batch }` key-value pairs."""

    assert isinstance(data, dict) and isinstance(placeholders, dict)
    assert set(placeholders.keys()).issubset(set(data.keys())), \
        'data keys %s \nnot compatible with placeholder keys %s' % (set(placeholders.keys()), set(data.keys()))

    def generator():
        batches = get_batches(data, batch_size, pad, bucket_order, bucket_structure, exact_epoch)
        # fixme: this is potentially inefficient as it might be called every time we retrieve a batch
        # todo: measure and fix if significant impact
        mapped = map(lambda xs: {placeholders[k]: xs[k] for k in placeholders}, batches)
        # for each key in placeholders dict, pair the placeholder with the corresponding batch dict value
        for x in mapped:
            yield x

    return GeneratorWithRestart(generator)

def mask_for_lengths(lengths, batch_size=None, max_length=None, dim2=None, mask_right=True,
                         value=-1000.0):

    """
    Creates a [batch_size x max_length] mask.
    :param lengths: int32 1-dim tensor of batch_size lengths
    :param batch_size: int32 0-dim tensor or python int
    :param max_length: int32 0-dim tensor or python int
    :param mask_right: if True, everything before "lengths" becomes zero and the
        rest "value", else vice versa
    :param value: value for the mask
    :return: [batch_size x max_length] mask of zeros and "value"s
    """
    if max_length is None:
        max_length = tf.reduce_max(lengths)
    if batch_size is None:
        batch_size = tf.shape(lengths)[0]

    if dim2 != None:
        # [batch_size x max_length x dim2]
        mask = tf.reshape(tf.tile(tf.range(0, max_length), [batch_size * dim2]), tf.stack([batch_size, max_length, dim2]))
        if mask_right:
            mask = tf.greater_equal(mask, tf.expand_dims(tf.expand_dims(lengths, 1), 1))
        else:
            mask = tf.less(mask, tf.expand_dims(tf.expand_dims(lengths, 1), 1))
    else:
        # [batch_size x max_length]
        mask = tf.reshape(tf.tile(tf.range(0, max_length), [batch_size]), tf.stack([batch_size, -1]))
        if mask_right:
            mask = tf.greater_equal(mask, tf.expand_dims(lengths, 1))
        else:
            mask = tf.less(mask, tf.expand_dims(lengths, 1))
    # otherwise we return a boolean mask
    if value != 0:
        mask = tf.cast(mask, tf.float32) * value
    return mask

def get_feed_dicts_new(data_train_np, placeholders, batch_size, inst_length):
    # around 8 times faster as with generator as it doesn't need to go over the whole training data every time during training
    data_train_batched = []
    realsamp = int(inst_length/batch_size)
    additionsamp = inst_length%batch_size
    if additionsamp != 0:
        realsamp += 1
    ids1 = choice(range(0, inst_length), inst_length-additionsamp, replace=False)  # sample without replacement so we get every sample once
    ids2 = choice(range(0, inst_length), additionsamp, replace=True)  # sample a few additional ones to fill up batch
    ids = np.append(ids1, ids2)

    start = 0
    for i in range(0, realsamp):
        batch_i = {}
        if i != 0:
            start = i * batch_size
        if i != realsamp:
            ids_sup = ids[start:((i+1)*batch_size)]
        else:
            ids_sup = ids[start:realsamp]
        for key, value in data_train_np.items():
            batch_i[placeholders[key]] = [data_train_np[key][ii] for ii in ids_sup]

        data_train_batched.append(batch_i)

    return data_train_batched

if __name__ == "__main__":
    train_dir = "/Users/Isabelle/Documents/UCLMR/semeval2017-orga/data/train2"
    dev_dir = "/Users/Isabelle/Documents/UCLMR/semeval2017-orga/data/dev/"
    test_dir = "/Users/Isabelle/Documents/UCLMR/semeval_articles/test_final2/"
    out_dir = "/tmp/"
    train_mode = False
    final_test = True

    if final_test == True:
        train_instances = read_ann(train_dir) + read_ann(dev_dir)
        dev_instances = read_ann(test_dir)
        dev_dir = test_dir
    else:
        train_instances = read_ann(train_dir)
        dev_instances = read_ann(dev_dir)

    print("Loaded {} train instances".format(len(train_instances)))
    print("Loaded {} dev instances".format(len(dev_instances)))

    vocab = Vocab()

    numinst = 10
    for instance in train_instances:
        for sent in instance.doc:
            for token in sent.tokens:
                vocab(token.word)

    print("Collected {} word types".format(len(vocab)))

    train_batchable = convert_to_batchable_format(train_instances, vocab)
    dev_batchable = convert_to_batchable_format(dev_instances, vocab)

    batch_size = 1024 #128

    data_all = {}
    for key, value in train_batchable.items():
        data_all[key] = train_batchable[key] + dev_batchable[key]

    data_np = numpify(data_all, pad=0)
    data_train_np = {}
    data_dev_np = {}
    for key, value in data_np.items():
        data_train_np[key] = data_np[key][0:len(train_batchable[key])]
        data_dev_np[key] = data_np[key][len(train_batchable[key]):len(train_batchable[key])+len(dev_batchable[key])]

    max_sent_len = len(data_train_np["relation_matrices"][0][0])

    if train_mode == False:
        placeholders = create_placeholders()

        dim = 100

        reset_output_dir()

        dev_feed_dicts = get_feed_dicts_new(data_dev_np, placeholders, batch_size, len(dev_instances))
        dev_preds = test_only(placeholders, dev_feed_dicts, vocab, max_sent_len,
                          emb_dim=dim, output_size=dim, a=1, b=1, c=1, useGoldKeyphr=False,
                          relations=False)

        for dev_batch in dev_preds:
            convert_batch_to_ann(dev_batch, dev_instances, out_dir=out_dir)

        calculateMeasures(dev_dir, out_dir, ignoremissing=False, remove_anno="types")

    else:

        for dim in [100]:
            for drop in [0.2]:#, 0.5]:
                for l2 in [0.2]:#, 0.5]:
                    for a in [1]:#, 2, 3]:
                        for b in [1]:#, 2, 3]:
                            print("Training with dim", dim, "drop", drop, "l2", l2, "a", a, "b", b)

                            placeholders = create_placeholders()

                            train_feed_dicts = get_feed_dicts_new(data_train_np, placeholders, batch_size, len(train_instances))

                            dev_feed_dicts = get_feed_dicts_new(data_dev_np, placeholders, batch_size, len(dev_instances))

                            reset_output_dir()

                            dev_preds = train(placeholders, train_feed_dicts, dev_feed_dicts, vocab, max_sent_len, max_epochs=250, emb_dim=dim, output_size=dim, l2=l2, dropout=drop, a=a, b=b, c=0, useGoldKeyphr=False, relations=False)

                            for dev_batch in dev_preds:
                                convert_batch_to_ann(dev_batch, dev_instances, out_dir=out_dir)

                            calculateMeasures(dev_dir, out_dir, ignoremissing=False, remove_anno="rel")

                            tf.reset_default_graph()


# -*- coding: utf-8 -*-
#
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# Using this computer program means that you agree to the terms
# in the LICENSE file included with this software distribution.
# Any use not explicitly granted by the LICENSE is prohibited.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# For comments or questions, please email us at pixie@tue.mpg.de
# For commercial licensing contact, please contact ps-license@tuebingen.mpg.de

import os, sys
import torch
import torchvision
import torch.nn.functional as F
import torch.nn as nn

import numpy as np
from skimage.io import imread
import cv2

from .models.encoders import ResnetEncoder, MLP, HRNEncoder
from .models.moderators import TempSoftmaxFusion
from .models.FLAME import FLAMETex
from .models.SMPLX import SMPLX
from .utils import util
from .utils import rotation_converter as converter
from .utils import tensor_cropper
from .utils.config import cfg

class PIXIE(object):
    def __init__(self, config=None, device='cuda:0'):
        if config is None:
            self.cfg = cfg
        else:
            self.cfg = config

        self.device = device
        # parameters setting
        self.param_list_dict = {}
        for lst in self.cfg.params.keys():
            param_list = cfg.params.get(lst)
            self.param_list_dict[lst] = {i:cfg.model.get('n_'+i) for i in param_list}

        # Build the models
        self._create_model()
        # Set up the cropping modules used to generate face/hand crops from the body predictions
        self._setup_cropper()

    def _setup_cropper(self):
        self.Cropper = {}
        for crop_part in ['head', 'hand']:
            data_cfg = self.cfg.dataset[crop_part]
            scale_size = (data_cfg.scale_min + data_cfg.scale_max)*0.5
            self.Cropper[crop_part] = tensor_cropper.Cropper(
                        crop_size=data_cfg.image_size,
                        scale=[scale_size, scale_size],
                        trans_scale = 0)

    def _create_model(self):
        self.model_dict = {}
        # Build all image encoders
        # Hand encoder only works for right hand, for left hand, flip inputs and flip the results back
        self.Encoder = {}
        for key in self.cfg.network.encoder.keys():
            if self.cfg.network.encoder.get(key).type == 'resnet50':
                self.Encoder[key] = ResnetEncoder().to(self.device)
            elif self.cfg.network.encoder.get(key).type == 'hrnet':
                self.Encoder[key] = HRNEncoder().to(self.device)
            self.model_dict[f'Encoder_{key}'] = self.Encoder[key].state_dict()

        # Build the parameter regressors
        self.Regressor = {}
        for key in self.cfg.network.regressor.keys():
            n_output = sum(self.param_list_dict[f'{key}_list'].values())
            channels = [2048] + self.cfg.network.regressor.get(key).channels + [n_output]
            if self.cfg.network.regressor.get(key).type == 'mlp':
                self.Regressor[key] = MLP(channels=channels).to(self.device)
            self.model_dict[f'Regressor_{key}'] = self.Regressor[key].state_dict()

        # Build the extractors
        # to extract separate head/left hand/right hand feature from body feature
        self.Extractor = {}
        for key in self.cfg.network.extractor.keys():
            channels = [2048] + self.cfg.network.extractor.get(key).channels + [2048]
            if self.cfg.network.extractor.get(key).type == 'mlp':
                self.Extractor[key] = MLP(channels=channels).to(self.device)
            self.model_dict[f'Extractor_{key}'] = self.Extractor[key].state_dict()

        # Build the moderators
        self.Moderator = {}
        for key in self.cfg.network.moderator.keys():
            share_part = key.split('_')[0]
            detach_inputs = self.cfg.network.moderator.get(key).detach_inputs
            detach_feature = self.cfg.network.moderator.get(key).detach_feature
            channels = [2048*2] + self.cfg.network.moderator.get(key).channels + [2]
            self.Moderator[key] = TempSoftmaxFusion(
                detach_inputs=detach_inputs, detach_feature=detach_feature,
                channels=channels).to(self.device)
            self.model_dict[f'Moderator_{key}'] = self.Moderator[key].state_dict()

        # Build the SMPL-X body model, which we also use to represent faces and
        # hands, using the relevant parts only
        self.smplx = SMPLX(self.cfg.model).to(self.device)
        self.part_indices = self.smplx.part_indices
        # Build the FLAME texture space
        if self.cfg.model.use_tex:
            self.flametex = FLAMETex(self.cfg.model).to(self.device)

        #-- resume model
        model_path = self.cfg.pretrained_modelpath
        if os.path.exists(model_path):
            checkpoint = torch.load(model_path)
            for key in self.model_dict.keys():
                util.copy_state_dict(self.model_dict[key], checkpoint[key])
        else:
            print(f'pixie trained model path: {model_path} does not exist!')
            exit()
        # eval mode
        for module in [self.Encoder, self.Regressor, self.Moderator, self.Extractor]:
            for net in module.values():
                net.eval()

    def decompose_code(self, code, num_dict):
        ''' Convert a flattened parameter vector to a dictionary of parameters
        '''
        code_dict = {}
        start = 0
        for key in num_dict:
            end = start+int(num_dict[key])
            code_dict[key] = code[:, start:end]
            start = end
        return code_dict

    def part_from_body(self, image, part_key, points_dict, crop_joints = None):
        ''' crop part(head/left_hand/right_hand) out from body data, joints also change accordingly
        '''
        assert part_key in ['head', 'left_hand', 'right_hand']
        assert 'smplx_kpt' in points_dict.keys()
        if part_key == 'head':
            # use face 68 kpts for cropping head image
            indices_key = 'face'
        elif part_key == 'left_hand':
            indices_key = 'left_hand'
        elif part_key == 'right_hand':
            indices_key = 'right_hand'

        # get points for cropping
        part_indices = self.part_indices[indices_key]
        if crop_joints is not None:
            points_for_crop = crop_joints[:, part_indices]
        else:
            points_for_crop = points_dict['smplx_kpt'][:, part_indices]

        # crop
        cropper_key = 'hand' if 'hand' in part_key else part_key
        points_scale = image.shape[-2:]
        cropped_image, tform = self.Cropper[cropper_key].crop(
                            image,
                            points_for_crop,
                            points_scale
                            )
        # transform points(must be normalized to [-1.1]) accordingly
        cropped_points_dict = {}
        for points_key in points_dict.keys():
            points = points_dict[points_key]
            cropped_points = self.Cropper[cropper_key].transform_points(points, tform, points_scale, normalize=True)
            cropped_points_dict[points_key] = cropped_points
        return cropped_image, cropped_points_dict

    @torch.no_grad()
    def encode(self, data, threthold=True, keep_local=True, copy_and_paste=False, body_only=False):
        ''' Encode images to smplx parameters
        Args:
            data: dict
                key: image_type (body/head/hand)
                value:
                    image: [bz, 3, 224, 224], range [0,1]
                    image_hd(needed if key==body): a high res version of image, only for cropping parts from body image
                    head_image: optinal, well-cropped head from body image
                    left_hand_image: optinal, well-cropped left hand from body image
                    right_hand_image: optinal, well-cropped right hand from body image
        Returns:
            param_dict: dict
                key: image_type (body/head/hand)
                value: param_dict
        '''
        for key in data.keys():
            assert key in ['body', 'head', 'hand']

        feature = {}
        param_dict = {}

        # Encode features
        for key in data.keys():
            part = key
            # encode feature
            feature[key] = {}
            feature[key][part] = self.Encoder[part](data[key]['image'])

            # for head/hand image
            if key == 'head' or key == 'hand':
                # predict head/hand-only parameters from part feature
                part_dict = self.decompose_code(self.Regressor[part](feature[key][part]), self.param_list_dict[f'{part}_list'])
                # if input is part data, skip feature fusion: share feature is the same as part feature
                # then predict share parameters
                feature[key][f'{key}_share'] = feature[key][key]
                share_dict = self.decompose_code(
                    self.Regressor[f'{part}_share'](feature[key][f'{part}_share']),
                    self.param_list_dict[f'{part}_share_list'])
                # compose parameters
                param_dict[key] = {**share_dict, **part_dict}

            # for body image
            if key == 'body':
                fusion_weight = {}
                f_body = feature['body']['body']
                # extract part feature
                for part_name in ['head', 'left_hand', 'right_hand']:
                    feature['body'][f'{part_name}_share'] = self.Extractor[f'{part_name}_share'](f_body)

                # -- check if part crops are given, if not, crop parts by coarse body estimation
                if 'head_image' not in data[key].keys() \
                    or 'left_hand_image' not in data[key].keys() \
                    or 'right_hand_image' not in data[key].keys():
                    #- run without fusion to get coarse estimation, for cropping parts
                    # body only
                    body_dict = self.decompose_code(self.Regressor[part](feature[key][part]), self.param_list_dict[part+'_list'])
                    # head share
                    head_share_dict = self.decompose_code(self.Regressor['head'+'_share'](feature[key]['head'+'_share']), self.param_list_dict['head'+'_share_list'])
                    # right hand share
                    right_hand_share_dict = self.decompose_code(self.Regressor['hand'+'_share'](feature[key]['right_hand'+'_share']), self.param_list_dict['hand'+'_share_list'])
                    # left hand share
                    left_hand_share_dict = self.decompose_code(self.Regressor['hand'+'_share'](feature[key]['left_hand'+'_share']), self.param_list_dict['hand'+'_share_list'])
                    # change the dict name from right to left
                    left_hand_share_dict['left_hand_pose'] = left_hand_share_dict.pop('right_hand_pose')
                    left_hand_share_dict['left_wrist_pose'] = left_hand_share_dict.pop('right_wrist_pose')
                    param_dict[key] = {**body_dict, **head_share_dict, **left_hand_share_dict, **right_hand_share_dict}
                    if body_only:
                        param_dict['moderator_weight'] = None
                        return param_dict
                    prediction_body_only = self.decode(param_dict[key], param_type='body')
                    # crop
                    for part_name in ['head', 'left_hand', 'right_hand']:
                        part = part_name.split('_')[-1]
                        points_dict = {
                            'smplx_kpt': prediction_body_only['smplx_kpt'],
                            'trans_verts': prediction_body_only['transformed_vertices']
                        }
                        cropped_image, cropped_joints_dict = self.part_from_body(data['body']['image_hd'], part_name, points_dict)
                        data[key][part_name+'_image'] = cropped_image

                # -- encode features from part crops, then fuse feature using the weight from moderator
                for part_name in ['head', 'left_hand', 'right_hand']:
                    part = part_name.split('_')[-1]
                    cropped_image = data[key][part_name+'_image']
                    # if left hand, flip it as if it is right hand
                    if part_name == 'left_hand':
                        cropped_image = torch.flip(cropped_image, dims=(-1,))
                    # run part regressor
                    f_part = self.Encoder[part](cropped_image)
                    part_dict = self.decompose_code(self.Regressor[part](f_part), self.param_list_dict[f'{part}_list'])
                    part_share_dict = self.decompose_code(self.Regressor[f'{part}_share'](f_part), self.param_list_dict[f'{part}_share_list'])
                    param_dict['body_'+part_name] = {**part_dict, **part_share_dict}

                    # moderator to assign weight, then integrate features
                    f_body_out, f_part_out, f_weight = self.Moderator[f'{part}_share'](feature['body'][f'{part_name}_share'], f_part, work=True)
                    if copy_and_paste:
                        # copy and paste strategy always trusts the results from part
                        feature['body'][f'{part_name}_share'] = f_part
                    elif threthold and part == 'hand':
                        # for hand, if part weight > 0.7 (very confident, then fully trust part)
                        part_w = f_weight[:,[1]]
                        part_w[part_w > 0.7] = 1.
                        f_body_out = feature['body'][f'{part_name}_share']*(1. - part_w) + f_part*part_w
                        feature['body'][f'{part_name}_share'] = f_body_out
                    else:
                        feature['body'][f'{part_name}_share'] = f_body_out
                    fusion_weight[part_name] = f_weight
                # save weights from moderator, that can be further used for optimization/running specific tasks on parts
                param_dict['moderator_weight'] = fusion_weight

                # -- predict parameters from fused body feature
                # head share
                head_share_dict = self.decompose_code(self.Regressor['head'+'_share'](feature[key]['head'+'_share']), self.param_list_dict['head'+'_share_list'])
                # right hand share
                right_hand_share_dict = self.decompose_code(self.Regressor['hand'+'_share'](feature[key]['right_hand'+'_share']), self.param_list_dict['hand'+'_share_list'])
                # left hand share
                left_hand_share_dict = self.decompose_code(self.Regressor['hand'+'_share'](feature[key]['left_hand'+'_share']), self.param_list_dict['hand'+'_share_list'])
                # change the dict name from right to left
                left_hand_share_dict['left_hand_pose'] = left_hand_share_dict.pop('right_hand_pose')
                left_hand_share_dict['left_wrist_pose'] = left_hand_share_dict.pop('right_wrist_pose')
                param_dict['body'] = {**body_dict, **head_share_dict, **left_hand_share_dict, **right_hand_share_dict}
                # copy tex param from head param dict to body param dict
                param_dict['body']['tex'] = param_dict['body_head']['tex']
                param_dict['body']['light'] = param_dict['body_head']['light']

                if keep_local:
                    # for local change that will not affect whole body and produce unnatral pose, trust part
                    param_dict[key]['exp'] = param_dict['body_head']['exp']
                    param_dict[key]['right_hand_pose'] = param_dict['body_right_hand']['right_hand_pose']
                    param_dict[key]['left_hand_pose'] = param_dict['body_left_hand']['right_hand_pose']

        return param_dict

    def convert_pose(self, param_dict, param_type):
        ''' Convert pose parameters to rotation matrix
        Args:
            param_dict: smplx parameters
            param_type: should be one of body/head/hand
        Returns:
            param_dict: smplx parameters
        '''
        assert param_type in ['body', 'head', 'hand']

        # convert pose representations: the output from network are continous repre or axis angle,
        # while the input pose for smplx need to be rotation matrix
        for key in param_dict:
            if "pose" in key and 'jaw' not in key:
                param_dict[key] = converter.batch_cont2matrix(param_dict[key])
        if param_type == 'body' or param_type == 'head':
            param_dict['jaw_pose'] = converter.batch_euler2matrix(param_dict['jaw_pose'])[:, None, :, :]

        # complement params if it's not in given param dict
        if param_type == 'head':
            batch_size = param_dict['shape'].shape[0]
            param_dict['abs_head_pose'] = param_dict['head_pose'].clone()
            param_dict['global_pose'] = param_dict['head_pose']
            param_dict['partbody_pose'] = self.smplx.body_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)[:,:self.param_list_dict['body_list']['partbody_pose']]
            param_dict['neck_pose'] = self.smplx.neck_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
            param_dict['left_wrist_pose'] = self.smplx.neck_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
            param_dict['left_hand_pose'] = self.smplx.left_hand_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
            param_dict['right_wrist_pose'] = self.smplx.neck_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
            param_dict['right_hand_pose'] = self.smplx.right_hand_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
        elif param_type == 'hand':
            batch_size = param_dict['right_hand_pose'].shape[0]
            param_dict['abs_right_wrist_pose'] = param_dict['right_wrist_pose'].clone()
            dtype = param_dict['right_hand_pose'].dtype; device=param_dict['right_hand_pose'].device
            x_180_pose = torch.eye(3, dtype=dtype, device=device).unsqueeze(0).repeat(1,1,1)
            x_180_pose[0,2,2] = -1.
            x_180_pose[0,1,1] = -1.
            param_dict['global_pose'] = x_180_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
            param_dict['shape'] = self.smplx.shape_params.expand(batch_size, -1)
            param_dict['exp'] = self.smplx.expression_params.expand(batch_size, -1)
            param_dict['head_pose'] = self.smplx.head_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
            param_dict['neck_pose'] = self.smplx.neck_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
            param_dict['jaw_pose'] = self.smplx.jaw_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
            param_dict['partbody_pose'] = self.smplx.body_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)[:,:self.param_list_dict['body_list']['partbody_pose']]
            param_dict['left_wrist_pose'] = self.smplx.neck_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
            param_dict['left_hand_pose'] = self.smplx.left_hand_pose.unsqueeze(0).expand(batch_size, -1, -1, -1)
        elif param_type == 'body':
            # the predcition from the head and hand share regressor is always absolute pose
            batch_size = param_dict['shape'].shape[0]
            param_dict['abs_head_pose'] = param_dict['head_pose'].clone()
            param_dict['abs_right_wrist_pose'] = param_dict['right_wrist_pose'].clone()
            param_dict['abs_left_wrist_pose'] = param_dict['left_wrist_pose'].clone()
            ## the body-hand share regressor is working for right hand
            ## so we assume body network get the flipped feature for the left hand. then get the parameters
            ## then we need to flip it back to left, which matches the input left hand
            param_dict['left_wrist_pose'] = util.flip_pose(param_dict['left_wrist_pose'])
            param_dict['left_hand_pose'] = util.flip_pose(param_dict['left_hand_pose'])
        else:
            exit()

        return param_dict

    def decode(self, param_dict, param_type):
        ''' Decode model parameters to smplx vertices & joints & texture
        Args:
            param_dict: smplx parameters
            param_type: should be one of body/head/hand
        Returns:
            predictions: smplx predictions
        '''
        if 'jaw_pose' in param_dict.keys() and len(param_dict['jaw_pose'].shape) == 2:
            self.convert_pose(param_dict, param_type)
        elif param_dict['right_wrist_pose'].shape[-1] == 6:
            self.convert_pose(param_dict, param_type)

        # concatenate body pose
        partbody_pose = param_dict['partbody_pose']
        param_dict['body_pose'] = torch.cat(
            [partbody_pose[:,:11],
            param_dict['neck_pose'],
            partbody_pose[:,11:11+2],
            param_dict['head_pose'],
            partbody_pose[:,13:13+4],
            param_dict['left_wrist_pose'],
            param_dict['right_wrist_pose']], dim=1)

        # change absolute head&hand pose to relative pose according to rest body pose
        if param_type == 'head' or param_type == 'body':
            param_dict['body_pose'] = self.smplx.pose_abs2rel(param_dict['global_pose'], param_dict['body_pose'], abs_joint='head')
        if param_type == 'hand' or param_type == 'body':
            param_dict['body_pose'] = self.smplx.pose_abs2rel(param_dict['global_pose'], param_dict['body_pose'], abs_joint='left_wrist')
            param_dict['body_pose'] = self.smplx.pose_abs2rel(param_dict['global_pose'], param_dict['body_pose'], abs_joint='right_wrist')

        if self.cfg.model.check_pose:
            # check if pose is natural (relative rotation), if not, set relative to 0 (especially for head pose)
            # xyz: pitch(positive for looking down), yaw(positive for looking left), roll(rolling chin to left)
            for pose_ind in [14]: #head [15-1, 20-1, 21-1]:
                curr_pose = param_dict['body_pose'][:, pose_ind]
                euler_pose = converter._compute_euler_from_matrix(curr_pose)
                for i, max_angle in enumerate([20,70,10]):
                    euler_pose_curr = euler_pose[:,i]
                    euler_pose_curr[euler_pose_curr!=torch.clamp(euler_pose_curr, min=-max_angle*np.pi/180, max=max_angle*np.pi/180)] = 0.
                param_dict['body_pose'][:,pose_ind] = converter.batch_euler2matrix(euler_pose)

        # SMPLX
        verts, landmarks, joints = self.smplx(
            shape_params=param_dict['shape'],
            expression_params=param_dict['exp'],
            global_pose=param_dict['global_pose'],
            body_pose=param_dict['body_pose'],
            jaw_pose=param_dict['jaw_pose'],
            left_hand_pose=param_dict['left_hand_pose'],
            right_hand_pose=param_dict['right_hand_pose'])
        smplx_kpt3d = joints.clone()

        # projection
        cam = param_dict[param_type + '_cam']
        trans_verts = util.batch_orth_proj(verts, cam)
        predicted_landmarks = util.batch_orth_proj(landmarks, cam)[:,:,:2]
        predicted_joints = util.batch_orth_proj(joints, cam)[:,:,:2]

        # save predcition
        prediction = {
                    'vertices': verts,
                    'transformed_vertices': trans_verts,
                    'face_kpt': predicted_landmarks,
                    'smplx_kpt': predicted_joints,
                    'smplx_kpt3d': smplx_kpt3d,
                    'joints': joints,
                    'cam': cam,
                    }
        # change the order of face keypoints, to be the same as "standard" 68 keypoints
        prediction['face_kpt'] = torch.cat([prediction['face_kpt'][:,-17:], prediction['face_kpt'][:,:-17]], dim=1)
        # return albedo map
        if 'tex' in param_dict.keys() and self.cfg.model.use_tex:
            albedo = self.flametex(param_dict['tex'])
            prediction['albedo'] = albedo
            prediction['light'] = param_dict['light'].reshape(-1, 9, 3)
        elif self.cfg.model.use_tex:
            texcode = torch.zeros([verts.shape[0], self.cfg.model.n_tex], device=self.device, dtype=torch.float32)
            albedo = self.flametex(texcode)
            prediction['albedo'] = albedo
        else:
            prediction['albedo'] = torch.zeros([verts.shape[0], 3, self.cfg.model.uv_size, self.cfg.model.uv_size], device=self.device, dtype=torch.float32)

        return prediction

    def decode_Tpose(self, param_dict):
        ''' return body mesh in T pose, support body and head param dict only
        '''
        verts, _, _ = self.smplx(
            shape_params=param_dict['shape'],
            expression_params=param_dict['exp'],
            jaw_pose=param_dict['jaw_pose'])
        return verts


import os
import cv2
import json
import numpy as np
import tensorflow as tf
import scipy.stats as st
from ..human import Human,BodyPart

class Post_Processor:
    def __init__(self,parts,limbs,colors,debug=False):
        self.cur_id=0
        self.parts=parts
        self.limbs=limbs
        self.colors=colors
        if(self.colors==None):
            self.colors=[[255,0,0]]*len(parts)
        self.n_pos=len(self.parts)
        self.n_limb=len(self.limbs)
        self.thres_conf=0.05
        self.thres_vec=0.05
        self.thres_vec_cnt=6
        self.thres_criterion2=0
        self.thres_part_cnt=4
        self.thres_human_score=0.3
        self.step_paf=10
        self.debug=debug

    def get_peak_map(self,conf_map):
        def _gauss_smooth(origin):
            sigma=3.0
            kernel_size=5
            smoothed=np.zeros(shape=origin.shape)
            channel_num=origin.shape[-1]
            for channel_idx in range(0,channel_num):
                smoothed[0,:,:,channel_idx]=cv2.GaussianBlur(origin[0,:,:,channel_idx],\
                    ksize=(kernel_size,kernel_size),sigmaX=sigma,sigmaY=sigma)
            return smoothed

        smoothed = _gauss_smooth(conf_map)
        max_pooled = tf.nn.pool(smoothed, window_shape=(3, 3), pooling_type='MAX', padding='SAME')
        return tf.where(tf.equal(smoothed, max_pooled), conf_map, tf.zeros_like(conf_map)).numpy()

    def process(self,conf_map,paf_map,img_h,img_w,data_format="channels_first"):
        if(data_format=="channels_first"):
            conf_map=np.transpose(conf_map,[1,2,0])
            paf_map=np.transpose(paf_map,[1,2,0])
        #conf_map
        if(conf_map.shape[0]!=img_h or conf_map.shape[1]!=img_w):
            conf_map=cv2.resize(conf_map,(img_w,img_h))
        conf_map=conf_map[np.newaxis,:,:,:]#TODO(JZ)
        #paf_map
        if(paf_map.shape[0]!=img_h or paf_map.shape[1]!=img_w):
            paf_map=cv2.resize(paf_map,(img_w,img_h))
        paf_map=paf_map[np.newaxis,:,:,:]
        peak_map=self.get_peak_map(conf_map)
        humans=self.process_paf(peak_map[0],conf_map[0],paf_map[0])
        return humans

    def process_paf(self,peak_map,conf_map,paf_map):
        #filter valid peaks
        peaks=[[] for part_idx in range(0,self.n_pos)]
        all_peaks=[]
        #print(f"test peak_map.shape:{peak_map.shape}")
        peak_ys,peak_xs,part_idxs=np.where(peak_map>self.thres_conf)
        for peak_idx,(part_idx,peak_y,peak_x) in enumerate(zip(part_idxs,peak_ys,peak_xs)):
            peak_score=conf_map[peak_y,peak_x,part_idx]
            peaks[part_idx].append(Peak(peak_idx,part_idx,peak_y,peak_x,peak_score))
            all_peaks.append(Peak(peak_idx,part_idx,peak_y,peak_x,peak_score))
        #start candidate connection
        if(self.debug):
            print("peak debug:")
            for part_idx in range(0,self.n_pos):
                print(f"found peak_{part_idx}:{len(peaks[part_idx])}")
            for part_idx in range(0,self.n_pos):
                print(f"peak {self.parts(part_idx)}:")
                for peak in peaks[part_idx]:
                    print(f"peak {self.parts(part_idx)} id:{peak.idx} x:{peak.x} y:{peak.y} score:{peak.score}")
                print()

        candidate_limbs=[[] for limb_idx in range(0,len(self.limbs))]
        for limb_idx,limb in enumerate(self.limbs):
            src_idx,dst_idx=limb
            peak_src_list=peaks[src_idx]
            peak_dst_list=peaks[dst_idx]
            if((len(peak_src_list)==0) or(len(peak_dst_list)==0)):
                continue
            self.debug_print(f"candidating: src:{self.parts(src_idx)} dst:{self.parts(dst_idx)}")
            for peak_src in peak_src_list:
                for peak_dst in peak_dst_list:
                    #calculate paf vector
                    vec_src=np.array([peak_src.y,peak_src.x])
                    vec_dst=np.array([peak_dst.y,peak_dst.x])
                    vec_limb=vec_dst-vec_src
                    lenth=np.sqrt(np.sum(vec_limb**2))
                    if(lenth<1e-12):
                        continue
                    vec_limb=vec_limb/lenth
                    paf_vectors=self.get_paf_vectors(limb_idx,vec_src,vec_dst,paf_map)
                    #calculate criterion
                    criterion1=0
                    scores=0.0
                    for step in range(self.step_paf):
                        score=np.sum(vec_limb*paf_vectors[step])
                        if(score>=self.thres_vec):
                            criterion1+=1
                        scores+=score
                    criterion2=scores/self.step_paf+min(0.0,0.5*conf_map.shape[0]/lenth-1.0)
                    criterion3=(peak_src.score+peak_dst.score)*0.1
                    #filter candidate limbs
                    self.debug_print(f"test start:id-{peak_src.idx} pos-{vec_src} end:id-{peak_dst.idx} pos-{vec_dst} c1:{criterion1} c2:{criterion2} c3:{criterion3}")
                    if(criterion1>self.thres_vec_cnt and criterion2>self.thres_criterion2):
                        candidate_limbs[limb_idx].append(Connection(peak_src.idx,peak_dst.idx,criterion2+criterion3))
        #filter chosen connection
        all_chosen_limbs=[[] for limb_idx in range(0,len(self.limbs))]
        for limb_idx in range(0,len(self.limbs)):
            sort_candidates=candidate_limbs[limb_idx]
            sort_candidates.sort(reverse=True)
            chosen_limbs=all_chosen_limbs[limb_idx]
            for candidate in sort_candidates:
                assigned=False
                for chosen_limb in chosen_limbs:
                    if(chosen_limb.peak_src_id==candidate.peak_src_id):
                        assigned=True
                    if(chosen_limb.peak_dst_id==candidate.peak_dst_id):
                        assigned=True
                    if(assigned):
                        break
                if(assigned):
                    continue
                chosen_limbs.append(candidate)
        #assemble human
        humans=[]
        for limb_idx,limb in enumerate(self.limbs):
            src_part_idx,dst_part_idx=limb
            chosen_limbs=all_chosen_limbs[limb_idx]
            for chosen_limb in chosen_limbs:
                peak_src_id,peak_dst_id=chosen_limb.peak_src_id,chosen_limb.peak_dst_id
                #print(f"assigning chosen_limb score:{chosen_limb.score}")
                touched_ids=[]
                for human_id,human in enumerate(humans):
                    if((human[src_part_idx]==peak_src_id) or (human[dst_part_idx]==peak_dst_id)):
                        touched_ids.append(human_id)
                if(len(touched_ids)==1):
                    human=humans[touched_ids[0]]
                    if(human[dst_part_idx]!=peak_dst_id):
                        #TODO: check why could the followers just take the previous larger score???
                        human[dst_part_idx]=peak_dst_id
                        human[19]+=1
                        human[18]+=all_peaks[peak_dst_id].score+chosen_limb.score
                elif(len(touched_ids)>=2):
                    membership=0
                    human_1=humans[touched_ids[0]]
                    human_2=humans[touched_ids[1]]
                    for part_idx in range(0,18):
                        if(human_1[part_idx]>=0 and human_2[part_idx]>=0):
                            membership=2
                    if(membership==0):
                        human_1[0:18]+=human_2[0:18]+1
                        human_1[18]+=human_2[18]+chosen_limb.score
                        human_1[19]+=human_2[19]
                        humans.pop(touched_ids[1])
                    elif(membership==2):
                        human_1[dst_part_idx]=peak_dst_id
                        human_1[19]+=1
                        human_1[18]+=all_peaks[peak_dst_id].score+chosen_limb.score
                elif(len(touched_ids)==0 and limb_idx<17):
                    human=np.zeros(shape=[20]).astype(np.float32)
                    human[:]=-1
                    human[src_part_idx]=peak_src_id
                    human[dst_part_idx]=peak_dst_id
                    human[18]=all_peaks[peak_src_id].score+all_peaks[peak_dst_id].score+chosen_limb.score
                    human[19]=2
                    humans.append(human)
        #return assembled human
        #print(f"test candidate human:{len(humans)}")
        ret_humans=[]
        for human_id,human in enumerate(humans):
            if(self.debug):
                self.debug_print(f"\ntest human filter human_id:{human_id} score:{human[18]/human[19]}  part_num:{human[19]}")
            #print(f"test candidate")
            for i in range(0,18):
                if(human[i]!=-1):
                    peak=all_peaks[int(human[i])]
                    self.debug_print(f"part:{self.parts(i)} loc_y:{peak.y} loc_x:{peak.x} socre:{peak.score}")
            if((human[18]/human[19]>=self.thres_human_score) and (human[19]>=self.thres_part_cnt)):
                ret_human=Human(self.parts,self.limbs,self.colors)
                ret_human.local_id=human_id
                ret_human.score=human[18]/human[19]
                for part_idx in range(0,self.n_pos-1):
                    if(human[part_idx]!=-1):
                        peak=all_peaks[int(human[part_idx])]
                        x,y,score=peak.x,peak.y,peak.score
                        ret_human.body_parts[part_idx]=BodyPart(parts=self.parts,u_idx=human[part_idx],part_idx=part_idx,\
                            x=x,y=y,score=score)
                ret_human.global_id=self.cur_id
                ret_humans.append(ret_human)
                self.cur_id+=1
        return ret_humans

    def get_paf_vectors(self,limb_id,vec_src,vec_dst,paf_map):
        def round(x):
            sign_x=np.where(x>0,1,-1)
            return (x+0.5*sign_x).astype(np.int)
        paf_vectors=np.zeros(shape=(self.step_paf,2))
        vec_limb=vec_dst-vec_src
        for step in range(0,self.step_paf):
            vec_loc_y,vec_loc_x=round(vec_src+vec_limb*step/self.step_paf)
            vec_paf_x=paf_map[vec_loc_y][vec_loc_x][limb_id*2]
            vec_paf_y=paf_map[vec_loc_y][vec_loc_x][limb_id*2+1]
            paf_vectors[step][0]=vec_paf_y
            paf_vectors[step][1]=vec_paf_x
        return paf_vectors

    def debug_print(self,msg):
        if(self.debug):
            print(msg)

class Peak:
    def __init__(self,peak_idx,part_idx,y,x,score):
        self.idx=peak_idx
        self.part_idx=part_idx
        self.y=y
        self.x=x
        self.score=score

class Connection:
    def __init__(self,peak_src_id,peak_dst_id,score):
        self.peak_src_id=peak_src_id
        self.peak_dst_id=peak_dst_id
        self.score=score

    def __lt__(self,other):
        return self.score<other.score

    def __eq__(self,other):
        return self.score==other.score


import multiprocessing as mp
from queue import Full as QueueFull
from queue import Empty as QueueEmpty
from threading import Thread
import ctypes
import numpy as np
import numpy
from functools import partial
from scipy.signal import fftconvolve
from h5py import File
from DRE.core.models import ModelsCube
from DRE.core.summary import Summary
from DRE.misc.progress_bar import progress
from DRE.misc.read_psf import get_psf
from astropy.io import fits
import os
import time
import datetime

class ModelCPU(ModelsCube):
    """
    A ModelCube like object that can be shared between processes to save memory,
    all attributes are the same as ModelsCube

    Methods
    -------
    to_shared_mem(array)
        sends the array to shared memory
    convolve(psf, n_proc=1, *args, **kwargs)
        convolves the models with the PSF using a pool of CPU workers
    """
    def __init__(self, models_file=None, out_compression='none', save_mosaics=False):
        super().__init__(models_file, out_compression)

        # send all attributes to shared memory, lock is not needed as is read-only
        self.models = self.to_shared_mem(self.models)
        self.log_r = self.to_shared_mem(self.log_r)
        self.ax_ratio = self.to_shared_mem(self.ax_ratio)
        self.angle = self.to_shared_mem(self.angle)
        self.save_mosaics = mp.Value(ctypes.c_bool, save_mosaics, lock=False)

    @staticmethod
    def to_shared_mem(array):
        """
        sends the array to shared memory using a multiprocessing.Array object as buffer,
        the array is assumed read-only so no lock is needed.

        Parameters
        ----------
        array : ndarray
            array to be send to shared memory
        """
        shape = array.shape
        shared_array_base = mp.Array(ctypes.c_float, int(np.prod(shape)), lock=False)
        shared_array = np.ctypeslib.as_array(shared_array_base)
        shared_array = shared_array.reshape(shape)
        shared_array[:] = array
        return shared_array

    def convolve(self, psf, n_proc=1, *args, **kwargs):
        """
        convolves the models with the PSF and stores them in the convolved_models attribute as a shared array,
        uses a multiprocessing.Pool for parallelization

        Parameters
        ----------
        psf : ndarray
            array representing the PSF
        n_proc : int
            number of CPU processes to use
        """

        convolve = partial(fftconvolve, in2=psf, mode='same', axes=(-2, -1))
        # flatten first dimensions e.g. (4, 10, 13, 21, 128, 128) -> (4 * 10 * 13 * 21, 128, 128)
        flatten_shape = (np.prod(self.models.shape[:-2]), * self.models.shape[-2:])
        with mp.Pool(n_proc) as pool:
            # convolve in parallel
            convolved = pool.map(convolve, self.models.reshape(flatten_shape))
        convolved = np.array(list(convolved)).reshape(self.models.shape)
        self.convolved_models = self.to_shared_mem(convolved)

class Parallelize:
    """
    this object manages the parallelization an I/O. It has two sub-threads, one for reading the input and another to
    write the output, and a list of sub-processes to do the fit in parallel. It assumes 'spawn' (Windows, Mac) start
    method instead of 'fork' (the default in Linux) to be multi-thread safe.

    Attributes
    ----------
    n_proc : int
        number of processes to use in the calculation
    processes : list
        list with the processes
    input_queue : multiprocessing.JoinableQueue
        multiprocessing-safe queue to feed the input from the feed_thread to the processes
    output_queue : multiprocessing.JoinableQueue
        multiprocessing-safe queue to feed the output from the processes to the output thread
    feed_thread : threading.Thread
        thread that reads the input files and feeds the input to the processes
    output_thread : threading.Thread
        thread that takes the output of the processes and writes them into a output file
    terminate : multiprocessing.Event
        event to stop the program by user request (Ctrl+C)

    Methods
    -------
    cpu_worker(model, input_queue, output_queue, terminate)
        it's a worker that runs in a new process, it takes an input from the input_queue and uses a ModelCPU to make the
        fit, then puts the output into the output_queue
    feed_workers(names, input_file)
        takes the input file (HD5F file with cuts), and puts the input in the input queue
    get_output(model, input_name, n_tasks, output_file, table, progress_status)
        gets the outputs from the output_queue and writes them into an output file
    start_processes(self, model, input_name, names, input_file, output_file, table, progress_status)
        starts the threads and processes
    stop_processes()
        joins the threads and processes
    abort()
        stops the processes by  user request (Ctrl+C)
    fit_file(model, input_name, input_file, output_file, psf, cats_dir, progress_status='')
        fits all the objects on an input file (HD5F file with cuts)
    fit_dir(model, input_dir='Cuts', output_dir='Chi', psf_dir='PSF', cats_dir='Sextracted')
        apply fit_file to each file in the input directory
    """

    def __init__(self, n_proc, max_size=100):
        self.n_proc = n_proc
        self.processes = []
        self.input_queue = mp.JoinableQueue(max_size)
        self.output_queue = mp.JoinableQueue()
        self.feed_thread = None
        self.output_thread = None
        self.terminate = mp.Event()

    @staticmethod
    def cpu_worker(model, input_queue, output_queue, terminate):
        try:
            for name, data, segment, noise in iter(input_queue.get, 'STOP'):
                chi_cube = model.dre_fit(data, segment, noise, backend=numpy)
                parameters = None
                mosaic = None
                success = not np.isnan(chi_cube).all()
                if success:
                    parameters = model.get_parameters(chi_cube)
                    ext, numb = name.split('_')
                    parameters['EXT_NUMBER'] = int(ext)
                    parameters['NUMBER'] = int(numb)
                    if model.save_mosaics:
                        mosaic = model.make_mosaic(data, segment, parameters['MODEL_IDX'])
                output_queue.put((name, success, chi_cube, parameters, mosaic))
                input_queue.task_done()
        except KeyboardInterrupt:
            terminate.set()

    def feed_workers(self, names, input_file):
        submitted = 0
        while submitted < len(names) and not self.terminate.is_set():
            with File(input_file, 'r') as input_h5f:
                name = names[submitted]
                data = input_h5f[name]
                # subscribes input to workers queue
                try:
                    self.input_queue.put((name, data['obj'][:], data['seg'][:], data['rms'][:]), timeout=2)
                    submitted += 1
                except QueueFull:
                    pass
        if not self.terminate.is_set():
            for i in range(self.n_proc):
                self.input_queue.put('STOP')

    def get_output(self, model, input_name, n_tasks, output_file, table, progress_status):
        completed = 0
        while completed < n_tasks and not self.terminate.is_set():
            try:
                name, success, chi_cube, params, mosaic = self.output_queue.get(timeout=2)
                if success:
                    table.append(params)
                    with File(output_file, 'a') as output_h5f:
                        output_h5f.create_dataset(f'{name}', data=chi_cube,
                                                  dtype='float32', **model.compression)
                    if model.save_mosaics:
                        os.makedirs(os.path.join("Mosaics", input_name), exist_ok=True)
                        mosaic_fits = fits.ImageHDU(data=mosaic)
                        mosaic_fits.writeto(os.path.join("Mosaics", input_name, f"{input_name}_{name}_mosaic.fits"),
                                            overwrite=True)
                self.output_queue.task_done()
                completed += 1
                progress(completed, n_tasks, progress_status)
            except QueueEmpty:
                pass

    def start_processes(self, model, input_name, names, input_file, output_file, table, progress_status):
        self.feed_thread = Thread(target=self.feed_workers, args=(names, input_file))
        self.output_thread = Thread(target=self.get_output, args=(model, input_name, len(names), output_file,
                                                                  table, progress_status))
        self.feed_thread.start()
        self.output_thread.start()
        self.processes = []
        for i in range(self.n_proc):
            p = mp.Process(target=self.cpu_worker, args=(model, self.input_queue, self.output_queue, self.terminate))
            self.processes.append(p)
            p.start()

    def stop_processes(self):
        if self.feed_thread.is_alive():
            self.feed_thread.join()
        for p in self.processes:
            if p.is_alive():
                p.join()
        self.output_thread.join()

    def abort(self):
        self.terminate.set()
        self.input_queue.cancel_join_thread()
        self.output_queue.cancel_join_thread()

    def fit_file(self, model, input_name, input_file, chi_file, output_dir, psf, cats_dir,
                 convolve=True, progress_status=''):
        with File(input_file, 'r') as input_h5f:
            names = list(input_h5f.keys())
        print(f"{progress_status}: {input_file}\t{len(names)} objects")
        # table with summary
        table = Summary(input_name)
        if convolve:
            # convolve with the psf
            print(f"{progress_status}: Convolving...")
            convolve_start = time.time()
            try:
                model.convolve(get_psf(psf), n_proc=self.n_proc)
            except FileNotFoundError:
                print(f"{progress_status}: Warning: Cannot find the PSF file {psf},\n\t skipping this tile")
                return
            print(f"{progress_status}: Convolved! ({time.time() - convolve_start:2.2f}s)")
        # fit in parallel
        job_start = time.time()
        try:
            self.start_processes(model, input_name, names, input_file, chi_file, table, progress_status)
            self.stop_processes()
            # save summary
            table.save(output_dir, cats_dir)
            time_delta = datetime.timedelta(seconds=(time.time() - job_start))
            obj_s = len(names) / (time.time() - job_start)
            print(f"\n{progress_status}: finished in {str(time_delta)[:10]}s ({obj_s:1.3f} obj/s)")
        except KeyboardInterrupt:
            print("\nAborted by user request")
            self.abort()

    def fit_dir(self, model, input_dir='Cuts', output_dir='Summary', chi_dir='Chi',
                psf_dir='PSF', cats_dir='Sextracted', convolve=True):
        start = time.time()
        # list with input files in input_dir
        files = os.listdir(input_dir)
        os.makedirs(chi_dir, exist_ok=True)
        for i, filename in enumerate(sorted(files)):
            input_file = os.path.join(input_dir, filename)
            name = os.path.basename(filename).replace('_cuts.h5', '')
            chi_file = os.path.join(chi_dir, f"{name}_chi.h5")
            psf = os.path.join(psf_dir, f"{name}.psf")
            if os.path.isfile(chi_file):
                os.remove(chi_file)
            # fit all cuts in each file
            self.fit_file(model, name, input_file, chi_file, output_dir, psf, cats_dir, convolve,
                          progress_status=f"({i + 1}/{len(files)})")
            if self.terminate.is_set():
                break
        time_delta = datetime.timedelta(seconds=(time.time() - start))
        print(f"DRE finished, total time: {str(time_delta)[:10]}s")


# -*- coding: utf-8 -*-
import numpy as np #numpy is used for a parameter input
from steganogan import SteganoGAN
from steganogan.loader import DataLoader
from steganogan.encoders import BasicEncoder, DenseEncoder
from steganogan.decoders import BasicDecoder, DenseDecoder
from steganogan.critics import BasicCritic

# Load the data
train = DataLoader('D:/dataset/train/', limit=np.inf, shuffle=True, batch_size=4)
validation = DataLoader('D:/dataset/val/', limit=np.inf, shuffle=True, batch_size=4)

# Create the SteganoGAN instance
steganogan = SteganoGAN(1, BasicEncoder, BasicDecoder, BasicCritic, hidden_size=32, cuda=True, verbose=True)

# Fit on the given data
# steganogan.fit(train, validation, epochs=5)

steganogan = SteganoGAN.load(architecture='demo', path=None, cuda=True, verbose=True)

# Encode a message in input.png
steganogan.encode('research/input.png', 'output.png', 'This is a super secret message')

# Decode the message from output.png
steganogan.decode('research/output.png')


from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import (
    Any,
    Dict,
    List,
    Mapping,
    Optional,
    Type,
    Union,
    overload,
)
from entity_gym.environment.environment import (
    ActionSpace,
    CategoricalActionSpace,
    Environment,
    EpisodeStats,
    ObsSpace,
    Observation,
    SelectEntityActionSpace,
    EntityType,
    ActionType,
)
import numpy as np
import numpy.typing as npt
from ragged_buffer import RaggedBufferF32, RaggedBufferI64, RaggedBufferBool

@dataclass
class VecSelectEntityActionMask:
    actors: RaggedBufferI64
    actees: RaggedBufferI64

    @overload
    def __getitem__(self, i: int) -> RaggedBufferI64:
        ...

    @overload
    def __getitem__(self, i: npt.NDArray[np.int64]) -> "VecSelectEntityActionMask":
        ...

    def __getitem__(
        self, i: Union[int, npt.NDArray[np.int64]]
    ) -> Union["VecSelectEntityActionMask", RaggedBufferI64]:
        if isinstance(i, int):
            return self.actors[i]
        else:
            return VecSelectEntityActionMask(self.actors[i], self.actees[i])

    def extend(self, other: Any) -> None:
        assert isinstance(
            other, VecSelectEntityActionMask
        ), f"Expected VecSelectEntityActionMask, got {type(other)}"
        self.actors.extend(other.actors)
        self.actees.extend(other.actees)

    def clear(self) -> None:
        self.actors.clear()
        self.actees.clear()

@dataclass
class VecCategoricalActionMask:
    actors: RaggedBufferI64
    mask: Optional[RaggedBufferBool]

    def __getitem__(
        self, i: Union[int, npt.NDArray[np.int64]]
    ) -> "VecCategoricalActionMask":
        if self.mask is not None and self.mask.size0() > 0:
            return VecCategoricalActionMask(self.actors[i], self.mask[i])
        else:
            return VecCategoricalActionMask(self.actors[i], None)

    def extend(self, other: Any) -> None:
        assert isinstance(
            other, VecCategoricalActionMask
        ), f"Expected CategoricalActionMaskBatch, got {type(other)}"
        if self.mask is not None and other.mask is not None:
            self.mask.extend(other.mask)
        elif self.mask is None and other.mask is None:
            pass
        elif self.mask is not None:
            self.mask.extend(
                RaggedBufferBool.from_flattened(
                    flattened=np.ones(
                        shape=(other.actors.items(), self.mask.size2()),
                        dtype=np.bool_,
                    ),
                    lengths=other.actors.size1(),
                )
            )
        elif other.mask is not None:
            self.mask = RaggedBufferBool.from_flattened(
                flattened=np.ones(
                    shape=(self.actors.items(), other.mask.size2()),
                    dtype=np.bool_,
                ),
                lengths=self.actors.size1(),
            )
            self.mask.extend(other.mask)
        else:
            raise Exception("Impossible!")
        self.actors.extend(other.actors)

    def clear(self) -> None:
        self.actors.clear()
        if self.mask is not None:
            self.mask.clear()

VecActionMask = Union[VecCategoricalActionMask, VecSelectEntityActionMask]

@dataclass
class VecObs:
    features: Dict[EntityType, RaggedBufferF32]
    action_masks: Dict[ActionType, VecActionMask]
    reward: npt.NDArray[np.float32]
    done: npt.NDArray[np.bool_]
    end_of_episode_info: Dict[int, EpisodeStats]

    def extend(self, b: "VecObs") -> None:
        num_envs = len(self.reward)
        for etype, feats in b.features.items():
            if etype not in self.features:
                self.features[etype] = empty_ragged_f32(
                    feats=feats.size2(), sequences=num_envs
                )
            self.features[etype].extend(feats)
        for atype, amask in b.action_masks.items():
            if atype not in self.action_masks:
                raise NotImplementedError()
            else:
                self.action_masks[atype].extend(amask)
        self.reward = np.concatenate((self.reward, b.reward))
        self.done = np.concatenate((self.done, b.done))
        num_envs = len(self.reward)
        for i, stats in b.end_of_episode_info.items():
            self.end_of_episode_info[i + num_envs] = stats

class VecEnv(ABC):
    @abstractmethod
    def env_cls(cls) -> Type[Environment]:
        """
        Returns the class of the underlying environment.
        """
        raise NotImplementedError

    @abstractmethod
    def reset(self, obs_config: ObsSpace) -> VecObs:
        raise NotImplementedError

    @abstractmethod
    def act(
        self, actions: Mapping[ActionType, RaggedBufferI64], obs_filter: ObsSpace
    ) -> VecObs:
        raise NotImplementedError

    @abstractmethod
    def render(self, **kwargs: Any) -> npt.NDArray[np.uint8]:
        raise NotImplementedError

    @abstractmethod
    def __len__(self) -> int:
        raise NotImplementedError

    def close(self) -> None:
        pass

def batch_obs(
    obs: List[Observation], obs_space: ObsSpace, action_space: Dict[str, ActionSpace]
) -> VecObs:
    """
    Converts a list of observations into a batch of observations.
    """
    features: Dict[EntityType, RaggedBufferF32] = {}
    action_masks: Dict[ActionType, VecActionMask] = {}
    reward = []
    done = []
    end_of_episode_info = {}

    # Initialize the entire batch with all entities and actions
    for entity_name, entity in obs_space.entities.items():
        feature_size = len(entity.features)
        features[entity_name] = RaggedBufferF32(feature_size)

    for action_name, space in action_space.items():
        if isinstance(space, CategoricalActionSpace):
            action_masks[action_name] = VecCategoricalActionMask(
                RaggedBufferI64(1),
                None,
            )
        elif isinstance(space, SelectEntityActionSpace):
            action_masks[action_name] = VecSelectEntityActionMask(
                RaggedBufferI64(1), RaggedBufferI64(1)
            )

    for i, o in enumerate(obs):
        for entity_type, entity in obs_space.entities.items():
            if entity_type not in features:
                features[entity_type] = RaggedBufferF32.from_flattened(
                    np.zeros((0, len(entity.features)), dtype=np.float32),
                    lengths=np.zeros(i, dtype=np.int64),
                )
            if entity_type in o.features:
                ofeats = o.features[entity_type]
                if not isinstance(ofeats, np.ndarray):
                    ofeats = np.array(ofeats, dtype=np.float32).reshape(
                        len(ofeats), len(obs_space.entities[entity_type].features)
                    )
                features[entity_type].push(ofeats)
            else:
                features[entity_type].push(
                    np.zeros((0, len(entity.features)), dtype=np.float32)
                )

        for atype, space in action_space.items():
            if atype not in o.actions:
                if atype in action_masks:
                    if isinstance(space, CategoricalActionSpace):
                        vec_action = action_masks[atype]
                        assert isinstance(vec_action, VecCategoricalActionMask)
                        vec_action.actors.push(np.zeros((0, 1), dtype=np.int64))
                        if vec_action.mask is not None:
                            vec_action.mask.push(
                                np.zeros((0, len(space.choices)), dtype=np.bool_)
                            )
                    elif isinstance(space, SelectEntityActionSpace):
                        vec_action = action_masks[atype]
                        assert isinstance(vec_action, VecSelectEntityActionMask)
                        vec_action.actors.push(np.zeros((0, 1), dtype=np.int64))
                        vec_action.actees.push(np.zeros((0, 1), dtype=np.int64))
                continue
            action = o.actions[atype]
            if atype not in action_masks:
                if isinstance(space, CategoricalActionSpace):
                    action_masks[atype] = VecCategoricalActionMask(
                        empty_ragged_i64(1, i), None
                    )
                elif isinstance(space, SelectEntityActionSpace):
                    action_masks[atype] = VecSelectEntityActionMask(
                        empty_ragged_i64(1, i), empty_ragged_i64(1, i)
                    )
            if isinstance(space, CategoricalActionSpace):
                vec_action = action_masks[atype]
                assert isinstance(vec_action, VecCategoricalActionMask)
                actor_indices = o._actor_indices(atype, obs_space)
                vec_action.actors.push(actor_indices.reshape(-1, 1))
                if action.mask is not None:
                    if vec_action.mask is None:
                        vec_action.mask = RaggedBufferBool.from_flattened(
                            np.ones((0, len(space.choices)), dtype=np.bool_),
                            np.zeros(i, dtype=np.int64),
                        )
                    amask = action.mask
                    if not isinstance(amask, np.ndarray):
                        amask = np.array(amask, dtype=np.bool_)
                    vec_action.mask.push(amask)
                elif vec_action.mask is not None:
                    vec_action.mask.push(
                        np.ones(
                            (len(actor_indices), len(space.choices)), dtype=np.bool_
                        )
                    )
            elif isinstance(space, SelectEntityActionSpace):
                vec_action = action_masks[atype]
                assert isinstance(vec_action, VecSelectEntityActionMask)
                actors = o._actor_indices(atype, obs_space).reshape(-1, 1)
                vec_action.actors.push(actors)
                if len(actors) > 0:
                    vec_action.actees.push(
                        o._actee_indices(atype, obs_space).reshape(-1, 1)
                    )
                else:
                    vec_action.actees.push(np.zeros((0, 1), dtype=np.int64))
            else:
                raise NotImplementedError()

        reward.append(o.reward)
        done.append(o.done)
        if o.end_of_episode_info:
            end_of_episode_info[len(reward) - 1] = o.end_of_episode_info

    return VecObs(
        features,
        action_masks,
        np.array(reward, dtype=np.float32),
        np.array(done, dtype=np.bool_),
        end_of_episode_info,
    )

def empty_ragged_f32(feats: int, sequences: int) -> RaggedBufferF32:
    return RaggedBufferF32.from_flattened(
        np.zeros((0, feats), dtype=np.float32),
        lengths=np.array([0] * sequences, dtype=np.int64),
    )

def empty_ragged_i64(feats: int, sequences: int) -> RaggedBufferI64:
    return RaggedBufferI64.from_flattened(
        np.zeros((0, feats), dtype=np.int64),
        lengths=np.array([0] * sequences, dtype=np.int64),
    )


import sqlalchemy as sa
from sqlalchemy.orm import relationship
from sqlalchemy import func, Index
from sqlalchemy.dialects import postgresql as psql

from .core import DBSession, Base
from .utils import fid_map
from .image import CalibratableImage
from .photometry import ForcedPhotometry
from .detections import Detection
from .subtraction import SingleEpochSubtraction
from .constants import MJD_TO_JD
from .thumbnails import Thumbnail
from .spatial import SpatiallyIndexed

import numpy as np

__all__ = ['Source']

class Source(Base, SpatiallyIndexed):
    id = sa.Column(sa.String, primary_key=True)
    # TODO should this column type be decimal? fixed-precison numeric
    ra = sa.Column(sa.Float)
    dec = sa.Column(sa.Float)

    offset = sa.Column(sa.Float, default=0.0)
    redshift = sa.Column(sa.Float, nullable=True)
    neighbor_info = sa.Column(psql.JSONB)

    altdata = sa.Column(psql.JSONB, nullable=True)
    score = sa.Column(sa.Float, nullable=True)
    accumulated_rb = sa.Column(sa.Float)

    thumbnails = relationship('Thumbnail', cascade='all')
    detections = relationship('Detection', cascade='all')
    forced_photometry = relationship('ForcedPhotometry', cascade='all')

    def add_linked_thumbnails(self):
        sdss_thumb = Thumbnail(public_url=self.sdss_url,
                               type='sdss')
        dr8_thumb = Thumbnail(public_url=self.desi_dr8_url,
                              type='dr8')
        DBSession().add_all([sdss_thumb, dr8_thumb])
        DBSession().commit()

    @property
    def sdss_url(self):
        """Construct URL for public Sloan Digital Sky Survey (SDSS) cutout."""
        return (f"http://skyservice.pha.jhu.edu/DR9/ImgCutout/getjpeg.aspx"
                f"?ra={self.ra}&dec={self.dec}&scale=0.3&width=200&height=200"
                f"&opt=G&query=&Grid=on")

    @property
    def desi_dr8_url(self):
        """Construct URL for public DESI DR8 cutout."""
        return (f"http://legacysurvey.org/viewer/jpeg-cutout?ra={self.ra}"
                f"&dec={self.dec}&size=200&layer=dr8&pixscale=0.262&bands=grz")

    def images(self, type=CalibratableImage):

        candidates = DBSession().query(type).filter(
            func.q3c_radial_query(type.ra,
                                  type.dec,
                                  self.ra, self.dec,
                                  0.64)
        ).filter(
            func.q3c_poly_query(self.ra, self.dec, type.poly)
        )

        return candidates.all()

    @property
    def best_detection(self):
        return DBSession().query(
            Detection
        ).join(Source).filter(
            Source.id == self.id
        ).order_by(
            (Detection.flux / Detection.fluxerr).desc()
        ).first()

    @property
    def light_curve(self):
        from astropy.table import Table
        lc_raw = []

        phot = DBSession().query(
            ForcedPhotometry.obsjd - MJD_TO_JD,
            ForcedPhotometry.filtercode,
            ForcedPhotometry.zp,
            ForcedPhotometry.flux,
            ForcedPhotometry.fluxerr,
            ForcedPhotometry.flags,
            ForcedPhotometry.id
        ).filter(
            ForcedPhotometry.source_id == self.id
        )

        for photpoint in phot:
            photd = {'mjd': photpoint[0],
                     'filter': 'ztf' + photpoint[1][-1],
                     'zp': photpoint[2],  # ap-correct the ZP
                     'zpsys': 'ab',
                     'flux': photpoint[3],
                     'fluxerr': photpoint[4],
                     'flags': photpoint[5],
                     'lim_mag': -2.5 * np.log10(5 * photpoint[4]) + photpoint[2],
                     'id': photpoint[-1]}
            lc_raw.append(photd)

        return Table(lc_raw)

    @property
    def unphotometered_images(self):
        subq = DBSession().query(ForcedPhotometry.id,
                                 ForcedPhotometry.image_id).filter(
            ForcedPhotometry.source_id == self.id
        ).subquery()

        q = DBSession().query(SingleEpochSubtraction).filter(
            func.q3c_radial_query(SingleEpochSubtraction.ra,
                                  SingleEpochSubtraction.dec,
                                  self.ra, self.dec,
                                  0.64)
        ).filter(
            func.q3c_poly_query(self.ra, self.dec, SingleEpochSubtraction.poly)
        ).outerjoin(
            subq, subq.c.image_id == SingleEpochSubtraction.id
        ).filter(
            subq.c.id == None
        )

        return q.all()

    def force_photometry(self, assume_background_subtracted=True):
        out = []
        for i in self.unphotometered_images:
            sci_path = f'/global/cscratch1/sd/dgold/zuds/{i.field:06d}/' \
                       f'c{i.ccdid:02d}/q{i.qid}/{fid_map[i.fid]}/' \
                       f'{i.basename}'

            mask_path = sci_path.replace('.fits', '.mask.fits')
            rms_path = sci_path.replace('.fits', '.rms.fits')

            fp = i.force_photometry(
                self, assume_background_subtracted=assume_background_subtracted,
                use_cutout=True, direct_load={'mask': mask_path, 'sci': sci_path,
                                              'rms': rms_path}
            )

            out.extend(fp)
        return out


import os
import sys
import multiprocessing
import cv2
import numpy as np
from PIL import Image
from openslide import OpenSlide, deepzoom
import matplotlib.pyplot as plt
from skimage import io
from tqdm import tqdm

import torch
import torch.nn.functional as F

sys.path.append('../GLNet-pytorch/')
from GLNet.models.glnet import GLNet as glnet
from GLNet.utils.enums import PhaseMode
from GLNet.utils.processing import (
    images_transform,
    resize,
    global2patch,
    patch2global,
    add_extra_pixels
)
from GLNet.utils.filters import apply_filters
from GLNet.models.functional import crop_global
from GLNet.utils.visualization import apply_mask
from GLNet.options import InferenceOptions

DATA_DIR = '/home/hainq/quanghai/Thesis/datasets/paip2019/training/'

def load_models(n_class, path_l2g, path_g):
    model = glnet(n_class)
    # model = nn.DataParallel(model)
    model.cuda()
    partial = torch.load(path_l2g)
    # partial = torch.load('../../checkpoints/fpn_mode2_paip.1008_lr1e4.pth')
    state = model.state_dict()
    pretrained_dict = {
        k: v for k, v in partial.items() if k in state
    }
    state.update(pretrained_dict)
    model.load_state_dict(state)
    model.eval()

    global_fixed = glnet(2)
    # global_fixed = nn.DataParallel(global_fixed)
    global_fixed.cuda()
    partial = torch.load(path_g)
    state = global_fixed.state_dict()
    pretrained_dict = {
        k: v for k, v in partial.items() if k in state and "local" not in k
    }
    # 2. overwrite entries in the existing state dict
    state.update(pretrained_dict)
    # 3. load the new state dict
    global_fixed.load_state_dict(state)
    global_fixed.eval()

    return model, global_fixed

def predict_on_single_patch(image_as_tensor, model, global_fixed, size_g=(1008, 1008), size_p=(1008, 1008), n_class=2):
    # image_as_tensor[0] = apply_filters(image_as_tensor[0])
    images_global = resize(image_as_tensor, size_g)
    scores = [
        np.zeros((1, n_class, image_as_tensor[i].size[1], image_as_tensor[i].size[0]))
        for i in range(len(image_as_tensor))
    ]

    images_glb = images_transform(images_global)

    patches, coordinates, templates, sizes, ratios = global2patch(
        image_as_tensor, size_p
    )

    predicted_ensembles = [
        np.zeros((len(coordinates[i]),n_class,size_p[0],size_p[1]))
        for i in range(len(image_as_tensor))
    ]

    for i in range(len(image_as_tensor)):
        j = 0
        while j < len(coordinates[i]):
            patches_var = images_transform(
                patches[i][j : j + 1]
            )  # b, c, h, w

            fm_patches, _ = model.module.collect_local_fm(
                images_glb[i : i + 1],
                patches_var,
                ratios[i],
                coordinates[i],
                [j, j + 1],
                len(image_as_tensor),
                global_model=global_fixed,
                template=templates[i],
                n_patch_all=len(coordinates[i]),
            )
            j += 1

    _, fm_global = model.forward(
        images_glb, None, None, None, mode=PhaseMode.GlobalFromLocal
    )

    for i in range(len(image_as_tensor)):
        j = 0
        # while j < n ** 2:
        while j < len(coordinates[i]):
            fl = fm_patches[i][j : j + 1].cuda()
            fg = crop_global(
                fm_global[i : i + 1],
                coordinates[i][j : j + 1],
                ratios[i])[0]
            fg = F.interpolate(
                fg, size=fl.size()[2:], mode="bilinear"
            )
            output_ensembles = model.module.ensemble(
                fl, fg
            )  # include cordinates
            # output_ensembles = F.interpolate(model.module.ensemble(fl, fg), size_p, **model.module._up_kwargs)

            # ensemble predictions
            predicted_ensembles[i][j : j + output_ensembles.size()[0]] += (
                F.interpolate(
                    output_ensembles,
                    size=size_p,
                    mode="nearest",
                )
                .data.cpu()
                .numpy()
            )
            j += 1

        scores[i] += np.rot90(
            np.array(
                patch2global(
                    predicted_ensembles[i : i + 1],
                    n_class,
                    sizes[i : i + 1],
                    coordinates[i : i + 1],
                    size_p,
                )
            ),
            k=0,
            axes=(3, 2),
        )

    predictions = [score.argmax(1)[0] for score in scores]
    return predictions

def predict_on_row(dz_obj, row, cols, model, global_fixed):
    for col in cols:
        tile = dz.get_tile(dz.level_count - 1, (col, row))
        if tile.size != (4096, 4096):
            tile = add_extra_pixels(tile, expected_shape=(4096, 4096))

def predict_on_svs(img_idx, model, global_fixed, save_pred=False):
    slide = OpenSlide(os.path.join(DATA_DIR, 'ws_images', f'{img_idx}.svs'))
    mask = io.imread(os.path.join(DATA_DIR, 'viable_masks', f'{img_idx}_viable.tif'))
    dz = deepzoom.DeepZoomGenerator(slide, tile_size=4096, overlap=0)
    img = slide.read_region((0, 0), 2, slide.level_dimensions[2])
    img = img.convert('RGB')
    img = np.array(img)

    cols, rows = dz.level_tiles[-1]

    num_processes = multiprocessing.cpu_count()
    if num_processes > cols:
        num_processes = cols

    # out = np.zeros_like(img)
    out = np.zeros((rows * 256, cols * 256, 3), dtype=np.uint8)

    for row in range(rows):
        for col in tqdm(range(cols)):
            tile = dz.get_tile(dz.level_count - 1, (col, row)) # col, row
            if tile.size != (4096, 4096):
                tile = add_extra_pixels(tile, expected_shape=(4096, 4096))
            pred = predict_on_single_patch([tile], model, global_fixed)
            pil_pred = Image.fromarray(pred[0].astype(np.uint8))

            pil_pred = pil_pred.resize((pil_pred.size[0] // 16, pil_pred.size[1] // 16), Image.NEAREST)
            tile = tile.resize((tile.size[0] // 16, tile.size[1] // 16), Image.BILINEAR).convert('RGB')

            pil_pred = np.array(pil_pred)
            tile = np.array(tile)

            applied_mask = apply_mask(tile, pil_pred)
    #         plt.imshow(applied_mask)

            edge_right = (col + 1) * tile.shape[0]
            edge_bot = (row + 1) * tile.shape[1]
    #         if edge_right > out.shape[0]:
    #             edge_right = out.shape[0]
    #         if edge_bot > out.shape[1]:
    #             edge_bot = out.shape[1]

            out[row * tile.shape[1] : edge_bot, \
                col * tile.shape[0] : edge_right, :] = applied_mask

    mask = cv2.resize(mask, (mask.shape[1] // 16, mask.shape[0] // 16))
    applied_mask = apply_mask(np.array(img), mask)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    ax1.axis('off')
    ax1.imshow(applied_mask)
    ax1.set_title('Ground Truth', fontsize=16)
    ax2.axis('off')
    ax2.imshow(out)
    ax2.set_title('Prediction', fontsize=16)
    if save_pred:
        fig.savefig(f'Pred_{img_idx}.png')
    return out

if __name__ == "__main__":
    args = InferenceOptions("Inference options").parse()
    model, global_fixed = load_models(args.n_class, args.path_l2g, args.path_g)
    predict_on_svs(args.img_idx, model, global_fixed, args.save_pred)


import numpy as np
import sounddevice as sd

MME = 0
sd.default.channels = 2
sd.default.dtype = 'int16'
sd.default.latency = 'low'
sd.default.samplerate = 48000

class PCMStream:
    def __init__(self):
        self.stream = None

    def read(self, num_bytes):
        # frame is 4 bytes
        frames = int(num_bytes / 4)
        data = self.stream.read(frames)[0]

        # convert to pcm format
        return data.tobytes()

    def change_device(self, num):
        if (self.stream is not None):
            self.stream.stop()
            self.stream.close()

        self.stream = sd.InputStream(device=num)
        self.stream.start()

def query_devices():
    index = 0
    options = {}

    for item in sd.query_devices():
        # pip version only supports MME api
        if (item.get('max_input_channels') > 0 and item.get('hostapi') == MME):
            options[item.get('name')] = index

        index += 1

    return options