# Kyle McCormick
# 2015-04-29

################################################################################
# GENERAL
################################################################################

import pyaudio
import wave
import math
import numpy
import struct
import threading
import time

pa = None
chunk_size = 256

def init():
    global pa
    if pa is None:
        pa = pyaudio.PyAudio()

def terminate():
    global pa
    if pa:
        pa.terminate()
        pa = None

def set_chunk_size(val):
    assert val > 0
    global chunk_size
    chunk_size = val

class LibraryError(Exception):
    pass

################################################################################
# SIGNALS
################################################################################

class Signal(object):
    def __init__(self, frame_rate, loop):
        assert frame_rate > 0
        self.frame_rate = frame_rate
        self.loop = loop
        self._is_playing = False
        self._stop_requested = False
        self._effects = []
    def play(self):
        th = threading.Thread(target=self._play_loop, args=())
        self._stop_requested = False
        th.start()
    def play_and_wait(self):
        self.play()
        while self._is_playing:
            time.sleep(0.001)
    def pause(self):
        self._stop_requested = True
        while self._is_playing:
            time.sleep(0.001)
    def is_playing(self):
        return self._is_playing
    def rewind(self):
        raise LibraryError('Signal subclass failed to override rewind')
    def add_effect(self, e):
        self._effects.append(e)
    def remove_effect(self, e):
        if self._effects.count(e) > 0:
            self._effects.remove(e)
    def _play_loop(self):
        self._is_playing = True
        stream = pa.open(format=pyaudio.paInt16,
                         channels=2,
                         rate=self.frame_rate,
                         output=True)
        left, right = numpy.zeros((chunk_size,), dtype=numpy.int16), \
                      numpy.zeros((chunk_size,), dtype=numpy.int16)
        has_more_data = self._get_samples(left, right)
        while has_more_data and not self._stop_requested:
            for i in range(0, len(left)):
                for e in self._effects:
                    e.apply_to_samples(i, left, right, self.frame_rate)
            data = ''
            for f in range(0, left.shape[0]):
                data += struct.pack('h', left[f])
                data += struct.pack('h', right[f])
            stream.write(data)
            left [range(0, chunk_size)] = 0
            right[range(0, chunk_size)] = 0
            has_more_data = self._get_samples(left, right)
        stream.stop_stream()
        stream.close()
        self._is_playing = False
    def _get_samples(self, left, right):
        raise LibraryError('Signal subclass failed to override _get_samples')

class SignalFromAudioFile(Signal):
    def __init__(self, fpath, loop=False):
        assert fpath != None
        self._infile = wave.open(fpath, 'rb')
        Signal.__init__(self, self._infile.getframerate(), loop)
        fmt = pa.get_format_from_width(self._infile.getsampwidth())
        if fmt != pyaudio.paInt16:
            raise IOError('can only load files with 16-bit integer encoding')
        self._num_chans = self._infile.getnchannels()
        if not self._num_chans in [1, 2]:
            raise IOError('can only load files with 1 or 2 channels')
        self._is_rewinding = False
        self._fpath = fpath
    def rewind(self):
        self._is_rewinding = True
        if self._infile:
            self._infile.close()
        self._infile = wave.open(self._fpath, 'rb')
        self._is_rewinding = False
    def _get_samples(self, left, right):
        if self._is_rewinding:
            return
        data = self._infile.readframes(left.shape[0])
        if data == '':
            if self.loop:
                self.rewind()
                data = self._infile.readframes(left.shape[0])
            else:
                self._infile.close()
                self._infile = None
                return None
        frames_read = len(data) / 2 / self._num_chans
        for ix_frame in range(0, frames_read):
            ix_sample = ix_frame * self._num_chans
            str1 = data[ix_sample*2+0] + data[ix_sample*2+1]
            val1 = struct.unpack('h', str1)[0]
            left[ix_frame] = val1
            if self._num_chans == 1:
                right[ix_frame] = val1
            else:
                str2 = data[ix_sample*2+2] + data[ix_sample*2+3]
                val2 = struct.unpack('h', str2)[0]
                right[ix_frame] = val2
        return frames_read > 0

class WaveType(object):
    SINE = 0
    SQUARE = 1
    TRIANGLE = 2
    SAWTOOTH = 3
    _MIN = SINE
    _MAX = SAWTOOTH

class Wave(Signal):
    def __init__(self, wave_type, freq, amplitude, loop=True):
        Signal.__init__(self, 44100, loop)
        assert WaveType._MIN <= wave_type <= WaveType._MAX
        assert freq > 0
        self.wave_type = wave_type
        self.freq = freq
        self.amplitude = amplitude
        self._frames_per_cycle = self.frame_rate / self.freq
        self._ix_frame = 0
    def rewind(self):
        self._ix_frame = 0
    def _calc_sample_val(self):
        d = self._ix_frame / float(self._frames_per_cycle)
        val = None
        if self.wave_type == WaveType.SINE:
            val = math.sin(2 * math.pi * d)
        elif self.wave_type == WaveType.SQUARE:
            val = (1 if d < 0.5 else -1)
        elif self.wave_type == WaveType.TRIANGLE:
            val = (d if d < 0.5 else (1 - d)) * 4 - 1
        elif self.wave_type == WaveType.SAWTOOTH:
            val = d * 2 - 1
        return self.amplitude * val
    def _get_samples(self, left, right):
        for i in range(0, left.shape[0]):
            left[i] = right[i] = self._calc_sample_val()
            self._ix_frame += 1
            if self._ix_frame >= self._frames_per_cycle:
                if self.loop:
                    self._ix_frame = 0
                else:
                    return False
        return True

class CompositeSignal(Signal):
    def __init__(self, *signals):
        assert len(signals) > 0
        fr = None
        for s in signals:
            if fr:
                assert fr == s.frame_rate
            fr = s.frame_rate
        Signal.__init__(self, fr, False)
        self.signals = list(signals)
        self.num_signals = len(signals)
    def rewind(self):
        for sig in self.signals:
            sig.rewind()
    def add_signal(self, signal):
        self.signals.append(signal)
        self.num_signals += 1
    def _get_samples(self, left, right):
        sub_left  = numpy.zeros((len(left),),  dtype=numpy.int16)
        sub_right = numpy.zeros((len(right),), dtype=numpy.int16)
        has_more_data = False
        for sig in self.signals:
            has_more_data = sig._get_samples(sub_left, sub_right) or has_more_data
            for i in range(0, len(sub_left)):
                for e in sig._effects:
                    e.apply_to_samples(i, sub_left, sub_right, sig.frame_rate)
            left  += sub_left  #/ self.num_signals
            right += sub_right #/ self.num_signals
        return has_more_data

################################################################################
# EFFECTS
################################################################################

class Effect(object):
    def apply_to_samples(self, i, left, right, frame_rate):
        raise LibraryError('Effect subclass failed to override apply_to_samples')

class Amplifier(Effect):
    def __init__(self, gain):
        self.gain = gain
    def apply_to_samples(self, i, left, right, frame_rate):
        left[i]  *= self.gain
        right[i] *= self.gain

class Fader(Effect):
    def __init__(self, left_gain, right_gain):
        self.left_gain = left_gain
        self.right_gain = right_gain
    def apply_to_samples(self, i, left, right, frame_rate):
        left[i]  *= self.left_gain
        right[i] *= self.right_gain

class ComplexFader(Effect):
    def __init__(self, l_to_l, l_to_r, r_to_l, r_to_r):
        self.l_to_l = l_to_l
        self.l_to_r = l_to_r
        self.r_to_l = r_to_l
        self.r_to_r = r_to_r
    def apply_to_samples(self, i, left, right, frame_rate):
        l = left[i]
        r = right[i]
        left[i]   = l * self.l_to_l + r * self.r_to_l
        right[i]  = l * self.l_to_r + r * self.r_to_r

class Overdriver(Effect):
    def __init__(self, gain, cutoff):
        assert cutoff > 0
        self.gain = gain
        self.cutoff = cutoff
    def apply_to_samples(self, i, left, right, frame_rate):
        l = float(left[i])
        left[i]  = numpy.sign(l) * min(self.cutoff, abs(l * self.gain))
        r = float(right[i])
        right[i] = numpy.sign(r) * min(self.cutoff, abs(r * self.gain))

class Oscillator(Effect):
    def __init__(self, freq, amplitude):
        assert freq > 0
        self.freq = freq
        self.amplitude = amplitude
        self._ix_frame = 0
    def apply_to_samples(self, i, left, right, frame_rate):
        frames_per_cycle = frame_rate / self.freq
        d = self._ix_frame / float(frames_per_cycle)
        a = self.amplitude ** math.sin(2 * math.pi * d)
        left[i] *= a
        right[i] *= a
        self._ix_frame += 1
        if self._ix_frame >= frames_per_cycle:
            self._ix_frame = 0

class Compressor(Effect):
    def __init__(self, low_thresh, low_factor, high_thresh, high_factor):
        assert low_thresh <= high_thresh
        self.low_thresh = low_thresh
        self.low_factor = float(low_factor)
        self.high_thresh = high_thresh
        self.high_factor = float(high_factor)
    def apply_to_samples(self, i, left, right, frame_rate):
        l = float(abs(left[i]))
        if l < self.low_thresh:
            l = self.low_thresh - (self.low_thresh - l) / self.low_factor
        elif l > self.high_thresh:
            l = self.high_thresh + (l - self.high_thresh) / self.high_factor
        left[i]   = numpy.sign(left[i]) * l
        r = float(abs(right[i]))
        if r < self.low_thresh:
            r = self.low_thresh - (self.low_thresh - r) / self.low_factor
        elif r > self.high_thresh:
            r = self.high_thresh + (r - self.high_thresh) / self.high_factor
        right[i]  = numpy.sign(right[i]) * r


import gensim
import generators as gens
from functions import make_dict, fill_batch, take_len, make_spell_miss_list
import numpy as np
import collections
from chainer import *
from chainer import functions as F
from chainer import optimizers as O
from chainer import links as L
import pickle
import random
import math
import time

train_txt = "train.txt"
ssweV = 'EandGmodel/eandgVoc.pkl'
ssweM = 'EandGmodel/eandg.model'
state_model = "EandGmodel/eandg.sta"
vocab_dict = 'EandGmodel/eandgVocabDict.pkl'

vocab_size = take_len(train_txt)
embed_size = 300
output_size = 2
hidden_size = 200
epoch = 13
window_size = 3
batch_size = 1
gpu = -1
r = 0.01
noise_size = 200

if gpu >= 0:
    cuda.check_cuda_available()
xp = cuda.cupy if gpu >= 0 else np

random.seed(0)
xp.random.seed(0)

id2word = {}
word2id = {}
word_freq = collections.defaultdict(lambda: len(word2id))
word2id["<unk>"] = 0
word2id["</s>"] = 1
word2id["EOS"] = -1
id2word[0] = '<unk>'
id2word[1] = '</s>'
id2word[-1] = 'EOS'

uni_dict = pickle.load(open('source_target.pkl','rb')) # error patterns
word2id, id2word, word_list, word_freq = make_dict(train_txt, word2id, id2word, word_freq)

