import argparse
import numpy as np
import math
import os
import time
from six.moves import cPickle  # six.moves is used to self-adjust the change of python 2 and python 3
import yaml
import itertools
from multiprocessing.dummy import Pool as ThreadPool

import torch
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau, StepLR

from model import yolo_layer
from utils.utils import get_random_data, preprocess_true_boxes, add_logger, update_values, set_tb_logger

class Train(object):
    def __init__(self, args, threadpool):

        # initial parameters and model
        self.args = args
        self.annotation_path = args.annotation_path
        self.input_shape = (args.height, args.width)

        self.anchors = args.anchors
        self.class_names = args.classes_names
        self.num_classes = len(self.class_names)
        self.model, self.infos = self._create_model()
        self.optimizer = self._get_optimizer(self.args, self.model)
        self.threadpool = threadpool

        # mkdir log file
        log_name = str(time.strftime("%Y%m%d%H%M%S", time.localtime())) + '_' + '_' + 'bs_' + str(self.args.batch_size)
        print("logging to %s ..." % (log_name))
        self.logger = set_tb_logger('logs', log_name)
        self._train_pipeline(self.annotation_path)
        self.init_lr = args.lr

    def _train_pipeline(self, sample_path):
        '''
            The training pipeline including training and validation samples generation, loss backward
            and checkpoints saving.
        '''
        # random split the data to training and validation samples
        data_gen, data_gen_val, max_batch_ind, max_val_batch_ind = self._train_data_generation(sample_path)

        self.args.n_iter = self.infos.get('n_iter', self.args.n_iter)
        start_epoch = self.infos.get('epoch', self.args.start_epoch)
        best_val_loss = self.infos.get('best_val_loss', 100000)

        # Learning rate self-adjusting depends on the validation loss
        scheduler = ReduceLROnPlateau(self.optimizer, factor=0.1, patience=2, verbose=True, eps=1e-12)
        # scheduler = StepLR(self.optimizer, step_size=3, gamma=0.6)

        # training approach
        for epoch in range(start_epoch, self.args.max_epochs):
            self._train(data_gen, max_batch_ind, epoch)
            val_loss = self._validation(epoch, data_gen_val, max_val_batch_ind)
            scheduler.step(val_loss)
            best_flag = False

            if best_val_loss is None or val_loss < best_val_loss:
                best_val_loss = val_loss
                best_flag = True

            # checkpoints saving
            self.checkpoint_save(self.infos, epoch, best_val_loss, best_flag)

    def cos_lr(self, epoch_start, epoch_max, base_lr):
        curr_lr = base_lr * (1 + math.cos(math.pi*epoch_start / epoch_max)) / 2
        return curr_lr

    def _train(self, data_gen, max_batch_ind, epoch):
        '''
        :param data_gen: generator for training data
        :param max_batch_ind: maximum slice for training data
        :param epoch:
        '''

        torch.set_grad_enabled(True)
        self.model.train()
        tmp_losses = 0
        batch_sum = 0
        start = time.time()
        for batch_ind in range(max_batch_ind):
            data = next(data_gen)
            img, label0, label1, label2 = data[0], data[1], data[2], data[3]
            img = torch.from_numpy(img).float().cuda()
            img = img.view(img.shape[0], img.shape[1], img.shape[2], img.shape[3]).permute(0, 3, 1, 2).contiguous()
            label0 = torch.from_numpy(label0)
            label1 = torch.from_numpy(label1)
            label2 = torch.from_numpy(label2)

            if self.args.use_cuda:
                img = img.cuda()
                label0 = label0.cuda()
                label1 = label1.cuda()
                label2 = label2.cuda()

            losses = self.model(img, label0, label1, label2)
            loss = add_logger(losses, self.logger, batch_ind + max_batch_ind * epoch, 'train')
            loss = loss.sum() / loss.numel()
            tmp_losses = tmp_losses + loss.item()

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            if batch_ind % self.args.display_interval == 0 and batch_ind != 0:
                end = time.time()
                batch_sum = batch_sum + self.args.display_interval
                tmp_losses_show = tmp_losses / batch_sum
                print("step {}/{} (epoch {}), loss: {:f} , lr:{:f}, time/batch = {:.3f}"
                      .format(batch_ind, max_batch_ind, epoch, tmp_losses_show, self.optimizer.param_groups[-1]['lr'],
                              (end - start)/self.args.display_interval))
                start = end

    def _validation(self, epoch, data_gen, max_val_batch_ind):
        torch.set_grad_enabled(False)
        self.model.eval()
        tmp_losses = 0

        for batch_idx in range(max_val_batch_ind):
            data = next(data_gen)
            img, label0, label1, label2 = data[0], data[1], data[2], data[3]
            img = torch.from_numpy(img).float()
            img = img.view(img.shape[0], img.shape[1], img.shape[2], img.shape[3]).permute(0, 3, 1, 2).contiguous()

            label0 = torch.from_numpy(label0)
            label1 = torch.from_numpy(label1)
            label2 = torch.from_numpy(label2)

            if self.args.use_cuda:
                img = img.cuda()
                label0 = label0.cuda()
                label1 = label1.cuda()
                label2 = label2.cuda()

            losses = self.model(img, label0, label1, label2)
            loss = add_logger(losses, self.logger, self.args.n_iter, 'val')
            tmp_losses = tmp_losses + loss.item()

        tmp_losses = tmp_losses / max_val_batch_ind
        print("============================================")
        print("Evaluation Loss (epoch {}), TOTAL_LOSS: {:.3f}".format(epoch, tmp_losses))
        print("============================================")
        return tmp_losses

    @staticmethod
    def _get_optimizer(args, net):
        params = []
        for key, value in dict(net.named_parameters()).items():
            if value.requires_grad:
                if 'backbone' in key:
                    params += [{'params':[value], 'lr':args.backbone_lr}]
                else:
                    params += [{'params':[value], 'lr':args.lr}]

        # Initialize optimizer class
        if args.optimizer == "adam":
            optimizer = optim.Adam(params, weight_decay=args.weight_decay)
        elif args.optimizer == "rmsprop":
            optimizer = optim.RMSprop(params, weight_decay=args.weight_decay)
        else:
            # Default to sgd
            optimizer = optim.SGD(params, momentum=0.9, weight_decay=args.weight_decay,
                                  nesterov=(args.optimizer == "nesterov"))
        return optimizer

    def _create_model(self):
        model = yolo_layer.yolov3layer(self.args)
        infos = {}
        if self.args.start_from != '':
            if self.args.load_best_score == 1:
                model_path = os.path.join(self.args.start_from, 'model-best.pth')
                info_path = os.path.join(self.args.start_from, 'infos-best.pkl')
            else:
                model_path = os.path.join(self.args.start_from, 'model.pth')
                info_path = os.path.join(self.args.start_from, 'infos.pkl')

            if os.path.exists(info_path):
                with open(info_path, 'rb') as f:
                    infos = cPickle.load(f)

            print('Loading the model from %s ...' %(model_path))
            model.load_state_dict(torch.load(model_path))

        if self.args.use_cuda:
            if self.args.mGPUs:
                model = torch.nn.DataParallel(model).cuda()
            else:
                model = model.cuda()
        return model, infos

    def _data_generator(self, annotation_lines, batch_size, input_shape, anchors, num_classes):
        '''data generator for fit_generator'''
        n = len(annotation_lines)
        i = 0
        while True:
            if i + batch_size > n:
                np.random.shuffle(annotation_lines)
                i = 0

            # Separate all image reader into different threads can speed up the program
            output = self.threadpool.starmap(get_random_data, zip(annotation_lines[i:i + batch_size],
                                                                  itertools.repeat(input_shape, batch_size)))
            image_data = list(zip(*output))[0]
            box_data = list(zip(*output))[1]
            i = i + batch_size
            image_data = np.array(image_data)
            box_data = np.array(box_data)
            y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)

            yield [image_data, *y_true]

    def _data_generator_wrapper(self, annotation_lines, batch_size, input_shape, anchors, num_classes):
        n = len(annotation_lines)
        if n == 0 or batch_size <= 0: return None
        return self._data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes)

    def _train_data_generation(self, sample_path):

        with open(sample_path) as f:
            lines = f.readlines()
        np.random.seed(10101)
        np.random.shuffle(lines)
        np.random.seed(None)
        val_split = self.args.val_split
        num_val = int(len(lines) * val_split)
        num_train = len(lines) - num_val

        batch_size = self.args.batch_size
        data_gen = self._data_generator_wrapper(lines[:num_train], batch_size, self.input_shape, self.anchors, self.num_classes)
        data_gen_val = self._data_generator_wrapper(lines[num_train:], batch_size, self.input_shape, self.anchors, self.num_classes)
        max_batch_ind = int(num_train / batch_size)
        max_val_batch_ind = int(num_val / batch_size)

        return data_gen, data_gen_val, max_batch_ind, max_val_batch_ind

    def checkpoint_save(self, infos, epoch, best_val_loss, best_flag=False):

        checkpoint_path = os.path.join(self.args.save_path, self.args.backbones_network)

        if not os.path.exists(checkpoint_path):
            os.makedirs(checkpoint_path)

        if self.args.mGPUs > 1:
            torch.save(self.model.module.state_dict(), os.path.join(checkpoint_path, 'model.pth'))
        else:
            torch.save(self.model.state_dict(), os.path.join(checkpoint_path, 'model.pth'))

        print("model saved to {}".format(checkpoint_path))
        infos['n_iter'] = self.args.n_iter
        infos['epoch'] = epoch
        infos['best_val_loss'] = best_val_loss
        infos['opt'] = self.args

        with open(os.path.join(checkpoint_path, 'infos.pkl'), 'wb') as f:
            cPickle.dump(infos, f)

        if best_flag:
            if self.args.mGPUs > 1:
                torch.save(self.model.module.state_dict(), os.path.join(checkpoint_path, 'model-best.pth'))
            else:
                torch.save(self.model.state_dict(), os.path.join(checkpoint_path, 'model-best.pth'))

            print("model saved to {} with best total loss {:.3f}".format(os.path.join(checkpoint_path, \
                                                                                'model-best.pth'), best_val_loss))

            with open(os.path.join(checkpoint_path, 'infos-best.pkl'), 'wb') as f:
                cPickle.dump(infos, f)

def main():
    args = make_args()
    with open(args.cfg_path, 'r') as handle:
        options_yaml = yaml.load(handle)
    update_values(options_yaml, vars(args))

    # set random seed to cpu and gpu
    if args.seed:
        torch.manual_seed(args.seed)
        if args.use_cuda:
            torch.cuda.manual_seed(args.seed)

    try:
        threadpool = ThreadPool(args.batch_size)
    except Exception as e:
        print(e)
        exit(1)

    Train(args, threadpool)

def make_args():
    # load the optional parameters and update new arguments
    parser = argparse.ArgumentParser()
    # # Data input settings
    parser.add_argument('--cfg_path', type=str, default='cfgs/Yolo_train.yml', help='load config')
    parser.add_argument('--use_cuda', type=bool, default=True, help='whether use gpu.')
    parser.add_argument('--mGPUs', type=bool, default=False, help='whether use mgpu.')
    return parser.parse_args()

if __name__ == '__main__':

    main()


from rnn_model import rnn_encoder, rnn_decoder
import numpy as np
import tensorflow as tf
from utils import Vocabulary
import os

# name, layerSize, hiddenSize, embedSize, sourceVocSize, targetVocSize, rnnType='lstm'
BATCH_SIZE = 32
LAYER_SIZE = 2
HIDDEN_SIZE = 512
EMBED_SIZE = 512
SOURCE_VOC_SIZE = 6000
TARGET_VOC_SIZE = 3000
RNN_TYPE = 'lstm'
LEARNING_RATE = 2e-3
voc = Vocabulary("./cmn.txt", EN_vocSize=SOURCE_VOC_SIZE, CN_vocSize=TARGET_VOC_SIZE)
voc.load_data()
EN_dataset = voc.EN_dataset
EN_seqnums = voc.EN_seqnums
CN_dataset = voc.CN_dataset
CN_seqnums = voc.CN_seqnums
CN_ID2WORD = voc.CN_id2word
CN_WORD2ID = voc.CN_word2id
dataset_nums = int(EN_dataset.shape[0])

def Eng2Chn(sess, sourceSeqVec, sourceSeqNums, targetSeqVec, targetSeqNums, probs, states, states_en, init_state):
    STATE = sess.run([states_en], feed_dict={sourceSeqVec:EN_dataset[5000:5001, :], sourceSeqNums: np.reshape(EN_seqnums[5000], [1])})
    WORD_VAL, STATE = sess.run([tf.argmax(probs, axis=1), states], feed_dict={targetSeqVec: np.reshape(CN_WORD2ID['<START>'], [1, 1]), targetSeqNums: np.array([1]), init_state: STATE})
    vector = []
    vector.append(int(WORD_VAL))
    count = 0
    while CN_ID2WORD[int(WORD_VAL)] != '<EOS>' and count < 15:
        WORD_VAL, STATE = sess.run([tf.argmax(probs, axis=1), states], feed_dict={targetSeqVec: np.reshape(WORD_VAL, [1, 1]), targetSeqNums: np.array([1]), init_state: STATE})
        vector.append(int(WORD_VAL))
        count += 1
    print(voc.vector2sentence(np.array(EN_dataset[5000, 1:-1]), flag=0))
    print(voc.vector2sentence(np.array(vector[:-1]), flag=1))

def train():
    sourceSeqVec = tf.placeholder(tf.int32, [None, None])#English
    targetSeqVec = tf.placeholder(tf.int32, [None, None])#Chinese
    sourceSeqNums = tf.placeholder(tf.int32, [None])
    targetSeqNums = tf.placeholder(tf.int32, [None])
    targetVecLabel = tf.placeholder(tf.int32, [None, None])
    encoder = rnn_encoder("encoder", LAYER_SIZE, HIDDEN_SIZE, EMBED_SIZE, SOURCE_VOC_SIZE, TARGET_VOC_SIZE, RNN_TYPE)
    decoder = rnn_decoder("decoder", LAYER_SIZE, HIDDEN_SIZE, EMBED_SIZE, TARGET_VOC_SIZE, RNN_TYPE)
    #Training phase
    states = encoder(sourceSeqVec, sourceSeqNums, BATCH_SIZE)
    probs, _, _ = decoder(targetSeqVec, targetSeqNums, BATCH_SIZE, states)
    labels = tf.one_hot(targetVecLabel, TARGET_VOC_SIZE)
    labels = tf.reshape(labels, [-1, TARGET_VOC_SIZE])
    loss = tf.reduce_mean(-tf.log(tf.reduce_sum(probs * labels, axis=1) + 1e-10))
    #Test phase
    states_en = encoder(sourceSeqVec, sourceSeqNums, 1)
    probs, states, init_state = decoder(targetSeqVec, targetSeqNums, 1)
    wordVal = tf.argmax(probs, axis=1)
    #Optimization
    Opt = tf.train.AdamOptimizer(LEARNING_RATE).minimize(loss)
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())
    saver = tf.train.Saver()
    # saver.restore(sess, "./save_para/.\\model.ckpt")
    if not os.path.exists("./save_para"):
        os.mkdir("./save_para")

    for i in range(20000):
        rand_select = np.random.randint(0, dataset_nums, [BATCH_SIZE])
        EN_batch = EN_dataset[rand_select]
        EN_batch_seqnums = EN_seqnums[rand_select]
        CN_batch = CN_dataset[rand_select]
        CN_batch_seqnums = CN_seqnums[rand_select]
        sess.run(Opt, feed_dict={sourceSeqVec: EN_batch, sourceSeqNums: EN_batch_seqnums, targetSeqVec: CN_batch[:, :-1], targetSeqNums: CN_batch_seqnums-1, targetVecLabel: CN_batch[:, 1:]})
        if i % 10 == 0:
            LOSS = sess.run(loss, feed_dict={sourceSeqVec: EN_batch, sourceSeqNums: EN_batch_seqnums, targetSeqVec: CN_batch[:, :-1], targetSeqNums: CN_batch_seqnums-1, targetVecLabel: CN_batch[:, 1:]})
            print("Iteration: %d, Loss: %f"%(i, LOSS))
        if i % 500 == 0:
            saver.save(sess, "./save_para/model.ckpt")
            Eng2Chn(sess, sourceSeqVec, sourceSeqNums, targetSeqVec, targetSeqNums, probs, states, states_en, init_state)

if __name__ == "__main__":
    train()


# -*- coding:utf-8 -*-
import numpy as np
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
from keras.models import load_model

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
config.gpu_options.visible_device_list = "4,5,6,7"
set_session(tf.Session(config=config))

model = load_model('keras_test.h5.')
num_samples = 10
height = 224
width = 224
num_classes = 1000
x = np.random.random((num_samples, height, width, 3))
print("X :")
print(x)
classes = model.predict(x, batch_size=256)
print("label :")
print(classes)


import sys
import os
from rdkit import Chem
from rdkit.Chem import AllChem
import numpy as np
import copy

def latent_space(smiles, N_BITS=512):
    m = Chem.MolFromSmiles(smiles)
    if m is None:
        return None
        raise ValueError('SMILES cannot be converted to a RDKit molecules:', smiles)
    m = Chem.AddHs(m)
    return np.array(AllChem.GetMorganFingerprintAsBitVect(m, 2, nBits=N_BITS))

def get_inps():

    f = open(sys.argv[1], 'r')
    lines = f.readlines()
    f.close()

    d = {}
    for line in lines:
        if len(line.strip().split()) > 1:
            a,b = line.strip().split()
            d[a] = b

    if d['script_dir'] == '.':
        d['script_dir'] = ''

    return d

d_inp = get_inps()

X, protein, score = [], [], []
count_er = 0

d = np.load('ligand_data.npy').item()
smiles, target, ic50 = d['smiles'], d['targets'], d['score']
for i in range (len(target)):
    if type(target[i]) == list:
        target[i] = str(target[i][0])
    else:
        target[i] = str(target[i].split(',')[0])

t_k_data = np.load(d_inp['targets_dir'] + 'target_data.npy').item()
#print (t_k_data['5tgz'].shape)
def t_k(tar, dat):
    if tar not in dat:
        #print (tar)
        return None
    return dat[tar]

pd = copy.copy(target)
pdbs = []

for i in range (len(target)):
    target[i] = t_k(target[i], t_k_data)

for i in range (len(smiles)):
    #print (smiles[i])
    if 0 and str(smiles[i]) in ['C[C@H]1CN2[C@@H]([C@@H](C)O1)C1(Cc3cc4c(noc4c(F)c23)C2=[N](C)N=NS2)C(=O)NC(=O)NC1=O',
				'CC(C)(C)OC(=O)N(C1CC1)C1=NC(=C[N]2=C(\C=C3/NC(=O)NC3=O)C=NC12)c1cccc(OC(F)(F)F)c1',
				'O=C1NC(=O)\C(N1)=C\C1=[N]2C=C(N=C(NC3CC3)C2N=C1)C#Cc1ccccc1',
				'O=C1NC(=O)\C(N1)=C\C1=[N]2C=C(N=C(NC3CC3)C2N=C1)c1cccc(CN2CCOCC2)c1',
				'FC(F)(F)Oc1cccc(c1)C1=C[N]2=C(\C=C3/NC(=O)NC3=O)C=NC2C(NC2CC2)=N1',
				'Fc1cccc(c1)C1=C[N]2=C(\C=C3/NC(=O)NC3=O)C=NC2C(NC2CC2)=N1',
				'CC(C)(c1nc(nc2N3CCOC[C@H]3COc12)[N]1=C(N)Nc2ccccc12)S(C)(=O)=O',
				'CC(C)(c1nc(nc2N3CCOC[C@H]3COc12)[N]1=C(Nc2ccccc12)N1CCOCC1)S(C)(=O)=O',
				'ONC(=O)CCCCn1cc(C(=O)Nc2cc(nn2)-c2ccccc2)c(=O)c2ccccc12',
				'[O-][N]1=CC(=CC=C1)c1cnc(NCc2c3CCOc3ccc2F)n2cnnc12',
				'c1ccc(c(c1)C=N(=Cc2ccccc2O)N(CCO)C3=[NH+]CCN3)O',
				'c1cc(ccc1C=N(=Cc2ccc(cc2)Cl)N(CCO)C3=[NH+]CCN3)Cl',
				'c1ccc(cc1)C2=CC(=N(=C(C#N)NC(=O)c3ccccc3)C(=C2)c4ccccc4)c5ccccc5',
				'F[Si-2](F)(F)(F)(F)F.[Na+].[Na+]',
				'[NH4+].[NH4+].F[Si-2](F)(F)(F)(F)F',
				'O=Cl=O']:
        continue
    try:
        try:
            x = latent_space(str(smiles[i]))
            if x is None:
                count_er += 1
                continue
        except KeyError:
            count_er += 1
            continue
    except IndexError:
        count_er += 1
        continue

    if target[i] is None:
        #print('protein error',pd[i], ic50[i])
        count_er += 1
        continue

    score.append(ic50[i])
    protein.append(target[i])
    X.append(x)
    pdbs.append(pd[i])

print(count_er)
np.save('final_data.npy', {'X':X, 'proteins':protein, 'score':score, 'pr_names':pdbs})