pickle.dump(word2id, open(vocab_dict, 'wb'))

def norm(data):
    return xp.sqrt(xp.sum(data**2))

def cos_sim(v1, v2):
    return xp.dot(v1, v2) / (norm(v1) * norm(v2))

def evaluate(model):
    cos_sum = 0
    m = model.copy()
    m.volatile = True
    s_l = make_spell_miss_list()
    count = 0
    for word in s_l:
        if word[0] in word2id and word[1] in word2id:
            count += 1
            cos_sum += cos_sim(m.x2e.W.data[word2id[word[0]]], m.x2e.W.data[word2id[word[1]]])

def convert_word(sentence):
    return [word2id[word] for word in sentence]

def make_sequences(sentences, tag):
    csequences = xp.array([sentences[i:i+window_size] for i in range(len(sentences)-window_size+1)], dtype=xp.int32)
    nls = list(csequences[:,1:2])
    noise = []
    for nl in nls:
        if int(*nl) in id2word and id2word[int(*nl)] in uni_dict:
            d = id2word[int(*nl)]
            ns = uni_dict[d]
            noise.append([[word2id[n[0]]] if n[0] in word2id else [word2id['<unk>']] for i,n in enumerate(sorted(ns.items(), key=lambda x:x[1], reverse=True)) if i < noise_size])
            for no in noise:
                while len(no) < noise_size:
                    no.append([random.randint(2, vocab_size)])
        else:
            noise.append(xp.random.randint(2,vocab_size, size=(noise_size, 1)))
    noise = xp.array(noise, dtype=xp.int32)
    sequences = xp.array([sentences[i:i+window_size] for i in range(len(sentences)-window_size+1)], dtype=xp.int32)
    zeros = xp.zeros((len(csequences),noise_size, window_size), dtype=xp.int32)
    csequences = xp.reshape(csequences, (len(csequences),1,window_size))
    csequences = csequences + zeros
    noise = xp.random.randint(vocab_size, size=(len(csequences),noise_size,1))
    csequences[:,:,1:2] = noise
    tags = xp.array([tag[i:i+window_size] for i in range(len(tag)-window_size+1)], dtype=xp.int32)
    tags = xp.sum(tags, axis=1)
    tags = xp.where(tags>0, xp.ones(len(tags), dtype=xp.int32), xp.zeros(len(tags), dtype=xp.int32))
    return sequences, csequences, tags

class SSWE_model(Chain):
    def __init__(self, vocab_size, embed_size, hidden_size, train=True):
        super(SSWE_model, self).__init__(
            x2e = L.EmbedID(vocab_size, embed_size),
            e2h = L.Linear(embed_size*window_size, hidden_size),
            h2h = L.Linear(hidden_size, hidden_size),
            h2o = L.Linear(hidden_size, 1),
            h2p = L.Linear(hidden_size, 2)
            )

    def __call__(self, sequence):
        self._reset_state()
        e = F.tanh(self.x2e(sequence))
        h = F.tanh(self.e2h(e))
        o = self.h2o(h)
        p = self.h2p(h)
        return o, p

    def _reset_state(self):
        self.zerograds()

    def initialize_embed(self, word2vec_model, word_list, word2id):
        for i in range(len(word_list)):
            word = word_list[i]
            if word in word2vec_model:
                self.x2e.W.data[i+2] = word2vec_model[word]

def forward(x, tags, model):
    accum_loss = Variable(xp.zeros((), dtype=xp.float32))
    x, c, t = make_sequences(x, tags)
    x = Variable(x)
    leng = len(c)
    c = Variable(xp.reshape(c, (leng*noise_size, window_size)))
    t = Variable(t)
    o, p = model(x)
    co, cp = model(c)
    a = F.softmax_cross_entropy(p, t)
    co = F.reshape(co, (leng,noise_size,1))
    o = F.broadcast_to( F.reshape(o, (leng,1,1)), (leng, noise_size, 1))
    b = F.sum(F.clipped_relu(1-co+o, 1e999))
    accum_loss = (1-r)*a +r* b
    return accum_loss

def train(model, opt):
    for i in range(epoch):
        print("epoch {}".format(i))
        start = time.time()
        total_loss = 0
        gen1 = gens.word_list(train_txt)
        gen2 = gens.batch(gens.sorted_parallel(gen1, embed_size*batch_size), batch_size)
        batchs = [ b for b in gen2]
        bl = list(range(len(batchs)))
        random.shuffle(bl)
        for n, j in enumerate(bl):
            if window_size > len(batchs[j][0])-4:
                continue
            tag0 = batchs[j][:]
            tags = [[int(c) for c in a[:-1]] for a in tag0]
            batch = fill_batch([b[-1].split() for b in batchs[j]])
            batch = [convert_word(b)  for b in batch]
            tags = xp.array(tags, dtype=xp.int32)
            accum_loss = forward(*batch, *tags, model)
            accum_loss.backward()
            opt.update()
            total_loss += accum_loss.data
        print(total_loss)
        evaluate(model)
        serializers.save_npz("{}{}".format(ssweM, i), model)
        print("time: {}".format(time.time() - start))

def main():
    model = SSWE_model(vocab_size, embed_size, hidden_size)
    if gpu >= 0:
        cuda.get_device(gpu).use()  # Make a specified GPU current
        model.to_gpu()  # Copy the model to the GPU
    model.x2e.W.data[0] = xp.zeros(embed_size, dtype=xp.float32)
    opt = O.Adam(alpha = 0.001)
    opt.setup(model)
    train(model, opt)
    with open(ssweV, 'wb') as f:
        pickle.dump(dict(word2id), f)
    state_d = {}
    state_d["vocab_size"] = vocab_size
    state_d["hidden_size"] = hidden_size
    state_d["embed_size"] = embed_size
    pickle.dump(state_d, open(state_model, "wb"))

if __name__ == '__main__':
    main()


'''Example script to generate text from Nietzsche's writings.

At least 20 epochs are required before the generated text
starts sounding coherent.

It is recommended to run this script on GPU, as recurrent
networks are quite computationally intensive.

If you try this script on new data, make sure your corpus
has at least ~100k characters. ~1M is better.
'''

from __future__ import print_function
from keras.callbacks import LambdaCallback
from keras.models import Sequential
from keras.models import load_model
from keras.layers import Dense
from keras.layers import Dropout
from keras.layers import LSTM
from keras.optimizers import RMSprop
from keras.utils.data_utils import get_file
import numpy as np
import random
import sys
import io, getopt, ast
from pathlib import Path

# ===========================================================

# parameters
dataset_path = "./datasets/nancy-drew-1-secret-of-old-clock.txt"

load_file = False
load_path = "./checkpoint.h5"
save_path = "./"
file_name = "checkpoint"

num_epochs = 5
checkpoints = list(range(num_epochs))
batch_size = 256
look_back = 40
step_size = 1

# Takes Command line inputs to override the above
if __name__ == "__main__":
    argv = sys.argv[1:]

    try:
        opts, args = getopt.getopt(argv,"hd:c:e:l:n:",["dataset=","checkpoints=","epochs=","load_model=","name="])
    except getopt.GetoptError:
        print ('train-model.py -d <dataset> -c <epochs to checkpoint> -e <# epochs> -l <model to load from> -n <name of files>')
        sys.exit(2)

    for opt, arg in opts:

        # Help Command
        if opt == '-h':
            print ('train-model.py -d <dataset> -c <epochs to checkpoint> -e <# epochs> -l <model to load from> -n <name of files>')
            sys.exit()

        # Num Epochs
        elif opt in ("-e","--epochs"):
            num_epochs = ast.literal_eval(arg)
            checkpoints = list(range(num_epochs))

        # Dataset Name
        elif opt in ("-d", "--dataset"):
            dataset_path = arg

        # Checkpoints
        elif opt in ("-c", "--checkpoints"):
            checkpoints = ast.literal_eval(arg)

        # Load Model
        elif opt in ("-l", "--load_model"):
            load_file = True
            load_file = arg

        elif opt in ("-n", "--name"):
            file_name = arg

# ===========================================================

with io.open(dataset_path, encoding='utf-8') as f:
    text = f.read().lower()
print('corpus length:', len(text))

chars = sorted(list(set(text)))
print('total chars:', len(chars))
char_indices = dict((c, i) for i, c in enumerate(chars))
indices_char = dict((i, c) for i, c in enumerate(chars))

# Outputs a charset
charset = sorted(list(set(text)))
# Make char's JS Readable
for i in range(len(charset)):
    if (charset[i] == '\n'):
        charset[i] = '\\n'
    if (charset[i] == '"'):
        charset[i] = '\\"'
# Generates the charset file
f = open(file_name + "-charset.txt","w+")
charset_final = '["'+ '", "'.join(charset) + '"]'
f.write(charset_final)
f.close()

print (charset_final)

# ===========================================================

# cut the text in semi-redundant sequences of maxlen characters
maxlen = look_back
sentences = []
next_chars = []
for i in range(0, len(text) - maxlen, step_size):
    sentences.append(text[i: i + maxlen])
    next_chars.append(text[i + maxlen])
print('nb sequences:', len(sentences))

print('Vectorization...')
x = np.zeros((len(sentences), maxlen, len(chars)), dtype=np.bool)
y = np.zeros((len(sentences), len(chars)), dtype=np.bool)
for i, sentence in enumerate(sentences):
    for t, char in enumerate(sentence):
        x[i, t, char_indices[char]] = 1
    y[i, char_indices[next_chars[i]]] = 1

# ===========================================================

# build the model: a single LSTM
print('Build model...')

my_file = Path(load_path)
if load_file and my_file.is_file():
    print("Found Checkpoint. Loading Saved model")
    model = load_model(load_path)
else:
    model = Sequential()
    model.add(LSTM(128, input_shape=(maxlen, len(chars))))
    model.add(Dense(len(chars), activation='softmax'))

    optimizer = RMSprop(lr=0.01)
    model.compile(loss='categorical_crossentropy', optimizer=optimizer)

# ===========================================================

def sample(preds, temperature=1.0):
    # helper function to sample an index from a probability array
    preds = np.asarray(preds).astype('float64')
    preds = np.log(preds) / temperature
    exp_preds = np.exp(preds)
    preds = exp_preds / np.sum(exp_preds)
    probas = np.random.multinomial(1, preds, 1)
    return np.argmax(probas)

# ===========================================================