def selection_10():

    # Library import
    import numpy
    import matplotlib
    import matplotlib.pyplot   as plt
    import matplotlib.gridspec as gridspec

    # Library version
    matplotlib_version = matplotlib.__version__
    numpy_version      = numpy.__version__

    # Histo binning
    xBinning = numpy.linspace(0.0,2000.0,81,endpoint=True)

    # Creating data sequence: middle of each bin
    xData = numpy.array([12.5,37.5,62.5,87.5,112.5,137.5,162.5,187.5,212.5,237.5,262.5,287.5,312.5,337.5,362.5,387.5,412.5,437.5,462.5,487.5,512.5,537.5,562.5,587.5,612.5,637.5,662.5,687.5,712.5,737.5,762.5,787.5,812.5,837.5,862.5,887.5,912.5,937.5,962.5,987.5,1012.5,1037.5,1062.5,1087.5,1112.5,1137.5,1162.5,1187.5,1212.5,1237.5,1262.5,1287.5,1312.5,1337.5,1362.5,1387.5,1412.5,1437.5,1462.5,1487.5,1512.5,1537.5,1562.5,1587.5,1612.5,1637.5,1662.5,1687.5,1712.5,1737.5,1762.5,1787.5,1812.5,1837.5,1862.5,1887.5,1912.5,1937.5,1962.5,1987.5])

    # Creating weights for histo: y11_PT_0
    y11_PT_0_weights = numpy.array([0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y11_PT_1
    y11_PT_1_weights = numpy.array([0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating a new Canvas
    fig   = plt.figure(figsize=(12,6),dpi=80)
    frame = gridspec.GridSpec(1,1,right=0.7)
    pad   = fig.add_subplot(frame[0])

    # Creating a new Stack
    pad.hist(x=xData, bins=xBinning, weights=y11_PT_0_weights+y11_PT_1_weights,\
             label="$signal4TeV$", rwidth=1.0,\
             color="#ce5e60", edgecolor="#ce5e60", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y11_PT_0_weights,\
             label="$signal1TeV$", rwidth=1.0,\
             color="#5954d8", edgecolor="#5954d8", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    # Axis
    plt.rc('text',usetex=False)
    plt.xlabel(r"p_{T} [ a_{1} ] ",\
               fontsize=16,color="black")
    plt.ylabel(r"$\mathrm{Events}$ $(\mathcal{L}_{\mathrm{int}} = 40.0\ \mathrm{fb}^{-1})$ ",\
               fontsize=16,color="black")

    # Boundary of y-axis
    ymax=(y11_PT_0_weights+y11_PT_1_weights).max()*1.1
    ymin=0 # linear scale
    #ymin=min([x for x in (y11_PT_0_weights+y11_PT_1_weights) if x])/100. # log scale
    plt.gca().set_ylim(ymin,ymax)

    # Log/Linear scale for X-axis
    plt.gca().set_xscale("linear")
    #plt.gca().set_xscale("log",nonposx="clip")

    # Log/Linear scale for Y-axis
    plt.gca().set_yscale("linear")
    #plt.gca().set_yscale("log",nonposy="clip")

    # Legend
    plt.legend(bbox_to_anchor=(1.05,1), loc=2, borderaxespad=0.)

    # Saving the image
    plt.savefig('../../HTML/MadAnalysis5job_0/selection_10.png')
    plt.savefig('../../PDF/MadAnalysis5job_0/selection_10.png')
    plt.savefig('../../DVI/MadAnalysis5job_0/selection_10.eps')

# Running!
if __name__ == '__main__':
    selection_10()


import array
import math
import time

import numpy as np
from enum import Enum, IntEnum
from aenum import auto
import os
import collections
import usb.core
import usb.util
from usb.backend import libusb1 as libusb

class SampleRate(IntEnum):
    """
    Enumeration for the Sample rate of the sounds in the Sound Card
    """

    #: 96KHz sample rate
    _96000HZ = 96000
    #: 192KHz sample rate
    _192000HZ = 192000

class DataType(IntEnum):
    """
    Type of the data to be send to the Sound Card
    """

    #: Integer 32 bits
    INT32 = 0,

    #: Single precision float
    FLOAT32 = 1

class SoundCardErrorCode(Enum):
    OK = 0,
    BAD_USER_INPUT = -1,

    HARP_SOUND_CARD_NOT_DETECTED = -1000,
    NOT_ABLE_TO_SEND_METADATA = auto(),
    NOT_ABLE_TO_READ_METADATA_COMMAND_REPLY = auto(),
    METADATA_COMMAND_REPLY_NOT_CORRECT = auto(),
    NOT_ABLE_TO_SEND_DATA = auto(),
    NOT_ABLE_TO_READ_DATA_COMMAND_REPLY = auto(),
    DATA_COMMAND_REPLY_NOT_CORRECT = auto(),
    NOT_ABLE_TO_SEND_READ_METADATA = auto(),
    NOT_ABLE_TO_READ_READ_METADATA_COMMAND_REPLY = auto(),
    READ_METADATA_COMMAND_REPLY_NOT_CORRECT = auto(),

    BAD_SOUND_INDEX = -1020,
    BAD_SOUND_LENGTH = auto(),
    BAD_SAMPLE_RATE = auto(),
    BAD_DATA_TYPE = auto(),
    DATA_TYPE_DO_NOT_MATCH = auto(),
    BAD_DATA_INDEX = auto(),

    PRODUCING_SOUND = -1030,
    STARTED_PRODUCING_SOUND = auto(),

    NOT_ABLE_TO_OPEN_FILE = -1040

class SoundMetadata(object):

    def __init__(self, sound_index, sound_length, sample_rate, data_type):
        """

        :param self:
        :param sound_index: Sound index in the soundcard (2 -> 31 since 0 and 1 are reserved)
        :param sound_length: Sound length in number of samples
        :param sample_rate: Sample rate
        :param data_type: 0 for Int32 and 1 for Float32 (not available right now)
        """
        self._sound_index = sound_index
        self._sound_length = sound_length
        self._sample_rate = sample_rate
        self._data_type = data_type

    def check_data(self):
        if self._sound_index < 2 or self._sound_index > 32:
            return SoundCardErrorCode.BAD_SOUND_INDEX

        if self._sound_length < 16:
            return SoundCardErrorCode.BAD_SOUND_LENGTH

        if self._sample_rate is not SampleRate._96000HZ and self._sample_rate is not SampleRate._192000HZ:
            return SoundCardErrorCode.BAD_SAMPLE_RATE

        if self._data_type is not DataType.INT32 and self._data_type is not DataType.FLOAT32:
            return SoundCardErrorCode.BAD_DATA_TYPE

        if self._sound_index == 0 and self._data_type is not DataType.FLOAT32:
            return SoundCardErrorCode.DATA_TYPE_DO_NOT_MATCH

        if self._sound_index == 1 and self._data_type is not DataType.FLOAT32:
            return SoundCardErrorCode.DATA_TYPE_DO_NOT_MATCH

        if self._sound_index > 1 and self._data_type is not DataType.INT32:
            return SoundCardErrorCode.DATA_TYPE_DO_NOT_MATCH

        return SoundCardErrorCode.OK

    def as_array(self):
        return np.array([self._sound_index, self._sound_length, self._sample_rate, self._data_type], dtype=np.int32)

class SoundCardModule(object):
    """
    Provides access to the Harp Sound Card. It allows to send and read the sounds in the Sound Card, through a normal
    USB connection.
    """

    def __init__(self, device=None):
        """
        If a libUSB's device is given, it will try to open it. If none is given it will try to connect to the first Sound Card that is connected to the computer.

        :param device: (Optional) libUSB device to use. If nothing is passed, it will try to connect automatically.
        """
        self._backend = libusb.get_backend()
        try:
            self._devices = list(usb.core.find(backend=self._backend, idVendor=0x04d8, idProduct=0xee6a, find_all=True))
        except OSError as e:
            pass

        self._dev = self._devices[0] if self._devices else None
        self._cfg = None
        self._port = None
        self._connected = False

        self.open(self._dev if device is None else device)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def open(self, device=None):
        """
        Opens the connection to the Sound Card. If no device is given, it will try to connect to the first Sound Card that is connected to the computer.

        :param device: (Optional) Already initialized libUSB's device to use.
        """
        if device is None:
            self._backend = libusb.get_backend()
            try:
                self._dev = usb.core.find(backend=self._backend, idVendor=0x04d8, idProduct=0xee6a)
            except OSError as e:
                self._dev = None
                pass
        else:
            self._dev = device

        if self._dev is None:
            print(
                "Unable to connect to the Sound Card through the USB port. You will be unable to send and receive sounds.")
        else:
            # set the active configuration. With no arguments, the first configuration will be the active one
            # note: some devices reset when setting an already selected configuration so we should check for it before
            self._cfg = self._dev.get_active_configuration()
            if self._cfg is None or self._cfg.bConfigurationValue != 1:
                self._dev.set_configuration(1)

        self._connected = True if self._dev else False

    @property
    def devices(self):
        return self._devices

    @property
    def connected(self):
        return self._connected

    def close(self):
        """
        Closes the connection with the Sound Card. It will close USB connection (to read and save sounds)
        """
        if self._dev:
            usb.util.dispose_resources(self._dev)

    def reset(self):
        """
        Resets the device, waits 700ms and tries to connect again so that the current instance of the SoundCard object can still be used.

        .. note:: Necessary at the moment after sending a sound.
        """
        if not self._dev:
            raise Exception("Sound card might not be connected. Please connect it before any operation.")

        # Reset command length:    'c' 'm' 'd' '0x88' + 'f'
        reset_cmd = [ord('c'), ord('m'), ord('d'), 0x88, ord('f')]
        # cmd = 'cmd' + chr(0x88) + 'f'
        wrt = self._dev.write(1, reset_cmd, 100)
        assert wrt == len(reset_cmd)

        time.sleep(700.0 / 1000.0)
        self.open()

    def read_sounds(self, output_folder=None, sound_index=None, clean_dst_folder=True):
        """
        Reads sounds from the sound card.

        .. note:: by default, it will clear the destination folder of all data. It will also write by default to a
                  "from_soundcard" folder in the working directory if none is given.

        :param output_folder: Destination folder's path.
        :param sound_index: If a sound_index is given, it will get only that sound, if nothing is passed it will gather all sounds from all indexes.
        :param clean_dst_folder: Flag that defines if the method should clean the destination folder or not

        """
        if not self._dev:
            raise Exception("Sound card might not be connected. Please connect it before any operation.")

        # admit that if the output_folder is None, write inside a 'from_soundcard' folder in the current directory
        if not output_folder:
            output_folder = os.path.join(os.getcwd(), 'from_soundcard')
            if not os.path.isdir(output_folder):
                os.makedirs(output_folder)
        else:
            # create folder if it doesn't exists
            if not os.path.exists(output_folder):
                os.makedirs(output_folder)

            if clean_dst_folder:
                for file in os.listdir(output_folder):
                    file_path = os.path.join(output_folder, file)
                    try:
                        if os.path.isfile(file_path):
                            os.unlink(file_path)
                    except Exception as e:
                        # probably a permissions error while deleting, ignore and try the next one
                        print("Error occurred when deleting file '{file_path}'. Ignoring error and continuing.".format(file_path=file_path))
                        continue

        if sound_index is None:
            for i in range(2, 32):
                self._from_soundcard(output_folder, i)
        else:
            self._from_soundcard(output_folder, sound_index)

        print("All files read!")

    def send_sound(self, wave_int, sound_index, sample_rate, data_type,
                   sound_filename=None, metadata_filename=None,
                   description_filename=None):
        """
            This method will send the sound to the Harp Sound Card as a byte array (int8)

            :param wave_int: NumPy array as int32 that represents the sound data
            :param sound_index:  The destination index in the Sound Card (>=2 and <= 32)
            :param sample_rate: The SampleRate enum value for either 96KHz or 192KHz
            :param data_type: The DataType enum value for either Int32 or Float32 (not implemented yet in the hardware)
            :param sound_filename: The name of the sound filename to be saved with the sound in the board (str)
            :param metadata_filename: The name of the metadata filename to be saved with the sound in the board (str)
            :param description_filename: The name of the description filename to be saved with the sound in the board (str)
        """
        self._to_soundcard(wave_int, sound_index, sample_rate, data_type, sound_filename, metadata_filename,
                          description_filename)

    def _from_soundcard(self, output_folder=None, sound_index=None):
        """
        Reads sounds from the sound card.

        :param output_folder: Destination folder's path.
        :param sound_index: If a sound_index is given, it will get only that sound, if nothing is passed it will gather all sounds from all indexes.

        """
        if not self._dev:
            raise Exception("Sound card might not be connected. Please connect it before any operation.")

        if sound_index is None or sound_index < 2 or sound_index > 31:
            raise Exception("sound_index must have a value between 2 and 31")

        metadata = self.__get_metadata_from_device(sound_index)

        if metadata is None:
            raise Exception('SoundCardModule: Error while getting metadata from device')

        # define prefix
        prefix = 'i'
        if sound_index < 9:
            prefix += '0' + str(sound_index) + '_'
        else:
            prefix += str(sound_index) + '_'

        sound_filename = metadata.sound_filename.decode('utf-8')
        metadata_filename = metadata.metadata_filename.decode('utf-8') if metadata.metadata_filename else None
        description_filename = metadata.description_filename.decode(
            'utf-8') if metadata.description_filename else None

        if prefix not in sound_filename:
            sound_filename = prefix + sound_filename
        if metadata_filename and prefix not in metadata_filename:
            metadata_filename = prefix + metadata_filename
        if description_filename and prefix not in description_filename:
            description_filename = prefix + description_filename

        if metadata.has_sound:
            with open(os.path.join(output_folder, sound_filename), 'w', encoding='utf8') as f:
                # TODO: read the sound so we can write it here
                f.write('TODO')

        if metadata.has_metadata:
            with open(os.path.join(output_folder, metadata_filename), 'wb') as f:
                # clean the zeros at the end
                f.write(metadata.metadata_array.tobytes().strip(b'\0'))

        if metadata.has_description:
            with open(os.path.join(output_folder, description_filename), 'wb') as f:
                f.write(metadata.description.tobytes().strip(b'\0'))

        # create summary info file
        if metadata.has_sound:
            with open(os.path.join(output_folder, sound_filename + '.metadata.txt'), 'w') as f:
                f.write('SOUND_INDEX = ' + str(sound_index))

                used_pos = math.ceil(metadata.sound_length / (33554432.0 * 2.0 / 32.0)) - 1
                if used_pos > 0:
                    f.write(", ")
                f.write(", ".join(str(sound_index + idx + 1) for idx in range(used_pos)))

                f.write("\n")

                f.write("TOTAL_SAMPLES = " + str(metadata.sound_length) + "\n")
                f.write(
                    "TOTAL_LENGTH_MS = " + str(int(metadata.sound_length / 2 / metadata.sample_rate * 1000)) + "\n")
                f.write("SAMPLE_RATE = " + str(metadata.sample_rate) + "\n")

                if metadata.data_type == 0:
                    f.write("DATA_TYPE = Int32\n")
                else:
                    f.write("DATA_TYPE = Float32\n")

                f.write("SOUND_FILENAME = " + sound_filename + "\n")

                if metadata.has_metadata:
                    f.write("USER_METADATA_FILENAME = " + metadata_filename + "\n")
                if metadata.has_description:
                    f.write("USER_DESCRIPTION_FILENAME = " + description_filename + "\n")

    def _to_soundcard(self, wave_int, sound_index, sample_rate, data_type,
                      sound_filename=None, metadata_filename=None,
                      description_filename=None):
        """
            This method will send the sound to the Harp Sound Card as a byte array (int8)

            :param wave_int: NumPy array as int32 that represents the sound data
            :param sound_index:  The destination index in the Sound Card (>=2 and <= 32)
            :param sample_rate: The SampleRate enum value for either 96KHz or 192KHz
            :param data_type: The DataType enum value for either Int32 or Float32 (not implemented yet in the hardware)
            :param sound_filename: The name of the sound filename to be saved with the sound in the board (str)
            :param metadata_filename: The name of the metadata filename to be saved with the sound in the board (str)
            :param description_filename: The name of the description filename to be saved with the sound in the board (str)
        """
        # confirm that the dev exists and is ready
        if not self._dev:
            raise EnvironmentError(
                'Sound card not initialized. Please call the initialize method before any operation.')

        int32_size = np.dtype(np.int32).itemsize
        # work with a int8 view of the wave_int (which is int32)
        wave_int8 = wave_int.view(np.int8)

        # get number of commands to send
        sound_file_size_in_samples = len(wave_int8) // 4
        commands_to_send = int(sound_file_size_in_samples * 4 // 32768 + (
            1 if ((sound_file_size_in_samples * 4) % 32768) is not 0 else 0))

        # Metadata command length: 'c' 'm' 'd' '0x80' + random + metadata  + 32768 + 2048 + 'f'
        metadata_cmd_header_size = 4 + int32_size + (4 * int32_size)
        metadata_cmd = np.zeros(metadata_cmd_header_size + 32768 + 2048 + 1, dtype=np.int8)

        metadata_cmd[0] = ord('c')
        metadata_cmd[1] = ord('m')
        metadata_cmd[2] = ord('d')
        metadata_cmd[3] = 0x80
        metadata_cmd[-1] = ord('f')

        rand_val = np.random.randint(-32768, 32768, size=1, dtype=np.int32)
        # copy that random data
        metadata_cmd[4: 4 + int32_size] = rand_val.view(np.int8)

        # create metadata info and add it to the metadata_cmd
        metadata = SoundMetadata(sound_index, sound_file_size_in_samples, sample_rate, data_type)
        if metadata.check_data() is not SoundCardErrorCode.OK:
            print("Input data incorrect, please correct it before proceeding.")
            return

        metadata_cmd[8: 8 + (4 * int32_size)] = metadata.as_array().view(np.int8)

        # add first data block of data to the metadata_cmd
        metadata_cmd_data_index = metadata_cmd_header_size
        metadata_cmd[metadata_cmd_data_index: metadata_cmd_data_index + 32768] = wave_int8[0: 32768]

        # prepare user_metadata
        # [0:169]     sound_filename
        # [170:339]   metadata_filename
        # [340:511]   description_filename
        # [512:1535]  metadata_filename content
        # [1536:2047] description_filename content
        user_metadata = np.zeros(2048, dtype=np.int8)
        user_metadata_index = metadata_cmd_data_index + 32768

        if sound_filename:
            tmp = bytearray()
            tmp.extend(map(ord, os.path.basename(sound_filename)))
            tmp_size = len(tmp) if len(tmp) < 169 else 169
            user_metadata[0:tmp_size] = tmp[0:tmp_size]

        if metadata_filename:
            tmp = bytearray()
            tmp.extend(map(ord, os.path.basename(metadata_filename)))
            tmp_size = len(tmp) if len(tmp) < 169 else 169
            user_metadata[170: 170 + tmp_size] = tmp[0:tmp_size]

            # get file contents, truncate data if required
            try:
                with open(metadata_filename, 'r', encoding='utf8') as f:
                    text = f.read()
                    text_tmp = bytearray()
                    text_tmp.extend(map(ord, text))
                    data_tmp = np.array(text_tmp)
                    data = data_tmp.view(np.int8)
                    data_size = len(data) if len(data) < 1023 else 1023
                    user_metadata[512: 512 + data_size] = data[0: data_size]
            except OSError as e:
                # TODO: should be a stronger error
                print("Error opening metadata file.")

        if description_filename:
            tmp = bytearray()
            tmp.extend(map(ord, os.path.basename(description_filename)))
            tmp_size = len(tmp) if len(tmp) < 169 else 169
            user_metadata[340: 340 + tmp_size] = tmp[0: tmp_size]

            # get file contents, truncate data if required
            try:
                with open(description_filename, 'r', encoding='utf8') as f:
                    text = f.read()
                    text_tmp = bytearray()
                    text_tmp.extend(map(ord, text))
                    data_tmp = np.array(text_tmp)
                    data = data_tmp.view(np.int8)
                    data_size = len(data) if len(data) < 511 else 511
                    user_metadata[1536: 1536 + data_size] = data[0: data_size]
            except OSError as e:
                print(e)
                # TODO: should be a stronger error
                print("Error opening description file.")

        # add user metadata (2048 bytes) to metadata_cmd
        metadata_cmd[user_metadata_index: user_metadata_index + 2048] = user_metadata

        # Metadata command reply: 'c' 'm' 'd' '0x80' + random + error
        metadata_cmd_reply = array.array('b', [0] * (4 + int32_size + int32_size))

        # send metadata_cmd and get it's reply
        try:
            res_write = self._dev.write(0x01, metadata_cmd.tobytes(), 100)
        except usb.core.USBError as e:
            # TODO: we probably should try again
            print("something went wrong while writing to the device")
            return

        assert res_write == len(metadata_cmd)

        try:
            ret = self._dev.read(0x81, metadata_cmd_reply, 1000)
        except usb.core.USBError as e:
            # TODO: we probably should try again

            print("something went wrong while reading from the device")
            return

        # get the random received and the error received from the reply command
        rand_val_received = int.from_bytes(metadata_cmd_reply[4: 4 + int32_size], byteorder='little', signed=True)
        error_received = int.from_bytes(metadata_cmd_reply[8: 8 + int32_size], byteorder='little', signed=False)

        assert rand_val_received == rand_val[0]
        assert error_received == 0

        # prepare command to send and to receive
        # Data command length:     'c' 'm' 'd' '0x81' + random + dataIndex + 32768 + 'f'
        data_cmd = np.zeros(4 + int32_size + int32_size + 32768 + 1, dtype=np.int8)
        data_cmd_data_index = 4 + int32_size + int32_size

        data_cmd[0] = ord('c')
        data_cmd[1] = ord('m')
        data_cmd[2] = ord('d')
        data_cmd[3] = 0x81
        data_cmd[-1] = ord('f')

        # Data command reply:     'c' 'm' 'd' '0x81' + random + error
        data_cmd_reply = array.array('b', [0] * (4 + int32_size + int32_size))

        # loop to send the rest of the commands
        # check reply for each command sent
        for i in range(1, commands_to_send):
            # it has to be as an np.array of int32 so that we can get a view as int8s
            rand_val = np.random.randint(-32768, 32768, size=1, dtype=np.int32)
            # copy that random data
            data_cmd[4: 4 + int32_size] = rand_val.view(np.int8)

            # write dataIndex to the data_cmd (2 ints size)
            data_cmd[8: 8 + int32_size] = np.array([i], dtype=np.int32).view(np.int8)

            # write data from wave_int to cmd
            wave_idx = i * 32768
            data_block = wave_int8[wave_idx: wave_idx + 32768]
            data_cmd[data_cmd_data_index: data_cmd_data_index + len(data_block)] = data_block

            # send data to device
            try:
                res_write = self._dev.write(0x01, data_cmd.tobytes(), 100)
            except usb.core.USBError as e:
                # TODO: we probably should try again
                print("something went wrong while writing to the device")
                return

            # TODO: we probably should try again
            assert res_write == len(data_cmd)

            try:
                ret = self._dev.read(0x81, data_cmd_reply, 400)
            except usb.core.USBError as e:
                # TODO: we probably should try again

                print("something went wrong while reading from the device")
                return

            # get the random received and the error received from the reply command
            rand_val_received = int.from_bytes(data_cmd_reply[4: 4 + int32_size], byteorder='little', signed=True)
            error_received = int.from_bytes(data_cmd_reply[8: 8 + int32_size], byteorder='little', signed=False)

            assert rand_val_received == rand_val[0]
            assert error_received == 0

    def __get_metadata_from_device(self, sound_index):
        int32_size = np.dtype(np.int32).itemsize

        # Read metadata command length: 'c' 'm' 'd' '0x84' + random + soundIndex + 'f'
        read_metadata_cmd = np.zeros(4 + int32_size + int32_size + 1, dtype=np.int8)
        read_metadata_cmd[0] = ord('c')
        read_metadata_cmd[1] = ord('m')
        read_metadata_cmd[2] = ord('d')
        read_metadata_cmd[3] = 0x84
        read_metadata_cmd[-1] = ord('f')

        rand_val = np.random.randint(-32768, 32768, size=1, dtype=np.int32)
        # copy that random data
        read_metadata_cmd[4: 4 + int32_size] = rand_val.view(np.int8)
        read_metadata_cmd[8: 8 + int32_size] = np.array([sound_index], dtype=np.int32).view(np.int8)

        # prepare to send command and receive the reply
        read_reply_cmd = array.array('b', [0] * (4 + 6 * int32_size + 2048))

        try:
            res_write = self._dev.write(0x01, read_metadata_cmd.tobytes(), 100)
        except usb.core.USBError as e:
            # TODO: we probably should try again
            print("something went wrong while writing to the device")
            return

        assert res_write == len(read_metadata_cmd)

        try:
            ret = self._dev.read(0x81, read_reply_cmd, 100)
        except usb.core.USBError as e:
            # TODO: we probably should try again

            print("something went wrong while reading from the device")
            return

        metadata = collections.namedtuple('Metadata',
                                          ['metadata_array', 'description', 'bit_mask', 'sound_length', 'data_type',
                                           'sample_rate', 'sound_filename', 'metadata_filename',
                                           'description_filename',
                                           'has_sound', 'has_metadata', 'has_description'])

        # get data from the reply array
        metadata.metadata_array = array.array('b', [0] * 1024)
        metadata.description = array.array('b', [0] * 512)

        # get the random received and the error received from the reply command
        rand_val_received = int.from_bytes(read_reply_cmd[4: 4 + int32_size], byteorder='little', signed=True)
        error_received = int.from_bytes(read_reply_cmd[8: 8 + int32_size], byteorder='little', signed=False)

        assert rand_val_received == rand_val[0]
        assert error_received == 0

        # bitmask
        metadata.bit_mask = int.from_bytes(read_reply_cmd[12:12 + int32_size + int32_size], byteorder='little',
                                           signed=True)

        metadata.has_sound = metadata.bit_mask & (1 << sound_index) == (1 << sound_index)

        metadata.sound_length = int.from_bytes(read_reply_cmd[16:16 + int32_size], byteorder='little', signed=True)
        metadata.sample_rate = int.from_bytes(read_reply_cmd[20:20 + int32_size], byteorder='little', signed=True)
        metadata.data_type = int.from_bytes(read_reply_cmd[24:24 + int32_size], byteorder='little', signed=True)

        metadata.sound_filename = read_reply_cmd[28:170].tobytes().strip(b'\0')

        metadata.has_metadata = False
        metadata.metadata_filename = ''
        if read_reply_cmd[28 + 170]:
            metadata.has_metadata = True
            metadata.metadata_array[0:1024] = read_reply_cmd[28 + 512:28 + 512 + 1024]
            metadata.metadata_filename = read_reply_cmd[28 + 170: 28 + 170 + 170].tobytes().strip(b'\0')

        metadata.has_description = False
        metadata.description_filename = ''
        if read_reply_cmd[28 + 170 + 170]:
            metadata.has_description = True
            metadata.description[0:512] = read_reply_cmd[28 + 512 + 1024:28 + 512 + 1024 + 512]
            metadata.description_filename = read_reply_cmd[28 + 170 + 170: 28 + 170 + 170 + 170].tobytes().strip(b'\0')

        return metadata


"""Calculate ph-ph interaction and phonons on grid."""
# Copyright (C) 2020 Atsushi Togo
# All rights reserved.
#
# This file is part of phono3py.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import warnings

import numpy as np
from phonopy.harmonic.dynamical_matrix import get_dynamical_matrix
from phonopy.structure.cells import (
    Primitive,
    compute_all_sg_permutations,
    sparse_to_dense_svecs,
)
from phonopy.structure.symmetry import Symmetry
from phonopy.units import AMU, EV, Angstrom, Hbar, THz, VaspToTHz

from phono3py.phonon3.real_to_reciprocal import RealToReciprocal
from phono3py.phonon3.reciprocal_to_normal import ReciprocalToNormal
from phono3py.phonon3.triplets import get_nosym_triplets_at_q, get_triplets_at_q
from phono3py.phonon.grid import (
    BZGrid,
    get_grid_points_by_rotations,
    get_ir_grid_points,
)
from phono3py.phonon.solver import run_phonon_solver_c, run_phonon_solver_py

class Interaction:
    """Calculate ph-ph interaction and phonons on grid.

    This class instance is the heart of phono3py calculation.
    Many data are stored.

    The following three steps have to be done manually.
    1) init_dynamical_matrix
    2) set_grid_point
    3) run

    Attributes
    ----------
    interaction_strength
    mesh_numbers
    is_mesh_symmetry
    fc3
    dynamical_matrix
    primitive
    primitive_symmetry
    bz_grid
    band_indices
    nac_params
    nac_q_direction
    zero_value_positions
    frequency_factor_to_THz
    lapack_zheev_uplo
    cutoff_frequency

    """

    def __init__(
        self,
        primitive: Primitive,
        bz_grid: BZGrid,
        primitive_symmetry: Symmetry,
        fc3=None,
        band_indices=None,
        constant_averaged_interaction=None,
        frequency_factor_to_THz=VaspToTHz,
        frequency_scale_factor=None,
        unit_conversion=None,
        is_mesh_symmetry=True,
        symmetrize_fc3q=False,
        cutoff_frequency=None,
        lapack_zheev_uplo="L",
    ):
        """Init method."""
        self._primitive = primitive
        self._bz_grid = bz_grid
        self._primitive_symmetry = primitive_symmetry

        self._band_indices = None
        self._set_band_indices(band_indices)
        self._constant_averaged_interaction = constant_averaged_interaction
        self._frequency_factor_to_THz = frequency_factor_to_THz
        self._frequency_scale_factor = frequency_scale_factor

        if fc3 is not None:
            self._set_fc3(fc3)

        # Unit to eV^2
        if unit_conversion is None:
            num_grid = np.prod(self.mesh_numbers)
            self._unit_conversion = (
                (Hbar * EV) ** 3
                / 36
                / 8
                * EV ** 2
                / Angstrom ** 6
                / (2 * np.pi * THz) ** 3
                / AMU ** 3
                / num_grid
                / EV ** 2
            )
        else:
            self._unit_conversion = unit_conversion
        if cutoff_frequency is None:
            self._cutoff_frequency = 0
        else:
            self._cutoff_frequency = cutoff_frequency
        self._is_mesh_symmetry = is_mesh_symmetry
        self._symmetrize_fc3q = symmetrize_fc3q
        self._lapack_zheev_uplo = lapack_zheev_uplo

        self._symprec = self._primitive_symmetry.tolerance

        self._triplets_at_q = None
        self._weights_at_q = None
        self._triplets_map_at_q = None
        self._ir_map_at_q = None
        self._interaction_strength = None
        self._g_zero = None

        self._phonon_done = None
        self._frequencies = None
        self._eigenvectors = None
        self._dm = None
        self._nac_params = None
        self._nac_q_direction = None

        self._band_index_count = 0

        svecs, multi = self._primitive.get_smallest_vectors()
        if self._primitive.store_dense_svecs:
            self._svecs = svecs
            self._multi = multi
        else:
            self._svecs, self._multi = sparse_to_dense_svecs(svecs, multi)

        self._masses = np.array(self._primitive.masses, dtype="double")
        self._p2s = np.array(self._primitive.p2s_map, dtype="int_")
        self._s2p = np.array(self._primitive.s2p_map, dtype="int_")

    def run(self, lang="C", g_zero=None):
        """Run ph-ph interaction calculation."""
        if (self._phonon_done == 0).any():
            self.run_phonon_solver()

        num_band = len(self._primitive) * 3
        num_triplets = len(self._triplets_at_q)

        self._interaction_strength = np.empty(
            (num_triplets, len(self._band_indices), num_band, num_band), dtype="double"
        )
        if self._constant_averaged_interaction is None:
            self._interaction_strength[:] = 0
            if lang == "C":
                self._run_c(g_zero)
            else:
                self._run_py()
        else:
            num_grid = np.prod(self.mesh_numbers)
            self._interaction_strength[:] = (
                self._constant_averaged_interaction / num_grid
            )

    @property
    def interaction_strength(self):
        """Return ph-ph interaction strength.

        Returns
        -------
        ndarray
            shape=(num_ir_grid_points, num_specified_band, num_band, num_band),
            dtype='double', order='C'

        """
        return self._interaction_strength

    def get_interaction_strength(self):
        """Return ph-ph interaction strength."""
        warnings.warn(
            "Use attribute, Interaction.interaction_strength "
            "instead of Interaction.get_interaction_strength().",
            DeprecationWarning,
        )
        return self.interaction_strength

    @property
    def mesh_numbers(self):
        """Return mesh numbers.

        Returns
        -------
        ndarray
           shape=(3, ), dtype='int_'

        """
        return self._bz_grid.D_diag

    def get_mesh_numbers(self):
        """Return mesh numbers."""
        warnings.warn(
            "Use attribute, Interaction.mesh_numbers "
            "instead of Interaction.get_mesh_numbers().",
            DeprecationWarning,
        )
        return self.mesh_numbers

    @property
    def is_mesh_symmetry(self):
        """Whether symmetry of grid is utilized or not."""
        return self._is_mesh_symmetry

    @property
    def fc3(self):
        """Return fc3."""
        return self._fc3

    def get_fc3(self):
        """Return fc3."""
        warnings.warn(
            "Use attribute, Interaction.fc3 " "instead of Interaction.get_fc3().",
            DeprecationWarning,
        )
        return self.fc3

    @property
    def dynamical_matrix(self):
        """Return DynamicalMatrix class instance."""
        return self._dm

    def get_dynamical_matrix(self):
        """Return DynamicalMatrix class instance."""
        warnings.warn(
            "Use attribute, Interaction.dynamical_matrix "
            "instead of Interaction.get_dynamical_matrix().",
            DeprecationWarning,
        )
        return self.dynamical_matrix

    @property
    def primitive(self):
        """Return Primitive class instance."""
        return self._primitive

    def get_primitive(self):
        """Return Primitive class instance."""
        warnings.warn(
            "Use attribute, Interaction.primitive "
            "instead of Interaction.get_primitive().",
            DeprecationWarning,
        )
        return self.primitive

    @property
    def primitive_symmetry(self):
        """Return Symmetry class instance of primitive cell."""
        return self._primitive_symmetry

    def get_triplets_at_q(self):
        """Return grid point triplets information.

        triplets_at_q is in BZ-grid.
        triplets_map_at_q is in GR-grid.
        ir_map_at_q is in GR-grid.
        See details at ``get_triplets_at_q``.

        """
        return (
            self._triplets_at_q,
            self._weights_at_q,
            self._triplets_map_at_q,
            self._ir_map_at_q,
        )

    @property
    def bz_grid(self):
        """Return BZGrid class instance."""
        return self._bz_grid

    @property
    def band_indices(self):
        """Return band indices.

        Returns
        -------
        ndarray
            shape=(num_specified_bands, ), dtype='int_'

        """
        return self._band_indices

    def get_band_indices(self):
        """Return band indices."""
        warnings.warn(
            "Use attribute, Interaction.band_indices "
            "instead of Interaction.get_band_indices().",
            DeprecationWarning,
        )
        return self.band_indices

    @property
    def nac_params(self):
        """Return NAC params."""
        return self._nac_params

    @property
    def nac_q_direction(self):
        """Return q-direction used for NAC at q->0.

        Direction of q-vector watching from Gamma point used for
        non-analytical term correction. This is effective only at q=0
        (physically q->0). The direction is given in crystallographic
        (fractional) coordinates.
        shape=(3,), dtype='double'.
        Default value is None, which means this feature is not used.

        """
        return self._nac_q_direction

    @nac_q_direction.setter
    def nac_q_direction(self, nac_q_direction):
        if nac_q_direction is None:
            self._nac_q_direction = None
        else:
            self._nac_q_direction = np.array(nac_q_direction, dtype="double")

    def get_nac_q_direction(self):
        """Return q-direction used for NAC at q->0."""
        warnings.warn(
            "Use attribute, Interaction.nac_q_direction "
            "instead of Interaction.get_nac_q_direction().",
            DeprecationWarning,
        )
        return self.nac_q_direction

    def set_nac_q_direction(self, nac_q_direction=None):
        """Set NAC q-point direction valid at q->0."""
        warnings.warn(
            "Use attribute, Interaction.nac_q_direction "
            "instead of Interaction.set_nac_q_direction().",
            DeprecationWarning,
        )
        self.nac_q_direction = nac_q_direction

    @property
    def zero_value_positions(self):
        """Return zero ph-ph interaction elements information.

        Returns
        -------
        shape is same as that of interaction_strength, dtype='byte', order='C'

        """
        return self._g_zero

    def get_zero_value_positions(self):
        """Return zero ph-ph interaction elements information."""
        warnings.warn(
            "Use attribute, Interaction.zero_value_positions "
            "instead of Interaction.get_zero_value_positions().",
            DeprecationWarning,
        )
        return self.zero_value_positions

    def get_phonons(self):
        """Return phonons on grid.

        Returns
        -------
        tuple
            frequencies : ndarray
                Phonon frequencies on grid.
                shape=(num_bz_grid, num_band), dtype='double', order='C'
            eigenvectors : ndarray
                Phonon eigenvectors on grid.
                shape=(num_bz_grid, num_band, num_band),
                dtype="c%d" % (np.dtype('double').itemsize * 2), order='C'
            phonon_done : ndarray
                1 if phonon at a grid point is calcualted, otherwise 0.
                shape=(num_bz_grid, ), dtype='byte'

        """
        return self._frequencies, self._eigenvectors, self._phonon_done

    @property
    def frequency_factor_to_THz(self):
        """Return phonon frequency conversion factor to THz."""
        return self._frequency_factor_to_THz

    def get_frequency_factor_to_THz(self):
        """Return phonon frequency conversion factor to THz."""
        warnings.warn(
            "Use attribute, Interaction.frequency_factor_to_THz ",
            "instead of Interaction.get_frequency_factor_to_THz().",
            DeprecationWarning,
        )
        return self.frequency_factor_to_THz

    @property
    def lapack_zheev_uplo(self):
        """Return U or L for lapack zheev solver."""
        return self._lapack_zheev_uplo

    def get_lapack_zheev_uplo(self):
        """Return U or L for lapack zheev solver."""
        warnings.warn(
            "Use attribute, Interaction.lapack_zheev_uplo "
            "instead of Interaction.get_lapack_zheev_uplo().",
            DeprecationWarning,
        )
        return self.lapack_zheev_uplo

    @property
    def cutoff_frequency(self):
        """Return cutoff phonon frequency to judge imaginary phonon."""
        return self._cutoff_frequency

    def get_cutoff_frequency(self):
        """Return cutoff phonon frequency to judge imaginary phonon."""
        warnings.warn(
            "Use attribute, Interaction.cutoff_frequency "
            "instead of Interaction.get_cutoff_frequency().",
            DeprecationWarning,
        )
        return self.cutoff_frequency

    def get_averaged_interaction(self):
        """Return sum over phonon triplets of interaction strength.

        See Eq.(21) of PRB 91, 094306 (2015)

        """
        # v[triplet, band0, band, band]
        v = self._interaction_strength
        w = self._weights_at_q
        v_sum = np.dot(w, v.sum(axis=2).sum(axis=2))
        return v_sum / np.prod(v.shape[2:])

    def get_primitive_and_supercell_correspondence(self):
        """Return atomic pair information."""
        return (self._svecs, self._multi, self._p2s, self._s2p, self._masses)

    def get_unit_conversion_factor(self):
        """Return unit conversion factor."""
        return self._unit_conversion

    def get_constant_averaged_interaction(self):
        """Return constant averaged interaction."""
        return self._constant_averaged_interaction

    def set_interaction_strength(self, pp_strength, g_zero=None):
        """Set interaction strength."""
        self._interaction_strength = pp_strength
        self._g_zero = g_zero

    def set_grid_point(self, grid_point, store_triplets_map=False):
        """Set grid point and prepare grid point triplets."""
        if not self._is_mesh_symmetry:
            (
                triplets_at_q,
                weights_at_q,
                triplets_map_at_q,
                ir_map_at_q,
            ) = get_nosym_triplets_at_q(grid_point, self._bz_grid)
        else:
            (
                triplets_at_q,
                weights_at_q,
                triplets_map_at_q,
                ir_map_at_q,
            ) = get_triplets_at_q(grid_point, self._bz_grid, swappable=True)

            # Special treatment of symmetry is applied when q_direction is
            # used.
            if self._nac_q_direction is not None:
                if (self._bz_grid.addresses[grid_point] == 0).all():
                    self._phonon_done[grid_point] = 0
                    self.run_phonon_solver(
                        np.array(
                            [
                                grid_point,
                            ],
                            dtype="int_",
                        )
                    )
                    rotations = []
                    for i, r in enumerate(self._bz_grid.reciprocal_operations):
                        dq = self._nac_q_direction
                        dq /= np.linalg.norm(dq)
                        diff = np.dot(r, dq) - dq
                        if (abs(diff) < 1e-5).all():
                            rotations.append(self._bz_grid.rotations[i])
                    (
                        triplets_at_q,
                        weights_at_q,
                        triplets_map_at_q,
                        ir_map_at_q,
                    ) = get_triplets_at_q(
                        grid_point,
                        self._bz_grid,
                        reciprocal_rotations=rotations,
                        is_time_reversal=False,
                    )

        reciprocal_lattice = np.linalg.inv(self._primitive.cell)
        for triplet in triplets_at_q:
            sum_q = (self._bz_grid.addresses[triplet]).sum(axis=0)
            if (sum_q % self.mesh_numbers != 0).any():
                print("============= Warning ==================")
                print("%s" % triplet)
                for tp in triplet:
                    print(
                        "%s %s"
                        % (
                            self._bz_grid.addresses[tp],
                            np.linalg.norm(
                                np.dot(
                                    reciprocal_lattice,
                                    self._bz_grid.addresses[tp]
                                    / self.mesh_numbers.astype("double"),
                                )
                            ),
                        )
                    )
                print("%s" % sum_q)
                print("============= Warning ==================")

        self._triplets_at_q = triplets_at_q
        self._weights_at_q = weights_at_q

        if store_triplets_map:
            self._triplets_map_at_q = triplets_map_at_q
            self._ir_map_at_q = ir_map_at_q

    def init_dynamical_matrix(
        self,
        fc2,
        supercell,
        primitive,
        nac_params=None,
        solve_dynamical_matrices=True,
        decimals=None,
    ):
        """Prepare for phonon calculation on grid.

        solve_dynamical_matrices : bool
           When False, phonon calculation will be postponed.

        """
        self._allocate_phonon()
        self._nac_params = nac_params
        self._dm = get_dynamical_matrix(
            fc2,
            supercell,
            primitive,
            nac_params=nac_params,
            frequency_scale_factor=self._frequency_scale_factor,
            decimals=decimals,
            symprec=self._symprec,
        )

        self._phonon_done[0] = 0
        if solve_dynamical_matrices:
            self.run_phonon_solver()
        else:
            self.run_phonon_solver(np.array([0], dtype="int_"))

        if (self._bz_grid.addresses[0] == 0).all():
            if np.sum(self._frequencies[0] < self._cutoff_frequency) < 3:
                for i, f in enumerate(self._frequencies[0, :3]):
                    if not (f < self._cutoff_frequency):
                        self._frequencies[0, i] = 0
                        print("=" * 26 + " Warning " + "=" * 26)
                        print(
                            " Phonon frequency of band index %d at Gamma "
                            "is calculated to be %f." % (i + 1, f)
                        )
                        print(" But this frequency is forced to be zero.")
                        print("=" * 61)

    def set_phonon_data(self, frequencies, eigenvectors, bz_grid_addresses):
        """Set phonons on grid."""
        if bz_grid_addresses.shape != self._bz_grid.addresses.shape:
            raise RuntimeError(
                "Input grid address size is inconsistent. " "Setting phonons faild."
            )

        if (self._bz_grid.addresses - bz_grid_addresses).all():
            raise RuntimeError(
                "Input grid addresses are inconsistent. " "Setting phonons faild."
            )
        else:
            self._phonon_done[:] = 1
            self._frequencies[:] = frequencies
            self._eigenvectors[:] = eigenvectors

    def run_phonon_solver(self, grid_points=None):
        """Run phonon solver at BZ-grid points."""
        if grid_points is None:
            self._get_phonons_at_all_bz_grid_points()
        else:
            _grid_points = grid_points
            self._run_phonon_solver_c(_grid_points)

    def _get_phonons_at_all_bz_grid_points(self, expand_phonons=False):
        """Get phonons at all BZ-grid points."""
        if expand_phonons:
            self._expand_phonons()
        else:
            grid_points = np.arange(len(self._bz_grid.addresses), dtype="int_")
            self._run_phonon_solver_c(grid_points)

    def _expand_phonons(self):
        """Phonons at ir-grid-points are copied by proper rotations.

        Some phonons that are not covered by rotations are solved.

        The following data are updated.
            self._frequencies
            self._eigenvectors
            self._phonon_done

        """
        ir_grid_points, _, _ = get_ir_grid_points(self._bz_grid)
        ir_bz_grid_points = self._bz_grid.grg2bzg[ir_grid_points]
        self._run_phonon_solver_c(ir_bz_grid_points)

        d2r_map = self._get_reciprocal_rotations_in_space_group_operations()

        # perms.shape = (len(spg_ops), len(primitive)), dtype='intc'
        perms = compute_all_sg_permutations(
            self._primitive.scaled_positions,
            self._bz_grid.symmetry_dataset["rotations"],
            self._bz_grid.symmetry_dataset["translations"],
            np.array(self._primitive.cell.T, dtype="double", order="C"),
            symprec=self._symprec,
        )

        for d_i, r_i in enumerate(d2r_map):
            r = self._bz_grid.rotations[r_i]
            r_cart = self._bz_grid.rotations_cartesian[r_i]
            for irgp in ir_bz_grid_points:
                bzgp = get_grid_points_by_rotations(
                    irgp,
                    self._bz_grid,
                    reciprocal_rotations=[
                        r,
                    ],
                    with_surface=True,
                )[0]
                if self._phonon_done[bzgp]:
                    continue

                self._rotate_eigvecs(irgp, bzgp, r_cart, perms[d_i], d_i)

        bz_grid_points_solved = self._get_phonons_at_minus_q()
        if bz_grid_points_solved:
            print("DEBUG: BZ-grid points additionally solved " "than ir-grid-points.")
            print(bz_grid_points_solved)

    def _get_reciprocal_rotations_in_space_group_operations(self):
        """Collect reciprocal rotations that belong to space group operations.

        Exclude reciprocal rotations that are made by time reversal symmetry.

        Returns
        -------
        d2r_map : list
            Indices of reciprocal rotations.

        """
        d2r_map = []
        for r in self._bz_grid.symmetry_dataset["rotations"]:
            for i, rec_r in enumerate(self._bz_grid.reciprocal_operations):
                if (rec_r.T == r).all():
                    d2r_map.append(i)
                    break

        assert len(d2r_map) == len(self._bz_grid.symmetry_dataset["rotations"])

        return d2r_map

    def _rotate_eigvecs(self, orig_gp, bzgp, r_cart, perm, t_i):
        r"""Rotate eigenvectors at q to those Rq.

        e_j'(Rq) = R e_j(q) exp(-iRq.\tau)

        """
        Rq = np.dot(self._bz_grid.QDinv, self._bz_grid.addresses[bzgp])
        tau = self._bz_grid.symmetry_dataset["translations"][t_i]
        phase_factor = np.exp(-2j * np.pi * np.dot(Rq, tau))
        self._phonon_done[bzgp] = 1
        self._frequencies[bzgp, :] = self._frequencies[orig_gp, :]
        eigvecs = self._eigenvectors[orig_gp, :, :] * phase_factor
        for i, vec in enumerate(eigvecs.T):
            vec_perm = vec.reshape(-1, 3)[perm, :].T
            vec_rot = np.dot(r_cart, vec_perm).T.ravel()
            self._eigenvectors[bzgp, :, i] = vec_rot

    def _get_phonons_at_minus_q(self):
        """Phonons at -q are given by phonons at q.

        A few points may be uncovered by rotations. Those points are counted.

        Returns
        -------
        bz_grid_points_solved : list of int
            BZ-grid points where phonons that were additionally solved
            in this method.

        """
        r_inv = -np.eye(3, dtype="int_")
        bz_grid_points_solved = []
        for bzgp, done in enumerate(self._phonon_done):
            if done:
                continue

            # Get grid point at -q.
            bzgp_mq = get_grid_points_by_rotations(
                bzgp,
                self._bz_grid,
                reciprocal_rotations=[
                    r_inv,
                ],
                with_surface=True,
            )[0]

            if self._phonon_done[bzgp_mq] == 0:
                self._run_phonon_solver_c(
                    np.array(
                        [
                            bzgp_mq,
                        ],
                        dtype="int_",
                    )
                )
                bz_grid_points_solved.append(bzgp_mq)

            self._phonon_done[bzgp] = 1
            self._frequencies[bzgp, :] = self._frequencies[bzgp_mq, :]
            self._eigenvectors[bzgp, :, :] = np.conj(self._eigenvectors[bzgp_mq, :, :])

        return bz_grid_points_solved

    def delete_interaction_strength(self):
        """Delete large arrays loosely.

        Memory deallocation would rely on garbage collector of python.
        So this may not work as expected.

        """
        self._interaction_strength = None
        self._g_zero = None

    def _set_fc3(self, fc3):
        if (
            type(fc3) == np.ndarray
            and fc3.dtype == np.dtype("double")
            and fc3.flags.aligned
            and fc3.flags.owndata
            and fc3.flags.c_contiguous
            and self._frequency_scale_factor is None
        ):  # noqa E129
            self._fc3 = fc3
        elif self._frequency_scale_factor is None:
            self._fc3 = np.array(fc3, dtype="double", order="C")
        else:
            self._fc3 = np.array(
                fc3 * self._frequency_scale_factor ** 2, dtype="double", order="C"
            )

    def _set_band_indices(self, band_indices):
        num_band = len(self._primitive) * 3
        if band_indices is None:
            self._band_indices = np.arange(num_band, dtype="int_")
        else:
            self._band_indices = np.array(band_indices, dtype="int_")

    def _run_c(self, g_zero):
        import phono3py._phono3py as phono3c

        if g_zero is None or self._symmetrize_fc3q:
            _g_zero = np.zeros(
                self._interaction_strength.shape, dtype="byte", order="C"
            )
        else:
            _g_zero = g_zero

        phono3c.interaction(
            self._interaction_strength,
            _g_zero,
            self._frequencies,
            self._eigenvectors,
            self._triplets_at_q,
            self._bz_grid.addresses,
            self._bz_grid.D_diag,
            self._bz_grid.Q,
            self._fc3,
            self._svecs,
            self._multi,
            self._masses,
            self._p2s,
            self._s2p,
            self._band_indices,
            self._symmetrize_fc3q,
            self._cutoff_frequency,
        )
        self._interaction_strength *= self._unit_conversion
        self._g_zero = g_zero

    def _run_phonon_solver_c(self, grid_points):
        run_phonon_solver_c(
            self._dm,
            self._frequencies,
            self._eigenvectors,
            self._phonon_done,
            grid_points,
            self._bz_grid.addresses,
            self._bz_grid.QDinv,
            frequency_conversion_factor=self._frequency_factor_to_THz,
            nac_q_direction=self._nac_q_direction,
            lapack_zheev_uplo=self._lapack_zheev_uplo,
        )

    def _run_py(self):
        r2r = RealToReciprocal(
            self._fc3, self._primitive, self.mesh_numbers, symprec=self._symprec
        )
        r2n = ReciprocalToNormal(
            self._primitive,
            self._frequencies,
            self._eigenvectors,
            self._band_indices,
            cutoff_frequency=self._cutoff_frequency,
        )

        for i, grid_triplet in enumerate(self._triplets_at_q):
            print("%d / %d" % (i + 1, len(self._triplets_at_q)))
            r2r.run(self._bz_grid.addresses[grid_triplet])
            fc3_reciprocal = r2r.get_fc3_reciprocal()
            for gp in grid_triplet:
                self._run_phonon_solver_py(gp)
            r2n.run(fc3_reciprocal, grid_triplet)
            self._interaction_strength[i] = (
                np.abs(r2n.get_reciprocal_to_normal()) ** 2 * self._unit_conversion
            )

    def _run_phonon_solver_py(self, grid_point):
        run_phonon_solver_py(
            grid_point,
            self._phonon_done,
            self._frequencies,
            self._eigenvectors,
            self._bz_grid.addresses,
            self._bz_grid.QDinv,
            self._dm,
            self._frequency_factor_to_THz,
            self._lapack_zheev_uplo,
        )

    def _allocate_phonon(self):
        num_band = len(self._primitive) * 3
        num_grid = len(self._bz_grid.addresses)
        self._phonon_done = np.zeros(num_grid, dtype="byte")
        self._frequencies = np.zeros((num_grid, num_band), dtype="double", order="C")
        itemsize = self._frequencies.itemsize
        self._eigenvectors = np.zeros(
            (num_grid, num_band, num_band), dtype=("c%d" % (itemsize * 2)), order="C"
        )


import time

import win32con
import ctypes
import pywintypes
import win32api
import win32gui
import win32ui
import win32com.client
import numpy as np
import cv2.cv2 as cv
import numpy as np
from config_loader import config

import logging

logger = logging.getLogger(__name__)

def rescale_point(hwnd, point):
    dpi = ctypes.windll.user32.GetDpiForWindow(hwnd)
    x = point[0] + config.getint("Device", "EdgeOffsetX")
    y = point[1] + config.getint("Device", "EdgeOffsetY")
    # logger.debug("Rescale: %s -> %s", point, (x, y))
    return int(np.round(x * 96 / dpi)), int(np.round(y * 96 / dpi))

def click_at(hwnd, x, y, hold=0):
    # 向后台窗口发送单击事件，(x, y)为相对于窗口左上角的位置
    x, y = rescale_point(hwnd, (x, y))
    pos = win32api.MAKELONG(int(x), int(y))
    win32gui.SendMessage(hwnd, win32con.WM_LBUTTONDOWN, win32con.MK_LBUTTON, pos)
    time.sleep(hold)
    win32gui.SendMessage(hwnd, win32con.WM_LBUTTONUP, 0, pos)

def rand_point(p, diff=0.2):
    if not isinstance(diff, (list, tuple)):
        diff = [diff, diff]

    res = []
    for i in range(2):
        if isinstance(diff[i], float) and 0 < diff[i] <= 1:
            delta = p[i] * diff[i]
        else:
            delta = diff[i]
        logger.debug("p[%d]=%.0f+-%.1f", i, p[i], delta)
        if delta == 0:
            res.append(delta)
        else:
            res.append(np.random.triangular(p[i] - delta, p[i], p[i] + delta))
    logger.debug("Randomrize %s +- %s to %s", p, diff, res)
    return res

def rand_drag(hwnd, start, end, step=100, start_delay=0):
    start = rescale_point(hwnd, start)
    end = rescale_point(hwnd, end)
    win32gui.SendMessage(
        hwnd, win32con.WM_LBUTTONDOWN, win32con.MK_LBUTTON, win32api.MAKELONG(*start)
    )
    time.sleep(start_delay)

    dx = abs(start[0] - end[0])
    dy = abs(start[1] - end[1])
    dist = np.sqrt(dx ** 2 + dy ** 2)
    delta = [max(5, int(0.2 * dx)), max(5, int(0.2 * dy))]
    # 最少需要2个坐标
    count = max(int(dist / step), 2)
    points = np.zeros((count, 2))
    points[:, 0] = np.linspace(start[0], end[0], count)
    points[:, 1] = np.linspace(start[1], end[1], count)
    points = points.astype("int")
    for point in points:
        time.sleep(0.05)
        target = win32api.MAKELONG(*[int(n) for n in rand_point(point, delta)])
        win32gui.SendMessage(hwnd, win32con.WM_MOUSEMOVE, win32con.MK_LBUTTON, target)
    time.sleep(0.1)
    win32gui.SendMessage(hwnd, win32con.WM_LBUTTONUP, 0, target)

def drag(hwnd, start, end, step=100, start_delay=0):
    start = rescale_point(hwnd, start)
    end = rescale_point(hwnd, end)
    win32gui.SendMessage(
        hwnd, win32con.WM_LBUTTONDOWN, win32con.MK_LBUTTON, win32api.MAKELONG(*start)
    )
    time.sleep(start_delay)

    count = int(np.linalg.norm(np.array(start) - np.array(end)) / step)
    # 最少需要2个坐标
    count = max(count, 2)
    points = np.zeros((count, 2))
    points[:, 0] = np.linspace(start[0], end[0], count)
    points[:, 1] = np.linspace(start[1], end[1], count)
    points = points.astype("int")
    for point in points:
        time.sleep(0.05)
        target = win32api.MAKELONG(*point)
        win32gui.SendMessage(hwnd, win32con.WM_MOUSEMOVE, win32con.MK_LBUTTON, target)
    time.sleep(0.1)
    win32gui.SendMessage(hwnd, win32con.WM_LBUTTONUP, 0, target)

def rand_click(hwnd, rect, hold=0):
    x_min, y_min, x_max, y_max = rect
    x = np.random.triangular(x_min, (x_min + x_max) / 2, x_max)
    y = np.random.triangular(y_min, (y_min + y_max) / 2, y_max)
    click_at(hwnd, x, y, hold)

def make_foreground(hwnd, retry=True):
    logger.warning("make foreground")
    try:
        win32gui.SetForegroundWindow(hwnd)
        win32gui.SetWindowPos(
            hwnd,
            win32con.HWND_TOP,
            0,
            0,
            0,
            0,
            win32con.SWP_NOMOVE | win32con.SWP_NOSIZE | win32con.SWP_SHOWWINDOW,
        )
    except pywintypes.error as err:
        logger.warning(err)
        if retry:
            shell = win32com.client.Dispatch("WScript.Shell")
            shell.SendKeys("%")
            make_foreground(hwnd, False)

def get_window_hwnd(title):
    return win32gui.FindWindow(None, title)

def detect_window_size(hwnd):
    """获取窗口的大小信息"""
    left, top, right, bottom = win32gui.GetWindowRect(hwnd)
    dpi = ctypes.windll.user32.GetDpiForWindow(hwnd)
    w = right - left
    h = bottom - top
    if dpi != 96:
        w = int(w * dpi / 96)
        h = int(h * dpi / 96)
    return w, h

def get_window_shot(hwnd):
    # 对后台应用程序截图，程序窗口可以被覆盖，但如果最小化后只能截取到标题栏、菜单栏等。

    # 使用自定义的窗口边缘和大小设置
    dx = config.getint("Device", "EdgeOffsetX")
    dy = config.getint("Device", "EdgeOffsetY")
    w = config.getint("Device", "MainWidth")
    h = config.getint("Device", "MainHeight")
    window_w, window_h = detect_window_size(hwnd)
    if dx + w > window_w or dy + h > window_h:
        raise ValueError("截图区域超出窗口! 请检查配置文件")
    # logger.debug("截图: %dx%d at %dx%d", w, h, dx, dy)

    # 返回句柄窗口的设备环境、覆盖整个窗口，包括非客户区，标题栏，菜单，边框
    hwndDC = win32gui.GetWindowDC(hwnd)
    # 创建设备描述表
    mfcDC = win32ui.CreateDCFromHandle(hwndDC)
    # 创建内存设备描述表
    saveDC = mfcDC.CreateCompatibleDC()
    # 创建位图对象
    saveBitMap = win32ui.CreateBitmap()
    saveBitMap.CreateCompatibleBitmap(mfcDC, w, h)
    saveDC.SelectObject(saveBitMap)
    # 截图至内存设备描述表
    saveDC.BitBlt((0, 0), (w, h), mfcDC, (dx, dy), win32con.SRCCOPY)
    # 获取位图信息
    bmpinfo = saveBitMap.GetInfo()
    bmpdata = saveBitMap.GetBitmapBits(True)
    # 生成图像
    image_data = np.frombuffer(bmpdata, "uint8")
    image_data = image_data.reshape((bmpinfo["bmHeight"], bmpinfo["bmWidth"], 4))
    image_data = cv.cvtColor(image_data, cv.COLOR_BGRA2BGR)
    # 内存释放
    win32gui.DeleteObject(saveBitMap.GetHandle())
    saveDC.DeleteDC()
    mfcDC.DeleteDC()
    win32gui.ReleaseDC(hwnd, hwndDC)

    return image_data

def heartbeat():
    info = win32api.GetLastInputInfo()
    tick = win32api.GetTickCount()
    if tick - info > 30000:
        win32api.keybd_event(win32con.VK_CAPITAL, 0, 0, 0)
        win32api.keybd_event(win32con.VK_CAPITAL, 0, win32con.KEYEVENTF_KEYUP, 0)
        time.sleep(0.1)
        win32api.keybd_event(win32con.VK_CAPITAL, 0, 0, 0)
        win32api.keybd_event(win32con.VK_CAPITAL, 0, win32con.KEYEVENTF_KEYUP, 0)

if __name__ == "__main__":
    nox_hwnd = get_window_hwnd("夜神模拟器")
    x0, y0, x1, y1 = (630, 370, 680, 420)
    click_at(nox_hwnd, (x0 + x1) / 2, (y0 + y1) / 2)


import os
import time
import cv2
import numpy as np
import models
import datasets as ds
from datasets.lane_det import collate_fn
from options.options import parser
import paddle.fluid as fluid
from paddle.fluid.io import DataLoader

best_mIoU = 0

def main():
    global args, best_mIoU
    args = parser.parse_args()

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    if args.dataset == 'LaneDet':
        num_class = 20
    else:
        raise ValueError('Unknown dataset ' + args.dataset)

    # get places
    places = fluid.cuda_places()

    with fluid.dygraph.guard():
        model = models.ERFNet(num_class, [576, 1024])
        input_mean = model.input_mean
        input_std = model.input_std

        if args.resume:
            print(("=> loading checkpoint '{}'".format(args.resume)))
            checkpoint, _ = fluid.load_dygraph(args.resume)
            model.load_dict(checkpoint)
            print("=> checkpoint loaded successfully")
        else:
            print(("=> loading checkpoint '{}'".format('trained/ERFNet_trained')))
            checkpoint, _ = fluid.load_dygraph('trained/ERFNet_trained')
            model.load_dict(checkpoint)
            print("=> default checkpoint loaded successfully")

        # Data loading code
        test_dataset = ds.LaneDataSet(
            dataset_path='datasets/PreliminaryData',
            data_list=args.val_list,
            transform=[
                lambda x: cv2.resize(x, (1024, 576)),
                lambda x: x - np.asarray(input_mean)[None, None, :] / np.array(input_std)[None, None, :],
            ]
        )

        test_loader = DataLoader(
            test_dataset,
            places=places[0],
            batch_size=1,
            shuffle=False,
            num_workers=args.workers,
            collate_fn=collate_fn
        )

        ### evaluate ###
        mIoU = validate(test_loader, model)
        # print('mIoU: {}'.format(mIoU))
    return

def validate(test_loader, model):

    batch_time = AverageMeter()
    # IoU = AverageMeter()
    mIoU = 0

    # switch to evaluate mode
    model.eval()

    end = time.time()
    for i, (input, img_name, img_size) in enumerate(test_loader):

        # compute output
        output = model(input).numpy()

        # measure accuracy and record loss
        pred = np.argmax(output, 1)

        img_name = img_name.numpy().tolist()
        img_size = img_size.numpy().tolist()

        for j, (name, size) in enumerate(zip(img_name, img_size)):
            # scale = size[1] / pred.shape[2]
            # img = cv2.resize(pred[j], None, fx=scale, fy=scale, interpolation=cv2.INTER_NEAREST)
            # img = np.pad(img, ((size[0] - img.shape[0], 0), (0, 0)))
            img = cv2.resize(pred[j], tuple(size[:2][::-1]), interpolation=cv2.INTER_NEAREST)
            img = img.astype(np.uint8)
            file_name = os.path.join(args.output_dir, str(name) + '.png')
            cv2.imwrite(file_name, img)

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if (i + 1) % args.print_freq == 0:
            print(('Test: [{0}/{1}]\t' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'.format(
                i + 1, len(test_loader), batch_time=batch_time)))

    print('finished, #test:{}'.format(i + 1))

    return mIoU

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = None
        self.avg = None
        self.sum = None
        self.count = None

    def update(self, val, n=1):
        if self.val is None:
            self.val = val
            self.sum = val * n
            self.count = n
            self.avg = self.sum / self.count
        else:
            self.val = val
            self.sum += val * n
            self.count += n
            self.avg = self.sum / self.count

if __name__ == '__main__':
    main()


# coding=utf-8
import tensorflow as tf
from GANNetwork.cycleGAN import CycleGAN
from common.common import getImg, torch_decay, getFiles, saveImg, imgPool, encode
import numpy as np
import random
import threading, os, time, cv2
from tensorflow.python import pywrap_tensorflow

slim = tf.contrib.slim

tf.app.flags.DEFINE_string('input',
                           "horse2zebra/trainA/n02381460_175.jpg",
                           'The directory that input pictures are saved')
tf.app.flags.DEFINE_string('output', "out/14.jpg", 'The directory that validate pictures are saved')
tf.app.flags.DEFINE_string('checkpoint', "model/model_zebra.ckpt",
                           'The directory that trained network will be saved')
tf.app.flags.DEFINE_string('status',
                           "X2Y",
                           'The directory that input pictures are saved')
tf.app.flags.DEFINE_string('Norm', 'BATCH', 'Choose to use Batchnorm or instanceNorm')
tf.app.flags.DEFINE_bool('USE_E', False, 'Choose to use Edge or not')
tf.app.flags.DEFINE_integer('img_size', 256, 'The size of input img')
tf.app.flags.DEFINE_integer('ngf', 64, 'The number of gen filters in first conv layer')
tf.app.flags.DEFINE_integer('batch_size', 1, 'The batch size of testing')

FLAGS = tf.app.flags.FLAGS

def test1(dir, status='X2Y'):
    capture = cv2.VideoCapture(dir)
    capture.set(cv2.CAP_PROP_FRAME_WIDTH, FLAGS.img_size)
    capture.set(cv2.CAP_PROP_FRAME_HEIGHT, FLAGS.img_size)
    config = tf.ConfigProto(allow_soft_placement=True)
    config.gpu_options.allow_growth = True
    with tf.Session(config=config) as sess:
        _, content_img = capture.read()
        content = np.expand_dims(encode(content_img), 0)
        network = CycleGAN(FLAGS.batch_size, FLAGS.ngf, FLAGS.img_size, is_training=False, Norm=FLAGS.Norm)
        network.test(content.shape)
        sess.run(tf.global_variables_initializer())
        var_list = [var for var in tf.global_variables() if 'generator_net' in var.name or
                    'discriminator_net' in var.name or 'edge_net' in var.name]
        saver = tf.train.Saver(var_list=var_list)
        saver.restore(sess, FLAGS.checkpoint)
        print("restored all")
        while True:
            if status == 'X2Y':
                out = sess.run(network.Ygenerated, feed_dict={network.testA: content})
            else:
                out = sess.run(network.Xgenerated, feed_dict={network.testB: content})
            # result = np.clip(out, 0, 255).astype(np.uint8)
            cv2.imshow("test", result)
            c = cv2.waitKey(1)
            if c == 27:
                break
            _, content_img = capture.read()
            content = np.expand_dims(encode(content_img), 0)
        capture.release()
        cv2.destroyAllWindows()

def test2(status='X2Y'):
    start = time.time()
    config = tf.ConfigProto(allow_soft_placement=True)
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        content_img = cv2.resize(getImg(FLAGS.input), (FLAGS.img_size, FLAGS.img_size))
        content = np.expand_dims(encode(content_img), 0)
        network = CycleGAN(FLAGS.batch_size, FLAGS.ngf, FLAGS.img_size, is_training=True, use_E=FLAGS.USE_E,
                           Norm=FLAGS.Norm)
        network.test(content.shape)
        sess.run(tf.global_variables_initializer())
        var_list = [var for var in tf.global_variables() if
                    'generator_net' in var.name or 'discriminator_net' in var.name or 'edge_net' in var.name]
        saver = tf.train.Saver(var_list=var_list)
        saver.restore(sess, FLAGS.checkpoint)
        print("restored all")
        s = time.time()
        if status == 'X2Y':
            out = sess.run(network.Ygenerated, feed_dict={network.testA: content})
        else:
            out = sess.run(network.Xgenerated, feed_dict={network.testB: content})
        # result = np.clip(out[0], 0, 255).astype(np.uint8)
        print("Transform in {} s".format((time.time() - s)))
        saveImg(out[0], FLAGS.output)
        print("Finished all process in {} s".format(time.time() - start))

def check():
    reader = pywrap_tensorflow.NewCheckpointReader(FLAGS.checkpoint)
    var_to_shape_map = reader.get_variable_to_shape_map()
    for key in var_to_shape_map:
        if 'G_Model' in key:
            print("tensor_name: ", key)

if __name__ == '__main__':
    test2(FLAGS.status)
    # check()