def on_epoch_end(epoch, _):

    # Checkpointing the model
    for i in checkpoints:
        if epoch + 1 == i:
            print("Checkpointing the model...")
            model.save("%s-%d.h5" % (file_name,i))
    print()

    diversity = 1
    start_index = random.randint(0, len(text) - maxlen - 1)

    print('----- diversity:', diversity)

    generated = ''
    sentence = text[start_index: start_index + maxlen]
    generated += sentence

    print('----- Generating with seed: "' + sentence + '"')
    sys.stdout.write(generated)

    for i in range(400):
        x_pred = np.zeros((1, maxlen, len(chars)))
        for t, char in enumerate(sentence):
            x_pred[0, t, char_indices[char]] = 1.

        preds = model.predict(x_pred, verbose=0)[0]
        next_index = sample(preds, diversity)
        next_char = indices_char[next_index]

        generated += next_char
        sentence = sentence[1:] + next_char

        sys.stdout.write(next_char)
        sys.stdout.flush()

    '''
    print('----- Generating text after Epoch: %d' % epoch)

    start_index = random.randint(0, len(text) - maxlen - 1)

    for diversity in [0.2, 0.5, 1.0, 1.2]:
        print('----- diversity:', diversity)

        generated = ''
        sentence = text[start_index: start_index + maxlen]
        generated += sentence
        print('----- Generating with seed: "' + sentence + '"')
        sys.stdout.write(generated)

        for i in range(400):
            x_pred = np.zeros((1, maxlen, len(chars)))
            for t, char in enumerate(sentence):
                x_pred[0, t, char_indices[char]] = 1.

            preds = model.predict(x_pred, verbose=0)[0]
            next_index = sample(preds, diversity)
            next_char = indices_char[next_index]

            generated += next_char
            sentence = sentence[1:] + next_char

            sys.stdout.write(next_char)
            sys.stdout.flush()
        print()
    '''
# ===========================================================

print_callback = LambdaCallback(on_epoch_end=on_epoch_end)

# ===========================================================

# checks if the checkpoints array says to checkpoint initial model
if 0 in checkpoints:
    print("Checkpointing the model...")
    model.save("%s-%d.h5" % (file_name,0))

# begin the training process
model.fit(x, y,
          batch_size=batch_size,
          epochs=num_epochs,
          callbacks=[print_callback])

print("Source: \"%s\" \nEpochs: %d \nBatch Size: %d \nStep Size: %d \n Look Back: %d" % (dataset_path, num_epochs, batch_size, step_size, maxlen))

# ===========================================================


import argparse
import collections

import numpy as np

import torch
import torch.optim as optim
from   torchvision import transforms

from retinanet import model
from retinanet.dataloader import CSVDataset, collater, Resizer, AspectRatioBasedSampler, Augmenter, Normalizer
from pathlib import Path
from torch.utils.data import DataLoader

from retinanet import csv_eval

assert torch.__version__.split('.')[0] == '1'

print('CUDA available: {}'.format(torch.cuda.is_available()))

def main(args=None):

    csv_train   = 'my_annotations.csv'
    csv_classes = 'my_class_list.csv'
    csv_val     = 'my_annotations.csv'
    model_path  = 'model_weights_test.pt'

    depth       = 18
    batch_size  = 1
    epochs      = 2

    evaluate_validation_set = False

    # Training dataset
    dataset_train = CSVDataset(train_file=csv_train,
                               class_list=csv_classes,
                               transform=transforms.Compose([Normalizer(), Augmenter(), Resizer()]))

    # Validation dataset
    if csv_val is None:
        dataset_val = None
        print('No validation annotations provided.')
    else:
        dataset_val = CSVDataset(train_file=csv_val,
                                 class_list=csv_classes,
                                 transform=transforms.Compose([Normalizer(), Resizer()]))

    sampler = AspectRatioBasedSampler(dataset_train, batch_size=batch_size, drop_last=False)
    dataloader_train = DataLoader(dataset_train, num_workers=3, collate_fn=collater, batch_sampler=sampler)

    if dataset_val is not None:
        sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
        dataloader_val = DataLoader(dataset_val, num_workers=3, collate_fn=collater, batch_sampler=sampler_val)

    # Create the model
    if depth == 18:
        retinanet = model.resnet18(num_classes=dataset_train.num_classes(), pretrained=True)
    elif depth == 34:
        retinanet = model.resnet34(num_classes=dataset_train.num_classes(), pretrained=True)
    elif depth == 50:
        retinanet = model.resnet50(num_classes=dataset_train.num_classes(), pretrained=True)
    elif depth == 101:
        retinanet = model.resnet101(num_classes=dataset_train.num_classes(), pretrained=True)
    elif depth == 152:
        retinanet = model.resnet152(num_classes=dataset_train.num_classes(), pretrained=True)
    else:
        raise ValueError('Unsupported model depth, must be one of 18, 34, 50, 101, 152')

    use_gpu = True

    if use_gpu:
        if torch.cuda.is_available():
            retinanet = retinanet.cuda()

    if torch.cuda.is_available():
        retinanet = torch.nn.DataParallel(retinanet).cuda()
    else:
        retinanet = torch.nn.DataParallel(retinanet)

    retinanet.training = True

    optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)

    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True)

    loss_hist = collections.deque(maxlen=500)

    # History
    train_hist = []
    valid_hist = []
    eval_hist  = []
    train_hist_header = 'epoch,iteration,loss,classification loss,regression loss,running loss\n'
    valid_hist_header = 'epoch,loss,classification loss,regression loss\n'
    eval_hist_header  = 'epoch,train map,valid map,train iou,valid iou\n'
    train_hist.append(train_hist_header)
    valid_hist.append(valid_hist_header)
    eval_hist.append(eval_hist_header)

    retinanet.train()
    retinanet.module.freeze_bn()

    print('Num training images: {}'.format(len(dataset_train)))

    for epoch_num in range(epochs):

        retinanet.train()
        retinanet.module.freeze_bn()

        epoch_loss          = []
        validation_loss_tot = []
        validation_loss_cls = []
        validation_loss_reg = []

        # Training
        for iter_num, data in enumerate(dataloader_train):
            try:
                optimizer.zero_grad()

                if torch.cuda.is_available():
                    classification_loss, regression_loss = retinanet([data['img'].cuda().float(), data['annot']])
                else:
                    classification_loss, regression_loss = retinanet([data['img'].float(), data['annot']])

                classification_loss = classification_loss.mean()
                regression_loss = regression_loss.mean()

                loss = classification_loss + regression_loss

                if bool(loss == 0):
                    continue

                loss.backward()

                torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)

                optimizer.step()

                loss_hist.append(float(loss))

                epoch_loss.append(float(loss))

                if iter_num%1000==0:
                    print('Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(
                           epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(loss_hist)))

                # Construct history string
                train_info = str(epoch_num) + ',' + str(iter_num) + ',' + str(float(loss)) + ',' + str(float(classification_loss)) + ',' + str(float(regression_loss)) + ',' + str(np.mean(loss_hist)) + '\n'
                train_hist.append(train_info)

                del classification_loss
                del regression_loss
            except Exception as e:
                print(e)
                continue

        # Validation
        if csv_val is not None:
            for data_val in dataloader_val:
                try:
                    if torch.cuda.is_available():
                        classification_loss_val, regression_loss_val = retinanet([data_val['img'].cuda().float(), data_val['annot']])
                    else:
                        classification_loss_val, regression_loss_val = retinanet([data_val['img'].float(), data_val['annot']])

                    classification_loss_val = classification_loss_val.mean()
                    regression_loss_val     = regression_loss_val.mean()
                    loss_val                = classification_loss_val + regression_loss_val

                    if bool(loss_val == 0):
                        continue

                    validation_loss_tot.append(float(loss_val))
                    validation_loss_cls.append(float(classification_loss_val))
                    validation_loss_reg.append(float(regression_loss_val))

                    del classification_loss_val
                    del regression_loss_val
                    del loss_val
                except Exception as e:
                    print(e)
                    continue

            print('Epoch: {} | Validation loss: {:1.5f} | Validation Classification loss: {:1.5f} | Validation Regression loss: {:1.5f}'.format(
                   epoch_num, np.mean(validation_loss_tot), np.mean(validation_loss_cls), np.mean(validation_loss_reg)))

            # Construct history string
            valid_info = str(epoch_num) + ',' + str(np.mean(validation_loss_tot)) + ',' + str(np.mean(validation_loss_cls)) + ',' + str(np.mean(validation_loss_reg)) + '\n'
            valid_hist.append(valid_info)

        # Evaluate validation set
        if csv_val is not None and evaluate_validation_set:

            #print('Evaluating training set')
            #train_mAP_epoch, train_IoU_epoch = csv_eval.evaluate(dataset_train, retinanet, verbose=True)

            print('Evaluating validation set')
            valid_mAP_epoch, valid_IoU_epoch = csv_eval.evaluate(dataset_val, retinanet, verbose=True)

            # Construct history string
            eval_info = str(epoch_num) + ',' + '0' + ',' + str(valid_mAP_epoch[0][0]) + ',' + '0' + ',' + str(valid_IoU_epoch[0]) + '\n'
            eval_hist.append(eval_info)

        scheduler.step(np.mean(epoch_loss))

        #torch.save(retinanet.module, '{}_retinanet_{}.pt'.format(dataset, epoch_num))

    retinanet.eval()

    torch.save(retinanet, model_path)

    # Save training and validation history
    save_history = True
    if save_history:

        history_path = Path('results/history')
        if not history_path.parent.is_dir():
            history_path.parent.mkdir(parents=True)

        # Save training history
        with open("results/history/training_history.csv", "w") as csv_file:
            for info in train_hist:
                csv_file.write(info)

        # Save validation history
        with open("results/history/validation_history.csv", "w") as csv_file:
            for info in valid_hist:
                csv_file.write(info)

if __name__ == '__main__':
    main()


import numpy as np

from nefct.geometry.scanner_config import ScannerConfig2D, ScannerConfig3D
from nefct.geometry.detector_config import DetectorDirectAConfig, DetectorDirectBConfig

__all__ = (
    'cbct_a_detector', 'cbct_b_detector', 'cbct_2d_scanner', 'cbct_3d_scanner', 'cbct_angles',
    'cbct_2d_full_fan_scanner', 'cbct_3d_full_fan_scanner', 'xcat_volumes', 'xcat_flows')
cbct_a_detector = DetectorDirectAConfig(512, 0.388 * 1024, 147.44)
cbct_b_detector = DetectorDirectBConfig(512, 0.388 * 1024, 0)
cbct_2d_scanner = ScannerConfig2D('flat', 1500, 1000, cbct_a_detector)
cbct_3d_scanner = ScannerConfig3D('flat', 1500, 1000, cbct_a_detector, cbct_b_detector)
cbct_2d_full_fan_scanner = ScannerConfig2D('flat', 1500, 1000, cbct_a_detector.update(offset = 0))
cbct_3d_full_fan_scanner = ScannerConfig3D('flat', 1500, 1000, cbct_a_detector.update(offset = 0),
                                           cbct_b_detector)
cbct_angles = np.linspace(0, np.pi * 2, 599).astype(np.float32)
xcat_volumes = np.array([4.1, 4.9, 5.5, 5.7, 5.7, 5.1, 4.2, 3.4, 2.9, 2.4, 1.9,
                         1.5, 1., 0.5, 0.3, 0.3, 0.3, 0.2, 0.2, 0.2, 0.1, 0.1,
                         -0.1, -0.3, -0.3, -0.3, -0.2, -0.2, -0.1, -0., 0.1, 0.4, 0.6,
                         1., 1.7, 2.5, 3.3, 4., 4.7, 5.4, 6.1, 6.3, 5.6, 4.4,
                         3.3, 2.8, 2.3, 1.9, 1.4, 1.1, 0.8, 0.4, 0.2, 0.1, 0.2,
                         0.2, 0.2, 0.2, 0.2, 0.2, -0., -0.2, -0.4, -0.3, -0.2, -0.2,
                         0., 0.4, 1., 1.9, 2.6, 3.1, 3.8, 4.6, 5.5, 6.4, 7.2,
                         7.2, 6.2, 4.9, 3.7, 2.7, 2.1, 1.8, 1.4, 1.1, 0.8, 0.6,
                         0.4, 0.2, 0., -0.1, -0.1, 0., 0.1, 0.1, 0.1, 0.1, 0.,
                         -0.1, -0.4, -0.4, -0.1, 0.2, 0.7, 1.4, 2.2, 3.1, 3.8, 4.4,
                         5., 5.8, 6.8, 7.7, 7.7, 6.7, 5.3, 4.2, 3.4, 2.5, 1.8,
                         1.3, 1.1, 1., 0.8, 0.6, 0.4, 0.3, 0.1, -0., -0., 0.1,
                         0.1, 0.1, 0.2, 0.2, 0.2, 0., -0.2, -0.3, -0.1, 0.1, 0.3,
                         0.7, 1.3, 2., 2.8, 3.5, 3.9, 4.4, 5.1, 6., 6.8, 6.9,
                         6., 4.8, 3.9, 3., 2.3, 1.8, 1.5, 1.4, 1.1, 0.9, 0.7,
                         0.6, 0.4, 0.2, 0.1, 0.1, 0.1, 0.1, 0.2, 0.2, 0.1, 0.,
                         -0.1, -0.2, -0.1, 0.2, 0.7, 1.4, 2.3, 3.1, 4., 4.8, 5.5,
                         6.3, 7., 7.6, 7.2, 6.1, 4.8, 3.9, 3.2, 2.5, 1.7, 1.1,
                         0.7, 0.6, 0.5, 0.4, 0.4, 0.3, 0.3, 0.1, -0.1, -0.1, 0.,
                         0.1, 0.1, 0.1, 0., -0., -0.1, -0.2, -0.3, 0., 0.4, 1.,
                         1.7, 2.5, 3.2, 3.8, 4.1, 4.3, 4.4, 4.4, 3.9, 3.2, 2.6,
                         2.3, 2., 1.7, 1., 0.6, 0.4, 0.4, 0.4, 0.4, 0.3, 0.3,
                         0.2, 0.1, -0.1, -0.2, -0.2, -0.1, -0.1, -0.1, 0., 0.2, 0.3,
                         0.5, 0.8, 1.2, 2., 2.9, 3.7, 4.3, 5., 5.7, 6.6, 7.1,
                         6.8, 5.6, 4.4, 3.8, 3.2, 2.6, 2.1, 1.7, 1.3, 0.8, 0.5,
                         0.4, 0.4, 0.4, 0.4, 0.3, 0.3, 0.3, 0.1, -0.1, -0.2, -0.1,
                         -0., -0.1, -0.1, -0.1, -0., 0.3, 0.6, 1.1, 1.8, 2.6, 3.5,
                         4.2, 5., 5.8, 6.6, 7., 6.5, 5.2, 4., 3.4, 2.9, 2.4,
                         2., 1.6, 1.2, 0.8, 0.3, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1,
                         0., -0., -0.2, -0.5, -0.5, -0.3, -0.1, 0.2, 0.8, 1.4, 2.1,
                         2.9, 3.5, 3.9, 4.4, 5.1, 5.8, 6.1, 5.7, 4.7, 3.9, 3.1,
                         2.2, 1.6, 1.2, 1., 0.8, 0.6, 0.5, 0.3, 0.3, 0.1, -0.1,
                         -0.1, -0., -0., 0., -0., -0.1, -0.2, -0.1, 0., 0.4, 1.,
                         1.9, 2.8, 3.6, 4.2, 4.8, 5.4, 5.8, 5.9, 5.7, 5., 4.1,
                         3.5, 3.2, 2.8, 2.3, 1.8, 1.2, 0.7, 0.5, 0.5, 0.4, 0.3,
                         0.3, 0.2, 0., -0.2, -0.3, -0.3, -0.2, -0.2, -0.2, -0.1, -0.,
                         0.1, 0.2, 0.4, 0.8, 1.4, 2.2, 3.1, 3.8, 4.5, 5.1, 5.8,
                         6.3, 6.3, 5.6, 4.4, 3.5, 2.8, 2.2, 1.7, 1.2, 0.8, 0.4,
                         0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.2, -0.4, -0.7,
                         -0.6, -0.5, -0.4, -0.4, -0.3, -0.3, -0.1, 0., 0.2, 0.4, 0.8,
                         1.5, 2.3, 3.1, 3.8, 4.4, 5.1, 5.8, 5.9, 5.4, 4.5, 3.8,
                         3.2, 2.6, 2.1, 1.8, 1.4, 0.9, 0.5, 0.2, 0.2, 0.2, 0.2,
                         0.2, 0.1, 0.1, 0., -0.2, -0.5, -0.4, -0.3, -0.2, -0.2, -0.2,
                         -0.2, -0.2, -0.1, -0.1, 0., 0.3, 0.8, 1.6, 2.6, 3.5, 4.3,
                         5.2, 5.9, 6.4, 6.6, 6.8, 6.6, 5.8, 4.5, 3.7, 3.1, 2.6,
                         1.8, 1.1, 0.6, 0.4, 0.3, 0.2, 0.1, 0.1, -0., -0.1, -0.2,
                         -0.5, -0.5, -0.3, -0.3, -0.3, -0.2, -0.3, -0.4, -0.4, -0.5, -0.4,
                         0.1, 0.7, 1.6, 2.6, 3.6, 4.4, 5.4, 6.2, 7.1, 8., 9.,
                         10., 10.7, 10.4, 9.2, 7.8, 6.7, 5.3, 4., 3., 2.4, 1.9,
                         1.4, 1., 0.6, 0.3, 0.1, -0.3, -0.6, -0.6, -0.5, -0.5, -0.5,
                         -0.5, -0.6, -0.6, -0.7, -1., -1., -0.9, -0.8, -0.8, -0.6, -0.3,
                         0.1, 0.6, 1.1, 1.8, 2.7, 3.6, 4.4, 5.4, 6.6, 8., 9.3,
                         10.5, 11.6, 11.8, 11.4, 10.8, 9.5, 7.6, 5.8, 4.6, 3.6, 2.7,
                         1.7, 1.1, 0.7, 0.5, 0.3, 0.1, -0., -0.2, -0.4, -0.6, -1.,
                         -1.1, -0.9, -0.8, -0.7, -0.7], dtype = np.float32)
xcat_flows = np.hstack(([0], xcat_volumes[1:] - xcat_volumes[:-1]))


#!/usr/bin/env python

import multiprocessing
import subprocess
import argparse
import shutil
import copy
import glob
import os
from Bio import SeqIO
from Bio import Phylo
from scipy import cluster

def make_ref(input_file, output_file):
    """
    Parsnp doesn't like having multi-fastas as input files, so this method puts all contigs in a multi-fasta
    into a single contig.
    :param input_file: Path to your input multi-fasta
    :param output_file: Path to output fasta. Overwrites the file if it already exists.
    """
    contigs = SeqIO.parse(input_file, 'fasta')
    with open(output_file, 'w') as f:
        f.write('>reference\n')
        for s in contigs:
            f.write(str(s.seq) + '\n')

def prune_clade_from_tree(tree, clade_name):
    """
    Given a treefile, will prune out the specified clade.
    :param tree: Treefile read in with Bio.Phylo
    :param clade_name: Name of clade you want removed from the tree.
    :return: New Biopython tree object.
    """
    new_tree = copy.deepcopy(tree)
    if new_tree.find_any(clade_name) is not None:
        new_tree.prune(clade_name)
    else:
        raise AttributeError('Clade {} not found in tree. Cannot be pruned.'.format(clade_name))
    return new_tree

def create_clustering(treefile, desired_clusters=10):
    """
    :param treefile: Tree file, in newick format.
    :param desired_clusters: The total number of clusters you want to return.
    :return: A list of strains that represent each cluster.
    """
    num_clusters = -1
    tree = Phylo.read(treefile, "newick")
    terminals = list()
    clades = tree.find_clades()
    matrix = list()
    # Get a list of all the terminal branches
    for clade in clades:
        if clade.is_terminal():
            terminals.append(str(clade))

    # Create a matrix with distances between each tip and each other tip
    for i in range(len(terminals)):
        for j in range(i + 1, len(terminals)):
            dist = tree.distance(terminals[i], terminals[j])
            matrix.append(dist)

    count = 1
    print('Finding correct number of clusters.')
    # Create the linkage thingy so we can try clustering.
    z = cluster.hierarchy.linkage(matrix, method='average')
    # Try different cutoff levels for making clusters. Start at a very large cluster distance.
    # Check if enough clusters are created. If not enough, make the cutoff slightly smaller and repeat process.
    cluster_distance = 0.9
    while num_clusters < desired_clusters:
        clustering = cluster.hierarchy.fcluster(z, cluster_distance, criterion='distance')
        num_clusters = (max(clustering))
        cluster_distance -= 0.00003
        count += 1
    # Once we've found the desired number of clusters, create a 2d array where each element of the array is a cluster
    # and each element is made up of a list of strains that are part of that cluster.
    clusters = list()
    for i in range(num_clusters):
        clusters.append(list())
    for i in range(len(clustering)):
        clusters[clustering[i] - 1].append(terminals[i])

    # Now that we have a list of strains for each cluster, choose the strain that's the most like everything else
    # in each cluster (the most average-y thing I guess?) to be the representative.
    strains = list()
    for c in clusters:
        representative_strain = choose_representative_strain(c, tree)
        representative_strain = representative_strain.replace('.fasta', '')
        strains.append(representative_strain)

    return strains

def choose_representative_strain(cluster, tree):
    """
    Given a cluster, will find which tip in that cluster tends to be closest to all others (aka the best representative)
    and return that tip
    :param cluster: Cluster found.
    :param tree: Tree read in by Bio.Phylo
    :return: representative strain.
    """
    representative = 'NA'
    best_distance = 100000000.0  # Start off at an absolutely ridiculous value.
    # Iterate through
    for strain1 in cluster:
        total_length = 0.0
        for strain2 in cluster:
            if strain1 != strain2:
                total_length += tree.distance(strain1, strain2)
        if total_length < best_distance:
            best_distance = total_length
            representative = strain1
    return representative

def main(args):
    if not os.path.isdir(args.output_dir):
        os.makedirs(args.output_dir)
    if not os.path.isdir(os.path.join(args.output_dir, 'tmp')):
        os.makedirs(os.path.join(args.output_dir, 'tmp'))
    tmpdir = os.path.join(args.output_dir, 'tmp')
    # Check that we haven't requested more strains than are present in the tree.
    if args.desired_strains > len(glob.glob(os.path.join(args.input_dir, '*fasta'))):
        raise ValueError('Number of requested strains cannot exceed number of input strains.')
    # Next step: Choose a reference. Parsnp will complain if there are multiple contigs, so make it into one contig.
    try:
        reference_file = glob.glob(os.path.join(args.input_dir, '*.fasta'))[0]
    except IndexError:
        raise FileNotFoundError('No fasta files found in specified directory. '
                         'Specified directory was: {}'.format(args.input_dir))
    # Create a copy of the tree that parsnp won't complain about.
    make_ref(reference_file, os.path.join(tmpdir, 'reference.fasta'))
    # Run PARSNP on everything to make a tree that we'll get to parsing to determine which species are far apart.
    cmd = '/mnt/nas/Programs/Parsnp-Linux64-v1.2/parsnp -r {tmpdir}/reference.fasta -d {input} -c -o {output} -p {threads}'.format(tmpdir=tmpdir,
                                                                                threads=str(args.num_threads),
                                                                                output=args.output_dir,
                                                                                input=args.input_dir)
    subprocess.call(cmd, shell=True)
    # Remove the "reference" fasta that's actually the same as one of our strains and write the revised tree to file.
    pruned_tree = prune_clade_from_tree(Phylo.read(os.path.join(args.output_dir, 'parsnp.tree'), 'newick'),
                                        'reference.fasta.ref')
    Phylo.write(pruned_tree, os.path.join(args.output_dir, 'pruned.tree'), 'newick')
    # Now get the desired number of strains.
    strains = create_clustering(os.path.join(args.output_dir, 'pruned.tree'), desired_clusters=args.desired_strains)
    shutil.rmtree(tmpdir)
    with open(os.path.join(args.output_dir, 'strains.txt'), 'w') as output_file:
        for strain in strains:
            output_file.write(strain + '\n')

if __name__ == '__main__':
    cpu_count = multiprocessing.cpu_count()
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input_dir',
                        type=str,
                        required=True,
                        help='Directory with sequences to be compared in uncompressed FASTA format.')
    parser.add_argument('-t', '--num_threads',
                        type=int,
                        default=cpu_count,
                        help='Number of threads to run analysis on.')
    parser.add_argument('-o', '--output_dir',
                        type=str,
                        required=True,
                        help='Output directory.')
    parser.add_argument('-d', '--desired_strains',
                        type=int,
                        required=True,
                        help='Number of strains you want to pick out.')
    args = parser.parse_args()
    main(args)


import copy
import json
import math
import random
import sys
import time

import numpy as np
from xt.environment.environment import Environment
from xt.framework.register import Registers

try:
    from flow.controllers import (ContinuousRouter, IDMController, RLController,
                                SimCarFollowingController)

    from flow.core.params import (EnvParams, InitialConfig, NetParams,
                                SumoCarFollowingParams, SumoParams,
                                VehicleParams)
    from flow.envs import AccelEnv
    from flow.envs.multiagent import MultiAgentAccelPOEnv, MultiAgentEightEnv
    from flow.networks import FigureEightNetwork
    from flow.networks.figure_eight import ADDITIONAL_NET_PARAMS
    from flow.utils.registry import make_create_env
except (ModuleNotFoundError, ImportError) as err:
    print("import error")

@Registers.env
class MaEnvFigure8(Environment):
    """kybersim environment for lanechange case"""
    def init_env(self, env_info):
        """
        cEnvFigure8reate a flow environment instance

        :param: the config information of environment
        :return: the instance of environment
        """
        self.init_state = None
        # update the agents number and env api type.
        self.n_agents = 10
        self.api_type = "unified"

        # We place one autonomous vehicle and 13 human-driven vehicles in the network
        vehicles = VehicleParams()
        vehicles.add(
            veh_id='rl',
            # initial_speed=5,
            acceleration_controller=(RLController, {}),
            routing_controller=(ContinuousRouter, {}),
            car_following_params=SumoCarFollowingParams(
                speed_mode="obey_safe_speed",
                # speed_mode="aggressive",
                decel=1.5,
            ),
            num_vehicles=10)

        # sumo-related parameters (see flow.core.params.SumoParams)
        sim = SumoParams(
            sim_step=0.1,
            render=False,
            restart_instance=True,
        )

        # environment related parameters (see flow.core.params.EnvParams)
        env_params = EnvParams(
            horizon=1000,
            additional_params={
                'target_velocity': 20,
                'max_accel': 3,
                'max_decel': 3,
                'sort_vehicles': False
            },
        )

        # network-related parameters (see flow.core.params.NetParams and the
        # network's documentation or ADDITIONAL_NET_PARAMS component)
        net = NetParams(additional_params=ADDITIONAL_NET_PARAMS.copy(), )

        # vehicles to be placed in the network at the start of a rollout (see
        # flow.core.params.VehicleParams)

        network = FigureEightNetwork(name='multiagent_figure_eight',
                                     vehicles=vehicles,
                                     net_params=net)

        # parameters specifying the positioning of vehicles upon initialization/
        env = MultiAgentEightEnv(env_params, sim, network)

        return env

    def reset(self):
        """
        reset the environment, if there are illegal data in observation
        then use old data

        :param reset_arg: reset scene information
        :return: the observation of environment
        """

        state = self.env.reset()
        self.init_state = state
        return state

    def step(self, actions, agent_index=0):
        """
        send lanechange cmd to kybersim

        :param action: action （0-2）
        :param agent_index: the index of agent
        :return: state, reward, done, info
        """

        state_, reward, done, _ = self.env.step(actions)
        return state_, reward, done, _

    def get_init_state(self, agent_index=0):
        """
        get reset observation of one agent.

        :param agent_index: the index of agent
        :return: the reset observation of agent
        """
        return self.init_state

    def close(self):
        """
        close environment
        """
        try:  # do some thing you need
            self.env.terminate()
        except AttributeError:
            print("please complete your env close function")

    def get_env_info(self):
        """rewrite environment's basic information, indicate multi-agents"""
        self.reset()
        env_info = {
            "n_agents": self.n_agents,
            "api_type": self.api_type,
        }
        # update the agent ids, will used in the weights map.
        # default work well with the sumo multi-agents
        agent_ids = list(
            self.get_init_state().keys())  # if self.n_agents > 1 else [0]
        env_info.update({"agent_ids": agent_ids})

        return env_info


import sys
import os
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(BASE_DIR)
sys.path.append(os.path.join(ROOT_DIR, 'train'))
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
# import ipdb
from torch.nn import init
from model_util import NUM_HEADING_BIN, NUM_SIZE_CLUSTER, NUM_OBJECT_POINT
from model_util import point_cloud_masking, parse_output_to_tensors
from model_util import FrustumPointNetLoss
#from model_util import g_type2class, g_class2type, g_type2onehotclass
#from model_util import g_type_mean_size
#from model_util import NUM_HEADING_BIN, NUM_SIZE_CLUSTER
# from provider_fpointnet import compute_box3d_iou

# NUM_HEADING_BIN = 12
# NUM_SIZE_CLUSTER = 6 # one cluster for each type
# NUM_OBJECT_POINT = 2048

# g_type2class={'table':0,'chair':1,'desk':2,'shelf':3,'cabinet':4,'box':5}
# g_class2type = {g_type2class[t]:t for t in g_type2class}
# g_type2onehotclass = {'table':0,'chair':1,'desk':2,'shelf':3,'cabinet':4,'box':5}
# g_type_mean_size = {'table': np.array([0.89078179,1.4377661,0.71774923]),
#                     'chair': np.array([0.58838503,0.55581025,0.84174869]),
#                     'desk': np.array([0.72018677,1.4049219,0.79348079]),
#                     'shelf':np.array([0.43792965,1.41958785,0.94123454]),
#                     'cabinet':np.array([0.57207646,1.37302015,1.21174823]),
#                     'box':np.array([0.37163918,0.39581311,0.3400343 ])
#                     }

# g_mean_size_arr = np.zeros((NUM_SIZE_CLUSTER, 3)) # size clustrs
# for i in range(NUM_SIZE_CLUSTER):
#     g_mean_size_arr[i,:] = g_type_mean_size[g_class2type[i]]

class PointNetInstanceSeg(nn.Module):
    def __init__(self,n_classes=6,n_channel=4):
        '''v1 3D Instance Segmentation PointNet
        :param n_classes:6
        :param one_hot_vec:[bs,n_classes]
        '''
        super(PointNetInstanceSeg, self).__init__()
        self.conv1 = nn.Conv1d(n_channel, 64, 1)
        self.conv2 = nn.Conv1d(64, 64, 1)
        self.conv3 = nn.Conv1d(64, 64, 1)
        self.conv4 = nn.Conv1d(64, 128, 1)
        self.conv5 = nn.Conv1d(128, 1024, 1)
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(64)
        self.bn3 = nn.BatchNorm1d(64)
        self.bn4 = nn.BatchNorm1d(128)
        self.bn5 = nn.BatchNorm1d(1024)

        self.n_classes = n_classes
        self.dconv1 = nn.Conv1d(1088+n_classes, 512, 1)
        self.dconv2 = nn.Conv1d(512, 256, 1)
        self.dconv3 = nn.Conv1d(256, 128, 1)
        self.dconv4 = nn.Conv1d(128, 128, 1)
        self.dropout = nn.Dropout(p=0.5)
        self.dconv5 = nn.Conv1d(128, 2, 1)
        self.dbn1 = nn.BatchNorm1d(512)
        self.dbn2 = nn.BatchNorm1d(256)
        self.dbn3 = nn.BatchNorm1d(128)
        self.dbn4 = nn.BatchNorm1d(128)

    def forward(self, pts, one_hot_vec): # bs,4,n
        '''
        :param pts: [bs,4,n]: x,y,z,intensity
        :return: logits: [bs,n,2],scores for bkg/clutter and object
        '''
        bs = pts.size()[0]
        n_pts = pts.size()[2]

        out1 = F.relu(self.bn1(self.conv1(pts))) # bs,64,n
        out2 = F.relu(self.bn2(self.conv2(out1))) # bs,64,n
        out3 = F.relu(self.bn3(self.conv3(out2))) # bs,64,n
        out4 = F.relu(self.bn4(self.conv4(out3)))# bs,128,n
        out5 = F.relu(self.bn5(self.conv5(out4)))# bs,1024,n
        global_feat = torch.max(out5, 2, keepdim=True)[0] #bs,1024,1

        expand_one_hot_vec = one_hot_vec.view(bs,-1,1)#bs,3,1
        expand_global_feat = torch.cat([global_feat, expand_one_hot_vec],1)#bs,1027,1
        expand_global_feat_repeat = expand_global_feat.view(bs,-1,1)\
                .repeat(1,1,n_pts)# bs,1027,n
        concat_feat = torch.cat([out2,\
            expand_global_feat_repeat],1)
        # bs, (641024+3)=1091, n

        x = F.relu(self.dbn1(self.dconv1(concat_feat)))#bs,512,n
        x = F.relu(self.dbn2(self.dconv2(x)))#bs,256,n
        x = F.relu(self.dbn3(self.dconv3(x)))#bs,128,n
        x = F.relu(self.dbn4(self.dconv4(x)))#bs,128,n
        x = self.dropout(x)
        x = self.dconv5(x)#bs, 2, n

        seg_pred = x.transpose(2,1).contiguous()#bs, n, 2
        return seg_pred

class PointNetEstimation(nn.Module):
    def __init__(self,n_classes=6):
        '''v1 Amodal 3D Box Estimation Pointnet
        :param n_classes:6
        :param one_hot_vec:[bs,n_classes]
        '''
        super(PointNetEstimation, self).__init__()
        self.conv1 = nn.Conv1d(3, 128, 1)
        self.conv2 = nn.Conv1d(128, 128, 1)
        self.conv3 = nn.Conv1d(128, 256, 1)
        self.conv4 = nn.Conv1d(256, 512, 1)
        self.bn1 = nn.BatchNorm1d(128)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(256)
        self.bn4 = nn.BatchNorm1d(512)

        self.n_classes = n_classes

        self.fc1 = nn.Linear(512+n_classes, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256,3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER*4)
        self.fcbn1 = nn.BatchNorm1d(512)
        self.fcbn2 = nn.BatchNorm1d(256)

    def forward(self, pts,one_hot_vec): # bs,3,m
        '''
        :param pts: [bs,3,m]: x,y,z after InstanceSeg
        :return: box_pred: [bs,3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER*4]
            including box centers, heading bin class scores and residuals,
            and size cluster scores and residuals
        '''
        bs = pts.size()[0]
        n_pts = pts.size()[2]

        out1 = F.relu(self.bn1(self.conv1(pts))) # bs,128,n
        out2 = F.relu(self.bn2(self.conv2(out1))) # bs,128,n
        out3 = F.relu(self.bn3(self.conv3(out2))) # bs,256,n
        out4 = F.relu(self.bn4(self.conv4(out3)))# bs,512,n
        global_feat = torch.max(out4, 2, keepdim=False)[0] #bs,512

        expand_one_hot_vec = one_hot_vec.view(bs,-1)#bs,3
        expand_global_feat = torch.cat([global_feat, expand_one_hot_vec],1)#bs,515

        x = F.relu(self.fcbn1(self.fc1(expand_global_feat)))#bs,512
        x = F.relu(self.fcbn2(self.fc2(x)))  # bs,256
        box_pred = self.fc3(x)  # bs,3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER*4
        return box_pred

class STNxyz(nn.Module):
    def __init__(self,n_classes=6):
        super(STNxyz, self).__init__()
        self.conv1 = torch.nn.Conv1d(3, 128, 1)
        self.conv2 = torch.nn.Conv1d(128, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 256, 1)
        #self.conv4 = torch.nn.Conv1d(256, 512, 1)
        self.fc1 = nn.Linear(256+n_classes, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 3)

        init.zeros_(self.fc3.weight)
        init.zeros_(self.fc3.bias)

        self.bn1 = nn.BatchNorm1d(128)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(256)
        self.fcbn1 = nn.BatchNorm1d(256)
        self.fcbn2 = nn.BatchNorm1d(128)
    def forward(self, pts,one_hot_vec):
        bs = pts.shape[0]
        x = F.relu(self.bn1(self.conv1(pts)))# bs,128,n
        x = F.relu(self.bn2(self.conv2(x)))# bs,128,n
        x = F.relu(self.bn3(self.conv3(x)))# bs,256,n
        x = torch.max(x, 2)[0]# bs,256
        expand_one_hot_vec = one_hot_vec.view(bs, -1)# bs,3
        x = torch.cat([x, expand_one_hot_vec],1)#bs,259
        x = F.relu(self.fcbn1(self.fc1(x)))# bs,256
        x = F.relu(self.fcbn2(self.fc2(x)))# bs,128
        x = self.fc3(x)# bs,
        ###if np.isnan(x.cpu().detach().numpy()).any():
        ###    ipdb.set_trace()
        return x

class FrustumPointNetv1(nn.Module):
    def __init__(self,n_classes=6,n_channel=6):
        super(FrustumPointNetv1, self).__init__()
        self.n_classes = n_classes
        self.InsSeg = PointNetInstanceSeg(n_classes=6,n_channel=n_channel)
        self.STN = STNxyz(n_classes=6)
        self.est = PointNetEstimation(n_classes=6)

    def forward(self, pts, one_hot_vec):#bs,4,n
        # 3D Instance Segmentation PointNet
        logits = self.InsSeg(pts,one_hot_vec)#bs,n,2

        # Mask Point Centroid
        object_pts_xyz, mask_xyz_mean, mask = \
                 point_cloud_masking(pts, logits)###logits.detach()

        # if torch.lt(torch.sum(mask),100):
            # print("srus amsll")
            # print(torch.sum(mask))
            # one_hot_vec[0,0]=1
            # one_hot_vec[0,5]=0
            # logits = self.InsSeg(pts,one_hot_vec)#bs,n,2
            # object_pts_xyz, mask_xyz_mean, mask = \
            #      point_cloud_masking(pts, logits)
            # print(torch.sum(mask))

        # T-Net
        object_pts_xyz = object_pts_xyz.cuda()
        center_delta = self.STN(object_pts_xyz,one_hot_vec)#(32,3)
        stage1_center = center_delta + mask_xyz_mean#(32,3)

        # if(np.isnan(stage1_center.cpu().detach().numpy()).any()):
            # ipdb.set_trace()
        object_pts_xyz_new = object_pts_xyz - \
                    center_delta.view(center_delta.shape[0],-1,1).repeat(1,1,object_pts_xyz.shape[-1])

        # 3D Box Estimation
        box_pred = self.est(object_pts_xyz_new,one_hot_vec)#(32, 59)

        center_boxnet, \
        heading_scores, heading_residuals_normalized, heading_residuals, \
        size_scores, size_residuals_normalized, size_residuals = \
                parse_output_to_tensors(box_pred, logits, mask, stage1_center)

        center = center_boxnet + stage1_center #bs,3
        return logits, mask, stage1_center, center_boxnet, object_pts_xyz_new,\
            heading_scores, heading_residuals_normalized, heading_residuals, \
            size_scores, size_residuals_normalized, size_residuals, center

if __name__ == '__main__':
    #python models/pointnet.py
    points = torch.zeros(size=(32,4,1024),dtype=torch.float32)
    label = torch.ones(size=(32,3))
    model = FrustumPointNetv1()
    logits, mask, stage1_center, center_boxnet, \
            heading_scores, heading_residuals_normalized, heading_residuals, \
            size_scores, size_residuals_normalized, size_residuals, center \
            = model(points,label)
    print('logits:',logits.shape,logits.dtype)
    print('mask:',mask.shape,mask.dtype)
    print('stage1_center:',stage1_center.shape,stage1_center.dtype)
    print('center_boxnet:',center_boxnet.shape,center_boxnet.dtype)
    print('heading_scores:',heading_scores.shape,heading_scores.dtype)
    print('heading_residuals_normalized:',heading_residuals_normalized.shape,\
          heading_residuals_normalized.dtype)
    print('heading_residuals:',heading_residuals.shape,\
          heading_residuals.dtype)
    print('size_scores:',size_scores.shape,size_scores.dtype)
    print('size_residuals_normalized:',size_residuals_normalized.shape,\
          size_residuals_normalized.dtype)
    print('size_residuals:',size_residuals.shape,size_residuals.dtype)
    print('center:', center.shape,center.dtype)
    '''
    logits: torch.Size([32, 1024, 2]) torch.float32
    mask: torch.Size([32, 1024]) torch.float32
    stage1_center: torch.Size([32, 3]) torch.float32
    center_boxnet: torch.Size([32, 3]) torch.float32
    heading_scores: torch.Size([32, 12]) torch.float32
    heading_residual_normalized: torch.Size([32, 12]) torch.float32
    heading_residual: torch.Size([32, 12]) torch.float32
    size_scores: torch.Size([32, 8]) torch.float32
    size_residuals_normalized: torch.Size([32, 8, 3]) torch.float32
    size_residuals: torch.Size([32, 8, 3]) torch.float32
    center: torch.Size([32, 3]) torch.float32
    '''
    loss = FrustumPointNetLoss()
    mask_label = torch.zeros(32,1024).float()
    center_label = torch.zeros(32,3).float()
    heading_class_label = torch.zeros(32).long()
    heading_residuals_label = torch.zeros(32).float()
    size_class_label = torch.zeros(32).long()
    size_residuals_label = torch.zeros(32,3).float()
    output_loss = loss(logits, mask_label, \
                center, center_label, stage1_center, \
                heading_scores, heading_residuals_normalized, heading_residuals, \
                heading_class_label, heading_residuals_label, \
                size_scores,size_residuals_normalized,size_residuals,
                size_class_label,size_residuals_label)
    print('output_loss',output_loss)
    print()


import cv2
import numpy as np
import tensorflow as tf
from Mybase.losses import *
from Mybase.comp_utils import *
from Mybase.deep_metric_utils.embcls_layer import *
from .bbox import *

class EmbClsTargetLayer(object):
    '''
    The similarity between pixels p and q is computed as follows:
        σ(p, q) = 2 / (1 + exp(||ep - eq||2^2))
                = 2 / (1 + exp(d))            d = ||ep - eq||2^2
                = 2 * (1 - 1/(1+exp(-d)))
                = 2 * (1 - sigmoid(d)) ∈ (0, 1]
                ??? sigmoid(ep * eq / sqrt(d)) ???
        ||tensor||2 = tf.norm(tensor, ord='euclidean', axis=-1, keepdims=False)
    We train the network by minimizing the following loss:
        Le = -1/|S| * ∑(p,q∈S) wpq*[1{yp=yq}*log(σ(p,q)) + 1{yp≠yq}*log(1-σ(p,q))]
        tf.nn.sigmoid_cross_entropy_with_logits(labels=z, logits=x)
        = z * -log(sigmoid(x)) + (1 - z) * -log(1 - sigmoid(x))
        = z * -log(1 / (1 + exp(-x))) + (1 - z) * -log(exp(-x) / (1 + exp(-x)))
        = x - x * z + log(1+)
        sigmoid(x) = 1 / (1 + exp(-x))
        where S is the set of pixels that we choose, yp is the instance label of pixel p,
        and wpq is the weight of the loss on the similarity between p and q. The weights
        wpq are set to values inversely proportional to the size of the instances p and q
        belong to, so the loss will not become biased towards the larger examples. We
        normalize weights so that ∑(p,q)wpq = 1(p,q所处的两个实例的并集中的计算, wpq = 1 / (Sp * Sq)).
        |S|为图片中所选择的实例的个数。
    In contrast to semantic segmentation, where the pixels themselves are classified, here
    we classify the mask that each pixel will generate if chosen as a seed. We see that most
    pixels inside the object make good seeds, but pixels near the boundaries are not so good.
    '''
    def __init__(self, cls_num=None, img_shp=None, fet_shp=None):

        self.cls_num     = cls_num
        self.img_shp     = img_shp
        self.fet_shp     = fet_shp
        self.fet_srd     = img_shp[0] / fet_shp[0]
        self.pix_smp_num = 16
        self.pix_sim_num = 4
        self.pix_msk_min = 0.5
        self.msk_ovp_min = 0.5
        self.EL          = EmbclsLayer(cls_num, img_shp, fet_shp)

    def generate_embcls_loss_img(self, elems=None):

        fet_emb, fet_cls, gbxs, gmks, gbx_num = elems
        #fet_emb-->(H, W, 64), fet_cls-->(H, W, 4, C+1)
        #若gbx_num==1，即无正样本，gbxs->(0,5)，gmks->(0, H, W)
        gbxs = gbxs[0:gbx_num]                                                             #(M+1, 5)    最后一个是无效边框
        gmks = gmks[0:gbx_num]                                                             #(M+1, H, W) 最后一个是无效Mask

        #把gmks根据gbxs融进fet_shp大小的特征图中
        def recover_mask(elems=None):
            gbx, gmk = elems                                                               #(5), (H‘, W’)
            crd = tf.cast(gbx[0:-1]/self.fet_srd, dtype=tf.int32)                          #(4)
            ymn = crd[0]
            xmn = crd[1]
            ymx = crd[2]
            xmx = crd[3]
            leh = tf.stack([ymx-ymn+1, xmx-xmn+1], axis=0)
            gmk = tf.expand_dims(gmk, axis=-1)                                             #(H“, W”, 1)
            gmk = tf.image.resize_images(gmk, leh, method=tf.image.ResizeMethod.BILINEAR, align_corners=False) #(H“, W”, 1)
            gmk = tf.squeeze(gmk, axis=[-1])                                               #(H“, W”)
            gmk = tf.cast(gmk>=self.pix_msk_min, dtype=tf.float32)                         #(H“, W”)
            paddings = [[ymn, self.fet_shp[0]-ymx-1], [xmn, self.fet_shp[1]-xmx-1]]        #(H, W)
            gmk = tf.pad(gmk, paddings, 'CONSTANT')
            return gmk
        #当没有box时，tf.map_fn不会拆开，所以tf.image.resize_images是安全的
        elems     = [gbxs[0:-1], gmks[0:-1]]
        gmks      = tf.map_fn(recover_mask, elems, dtype=tf.float32, \
                              parallel_iterations=10, back_prop=False, swap_memory=True, infer_shape=True) #(M, H, W)
        #得到纯净的背景mask(不包含任何形式的crowd)
        gmk_bgd   = tf.cast(tf.equal(tf.reduce_sum(gmks, axis=0, keepdims=True), 0), dtype=tf.float32)     #(1, H, W)
        gmks      = tf.concat([gmks, gmk_bgd], axis=0)                                     #(M+1, H, W)
        #统计各mask的pixels数量
        pix_nums  = tf.reduce_sum(gmks, axis=[1, 2])                                       #(M+1)
        #剔除不合理的mask(包含crowd)
        ncw_idxs0 = tf.where(pix_nums>=self.pix_smp_num)
        ncw_idxs1 = tf.where(gbxs[:, 4]>=0)                                                #最后一个是背景
        ncw_msks  = tf.zeros(shape=[gbx_num], dtype=tf.float32)
        ncw_msks0 = tensor_update(ncw_msks, ncw_idxs0, 1.0)
        ncw_msks1 = tensor_update(ncw_msks, ncw_idxs1, 1.0)
        ncw_msks  = ncw_msks0 * ncw_msks1
        ncw_idxs  = tf.where(tf.equal(ncw_msks, 1.0))
        gbxs      = tf.gather_nd(gbxs, ncw_idxs)                                           #(M, 5)
        gmks      = tf.gather_nd(gmks, ncw_idxs)                                           #(M, H, W)
        gbx_num   = tf.shape(gbxs)[0]                                                      #M
        pix_num   = gbx_num * self.pix_smp_num                                             #M*10

        ##########################Get the loss for embeddings!#########################
        #随机采样K个点
        def sample_pixels(elems=None):
            gmk      = elems                                                               #(H, W)
            pix_idxs = tf.where(tf.equal(gmk, 1.0))                                        #(N, 2)
            pix_idxs = tf.random_shuffle(pix_idxs)[:self.pix_smp_num]                      #(K, 2)
            return pix_idxs
        elems     = gmks
        pix_idxs  = tf.map_fn(sample_pixels, elems, dtype=tf.int64, \
                              parallel_iterations=10, back_prop=False, swap_memory=True, infer_shape=True) #(M, K, 2)
        pix_idxs  = tf.reshape(pix_idxs, [-1, 2])                                          #(M*K,  2)
        pix_embs  = tf.gather_nd(fet_emb, pix_idxs)                                        #(M*K, 64), (H, W, 64), (M*K, 2)
        #生成对应的prediciton
        '''
        sims_pst  = tf.matmul(pix_embs, pix_embs, transpose_b=True) / tf.sqrt(tf.cast(tf.shape(pix_embs)[-1], dtype=tf.float32))
        #sims_pst = tf.sigmoid(sims_pst)                                                   #(M*K, M*K)
        '''
        pix_embs0 = tf.reshape(pix_embs, [pix_num, 1, -1])                                 #(M*K,   1, 64)
        pix_embs1 = tf.reshape(pix_embs, [1, pix_num, -1])                                 #(  1, M*K, 64)
        #pix_embs = pix_embs0 - pix_embs1                                                  #(M*K, M*K, 64)
        pix_embs  = tf.abs(pix_embs0 - pix_embs1)                                          #(M*K, M*K, 64)
        pix_embs  = tf.clip_by_value(pix_embs, 1e-4, 5.0)                                  #(M*K, M*K, 64)
        pix_dsts  = tf.norm(pix_embs, ord='euclidean', axis=-1, keepdims=False) ** 2       #(M*K, M*K)
        #pix_dsts = tf.Print(pix_dsts, [pix_dsts], message=None, first_n=None, summarize=100)
        sims_pst  = 2.0 * (1.0 - tf.sigmoid(pix_dsts))                                     #(M*K, M*K)
        sims_pst  = tf.clip_by_value(sims_pst, 1e-4, 1-1e-4)
        #生成对应的target
        crd_offs  = tf.zeros(shape=[gbx_num, 1], dtype=tf.int32) + self.pix_smp_num        #(M,   1)
        crd_offs  = tf.cumsum(crd_offs, axis=0, exclusive=True, reverse=False)             #(M,   1)
        ycd_idxs  = tf.expand_dims(tf.range(self.pix_smp_num), axis=-1)                    #(K,   1)
        ycd_idxs  = tf.reshape(tf.tile(ycd_idxs, [1, self.pix_smp_num]), [1, -1])          #(1, K*K)
        ycd_idxs  = tf.add(ycd_idxs, crd_offs)                                             #(M, K*K)
        xcd_idxs  = tf.expand_dims(tf.range(self.pix_smp_num), axis= 0)                    #(1,   K)
        xcd_idxs  = tf.reshape(tf.tile(xcd_idxs, [self.pix_smp_num, 1]), [1, -1])          #(1, K*K)
        xcd_idxs  = tf.add(xcd_idxs, crd_offs)                                             #(M, K*K)
        crd_idxs  = tf.stack([ycd_idxs, xcd_idxs], axis=-1)                                #(M, K*K, 2)
        crd_idxs  = tf.reshape(crd_idxs, [-1, 2])                                          #(M*K*K,  2)
        sims_pre  = tf.zeros(shape=[pix_num, pix_num], dtype=tf.float32)                   #(M*K, M*K)
        sims_pre  = tensor_update(sims_pre, crd_idxs, 1.0)                                 #(M*K, M*K)
        #生成对应的loss
        '''
        sims_los  = tf.nn.sigmoid_cross_entropy_with_logits(labels=sims_pre, logits=sims_pst) #(M*K, M*K)
        sims_los  = tf.reduce_sum(sims_los)
        '''
        sims_los  = -(sims_pre * tf.log(sims_pst) + (1.0-sims_pre) * tf.log(1.0-sims_pst)) #(M*K*M*K)
        sims_los  = tf.reduce_sum(sims_los)
        ##########################Get the loss for classfication!#########################
        #再次随机采样K个点
        elems     = gmks
        pix_idxs  = tf.map_fn(sample_pixels, elems, dtype=tf.int64, \
                              parallel_iterations=10, back_prop=False, swap_memory=True, infer_shape=True) #(M, K, 2)
        #生成对应的prediciton
        prbs_pst  = tf.gather_nd(fet_cls, pix_idxs)                                        #(M, K, 4, C), (H, W, 4, C), (M, K, 2)
        #生成对应的target
        pix_idxs  = tf.reshape(pix_idxs, [-1, 2])                                          #(M*K, 2)
        msks      = self.EL.generate_msks_vld(pix_idxs, fet_emb)                           #(M*K, 4, H, W)
        msks      = tf.reshape(msks, [gbx_num, self.pix_smp_num, -1, self.fet_shp[0], self.fet_shp[1]]) #(M, K, 4, H, W)
        #计算每个点可能对应的类别标签
        prbs_pre0 = tf.cast(gbxs[:, -1], dtype=tf.int32)                                   #(M)
        prbs_pre0 = prbs_pre0[:, tf.newaxis, tf.newaxis]                                   #(M, 1, 1)
        prbs_pre0 = tf.tile(prbs_pre0, [1, self.pix_smp_num, self.pix_sim_num])            #(M, K, 4)
        prbs_pre1 = tf.zeros(shape=tf.shape(prbs_pre0), dtype=tf.int32)                    #(M, K, 4)
        #计算由K个点生成的mask和gmks之间的IOU
        gmks      = gmks[:, tf.newaxis, tf.newaxis, :, :]                                  #(M, 1, 1, H, W)
        gmks      = tf.cast(gmks, dtype=tf.bool)                                           #(M, 1, 1, H, W)
        msks      = tf.cast(msks, dtype=tf.bool)                                           #(M, K, 4, H, W)
        msk_iscs  = tf.cast(tf.logical_and(msks, gmks), dtype=tf.float32)                  #(M, K, 4, H, W)
        msk_uins  = tf.cast(tf.logical_or (msks, gmks), dtype=tf.float32)                  #(M, K, 4, H, W)
        msk_iscs  = tf.reduce_sum(msk_iscs, axis=[3, 4])                                   #(M, K, 4)
        msk_uins  = tf.reduce_sum(msk_uins, axis=[3, 4])                                   #(M, K, 4)
        msk_ovps  = tf.where(msk_uins>0, msk_iscs/msk_uins, tf.zeros(shape=tf.shape(msk_uins), dtype=tf.float32)) #(M, K, 4)
        prbs_pre  = tf.where(msk_ovps>=self.msk_ovp_min, prbs_pre0, prbs_pre1)             #(M, K, 4)
        #生成对应的loss
        prbs_los  = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=prbs_pre, logits=prbs_pst) #(M, K, 4)
        prbs_los  = tf.reduce_sum(prbs_los)
        #pix_num  = tf.Print(pix_num, [pix_num, sims_pst], message=None, first_n=None, summarize=100)
        pix_num0  = pix_num ** 2
        pix_num1  = pix_num *  4
        return sims_los, prbs_los, pix_num0, pix_num1

    def generate_embcls_loss(self, fet_emb=None, fet_cls=None, gbxs=None, gmks=None, gbx_nums=None):

        elems = [fet_emb, fet_cls, gbxs, gmks, gbx_nums]
        sims_loss, prbs_loss, pix_nums0, pix_nums1 = \
            tf.map_fn(self.generate_embcls_loss_img, elems, dtype=(tf.float32, tf.float32, tf.int32, tf.int32),
                      parallel_iterations=10, back_prop=True, swap_memory=True, infer_shape=True)
        sims_los = tf.reduce_sum(sims_loss)
        prbs_los = tf.reduce_sum(prbs_loss)
        pix_num0 = tf.cast(tf.reduce_sum(pix_nums0), dtype=tf.float32)
        pix_num1 = tf.cast(tf.reduce_sum(pix_nums1), dtype=tf.float32)
        sims_los = tf.cond(pix_num0>0, lambda: sims_los/pix_num0, lambda: tf.constant(0.0))
        prbs_los = tf.cond(pix_num1>0, lambda: prbs_los/pix_num1, lambda: tf.constant(0.0))
        return sims_los, prbs_los


import warnings
import h5py
import numpy as np
from scipy.spatial import KDTree
from . import utils

def load(fname, **kwargs):
    """Load a .hdf5 GADGET output file."""
    particle_snap = False
    group_snap = False
    with h5py.File(fname) as f:
        try:
            f["Header"].attrs["Ngroups_Total"]
            group_snap = True
        except KeyError:
            try:
                f["PartType1"]
                particle_snap = True
            except KeyError:
                pass
    if group_snap:
        return HaloCatalog(fname, **kwargs)
    elif particle_snap:
        return ParticleData(fname, **kwargs)
    else:
        raise ValueError(f"File {fname} not recognized as a particle or group snapshot")

class Snapshot:
    """Base class to load GADGET-4 particle or group snapshot files.

    Assumptions about GADGET hdf5 file:
        single snapshot file contains all output data for given step
        file name is of the format {}_number.hdf5
        cosmological simulation (in comoving coordinates)
        default code units: h^-1 Mpc, km/s, 1e10 h^-1 Msun
    """

    def __init__(self, fname):
        self.file_name = fname
        try:
            self.snap_num = int(fname.split('/')[-1].split('_')[-1].split('.')[0])
        except ValueError:
            self.snap_num = 0
        with h5py.File(fname) as f:
            self.a = f["Header"].attrs["Time"]
            self.z = f["Header"].attrs["Redshift"]
            self.box_size = f["Header"].attrs["BoxSize"]
            self.softening = f["Parameters"].attrs["SofteningComovingClass0"]
            if self.softening * self.a > f["Parameters"].attrs["SofteningMaxPhysClass0"]:
                self.softening = f["Parameters"].attrs["SofteningMaxPhysClass0"] / self.a

            # Cosmology parameters
            self.OmegaLambda = f["Parameters"].attrs["OmegaLambda"]
            self.OmegaMatter = f["Parameters"].attrs["Omega0"]
            self.Hubble0 = f["Parameters"].attrs["Hubble"] \
                * f["Parameters"].attrs["HubbleParam"]
            self.h = self.Hubble0 / 100
            self.Hubble = self.Hubble0 * np.sqrt(self.OmegaMatter * self.a**-3
                        + (1-self.OmegaMatter-self.OmegaLambda) * self.a**-2
                        + self.OmegaLambda)

class ParticleData(Snapshot):
    """Load GADGET-4 particle snapshots.

    Additional assumptions:
       only dark matter particles - PartType1
       all particles have the same mass

    Note: all data is loaded directly into memory,
      might be a problem for very large simulations.
    For a 256**3 particle sim in single precision a ParticleData object will use ~0.44GB
      a 512**3 particle sim will use ~3.5GB.

    Constructing trees as well will use significantly more memory:
    https://github.com/scipy/scipy/issues/15065.
    For a 512**3 sim, the tree will add around 5GB of memory usage.
    """

    def __init__(self, fname, load_vels=True, load_ids=True, make_tree=False):
        super().__init__(fname)
        with h5py.File(fname) as f:
            self.n_parts = f["Header"].attrs["NumPart_Total"][1]
            self.part_mass = f["Header"].attrs["MassTable"][1] * 1e10
            self.mean_particle_sep = self.box_size / self.n_parts**(1/3)
            self.mean_matter_density = self.n_parts * self.part_mass / self.box_size**3
            self.flat_crit_density = self.mean_matter_density / self.OmegaMatter

            self.pos = f["PartType1"]["Coordinates"][:]
            self.vel = None
            self.ids = None
            if load_vels:
                self.vel = f["PartType1"]["Velocities"][:]
            if load_ids:
                self.ids = f["PartType1"]["ParticleIDs"][:]

            self.tree = None
            if make_tree:
                # make tree box size slightly larger to make sure it accomadates all the data
                tree_box = self.box_size * (1 + 1e-8)
                self.tree = KDTree(self.pos, boxsize=tree_box, leafsize=30)

    def query_radius(self, center, radius, count_only=False):
        """Return indicies of particles within radius of center."""
        if self.tree is None:
            warnings.warn("position tree not constructed, using brute force method.", RuntimeWarning)
            if radius > self.box_size:
                warnings.warn("radius is greater than box_size/2, periodic BC are not implemented with brute force method",
                        RuntimeWarning)
            mask = utils.get_sphere_mask(self.pos, center, radius)
            if count_only:
                return np.sum(mask)
            return np.nonzero(mask)[0]
        else:
            return self.tree.query_ball_point(center, radius, return_length=count_only)

    def query_ids(self, ids):
        """Return indicies of particles given list of ids."""
        if self.ids is None:
            raise RuntimeError("Cannot select particles, snapshot loaded with load_ids=False.")
        return np.nonzero(np.isin(self.ids, ids))[0]

class HaloCatalog(Snapshot):
    """Load GADGET-4 group snapshots (FoF and SUBFIND groups).
    """

    def __init__(self, fname, load_subhalos=False):
        super().__init__(fname)

        with h5py.File(fname) as f:
            self.n_halos = f["Header"].attrs["Ngroups_Total"]
            self.n_subhalos = 0

            if self.n_halos > 0:
                self.pos = f["Group"]["GroupPos"][:]
                self.vel = f["Group"]["GroupVel"][:]
                self.masses = f["Group"]["GroupMass"][:] * 1e10
                self.offsets = f["Group"]["GroupOffsetType"][:,1]
                self.lengths = f["Group"]["GroupLen"][:]

            if load_subhalos:
                try:
                    self.n_subhalos = f["Header"].attrs["Nsubhalos_Total"]
                    self.pos_sh = f["Subhalo"]["SubhaloPos"][:]
                    self.vel_sh = f["Subhalo"]["SubhaloVel"][:]
                    self.masses_sh = f["Subhalo"]["SubhaloMass"][:] * 1e10
                    self.offsets_sh = f["Subhalo"]["SubhaloOffsetType"][:,1]
                    self.lengths_sh = f["Subhalo"]["SubhaloLen"][:]
                except KeyError:
                    pass

    def get_particles(self, halo_i, subhalo=False):
        """Return indicies of particles that are part of the given halo."""
        offset = self.offsets_sh[halo_i] if subhalo else self.offsets[halo_i]
        length = self.lengths_sh[halo_i] if subhalo else self.lengths[halo_i]
        return np.arange(offset, offset+length)

    def get_particle_ids(self, halo_i, particle_data, subhalo=False):
        """Return ids of particles that are part of the given halo."""
        particles = self.get_particles(halo_i, subhalo=subhalo)
        return particle_data.ids[particles]

    def calc_fof_hmf(self, bins):
        """Calculate HMF of FoF clusters."""
        return utils.calc_hmf(bins, self.masses, self.box_size)

    def calc_subhalo_hmf(self, bins):
        """Calculate HMF of SUBFIND clusters."""
        return utils.calc_hmf(bins, self.masses_sh, self.box_size)


import numpy as np
import sys
import gizeh as gz
import random

class SketchGizeh(object):

    def __init__(self):
        self._surface = None

    def setup(self, width=1200, height=900, bg_color=(1, 1, 1)):
        self._width = width
        self._height = height
        self._surface = gz.Surface(width=width, height=height, bg_color=bg_color)

    def draw(self, strokes, scale=1.0, pos=[0, 0], stroke_width=6, color=[0, 0, 0]):
        """iterate through a list of strokes, drawing them on the canvas
        pos is normalised coprdinates in range (0,1)
        """
        try:
            for val in pos:
                if val < 0 or val > 1 or not isinstance(val, float):
                    raise ValueError('position coordinates should be float between (0,1)')
            scale *= np.mean([self._width, self._height]) / 255
            pos[0] = pos[0] * self._width - (scale * (255 / 2))
            pos[1] = pos[1] * self._height - (scale * (255 / 2))
            lines = self._convert_quickdraw_strokes_to_gizeh_group(strokes, color, stroke_width=stroke_width / scale)
            lines = lines.scale(scale).translate(xy=pos)
            lines.draw(self._surface)
        except ValueError as e:
            print(repr(e))

    def draw_person(self, dataset, scale=1.0, position=[0, 0], stroke_width=6):
        body_parts = {'face': [0, 0], 't-shirt': [0, 250], 'pants': [0, 480]}  # dict of parts + translation
        gz_body_parts = []
        for name, pos in body_parts.items():
            strokes = dataset.get_drawing(name, random.randint(1, 1000))
            strokes_gz = self._convert_quickdraw_strokes_to_gizeh_group(strokes, stroke_width=stroke_width / scale)
            strokes_gz = strokes_gz.translate(pos)
            gz_body_parts.append(strokes_gz)
        scale *= np.mean([self._width, self._height]) / 750
        pos[0] = position[0] * self._width - (scale * (255 / 2))
        pos[1] = position[1] * self._height - (scale * (750 / 2))
        gz_body_parts = gz.Group(gz_body_parts).scale(scale).translate(xy=pos)
        gz_body_parts.draw(self._surface)

    def _convert_quickdraw_strokes_to_gizeh_group(self, strokes, color=[0, 0, 0], stroke_width=5):
        lines_list = []
        for stroke in strokes:
            x, y = stroke
            points = list(zip(x, y))
            line = gz.polyline(points=points, stroke=color, stroke_width=stroke_width)
            lines_list.append(line)
        return gz.Group(lines_list)

    def draw_object_recognition_results(self, boxes, classes, scores, labels, dataset, threshold=0.5):
        """draw results of object recognition

        :return: list of objects drawn to the canvas
        """
        drawn_objects = []  # list of the objects drawn
        for i in range(boxes.shape[0]):
            if scores is None or scores[i] >= threshold:
                box = tuple(boxes[i].tolist())
                if classes[i] in labels.keys():
                    class_name = labels[classes[i]]['name']
                    drawn_objects.append(class_name)
                else:
                    raise ValueError('no label for index {}'.format(i))
                ymin, xmin, ymax, xmax = box
                centre = [np.mean([xmin, xmax]), np.mean([ymin, ymax])]
                size = np.mean([xmax - xmin, ymax - ymin])
                if class_name == 'person':
                    self.draw_person(dataset, scale=ymax - ymin, position=centre)
                else:
                    drawing = dataset.get_drawing(class_name, random.randint(1, 1000))
                    self.draw(drawing, scale=size, pos=centre)
        return drawn_objects

    def get_npimage(self):
        return self._surface.get_npimage()

    def save_png(self, path):
        self._surface.write_to_png(str(path))