import argparse
import pickle
import sys

import numpy as np

from data_prep import batch_iter, createOneHotMosei2way, get_raw_data, createOneHot

seed = 1234

np.random.seed(seed)
import tensorflow as tf
tf.enable_eager_execution()
from tqdm import tqdm

from model import LSTM_Model

from sklearn.metrics import f1_score , roc_auc_score

tf.set_random_seed(seed)

unimodal_activations = {}

def multimodal(unimodal_activations, data, classes, attn_fusion=True, enable_attn_2=False, use_raw=True):
    if use_raw:
        if attn_fusion:
            attn_fusion = False

        train_data, test_data, audio_train, audio_test, text_train, text_test, video_train, video_test, train_label, test_label, seqlen_train, seqlen_test, train_mask, test_mask = get_raw_data(
            data, classes)

        a_dim = audio_train.shape[-1]
        v_dim = video_train.shape[-1]
        t_dim = text_train.shape[-1]

    else:
        print("starting multimodal")
        # Fusion (appending) of features

        train_mask = unimodal_activations['train_mask']
        test_mask = unimodal_activations['test_mask']
        print('train_mask', train_mask.shape)

        array_default = np.empty((train_mask.shape[0], train_mask.shape[1], 100))
        t_dim = a_dim = v_dim = None #initializing the dims

        if ('text_train' in unimodal_activations):
            text_train = unimodal_activations['text_train']
            text_test = unimodal_activations['text_test']
            t_dim = text_train.shape[-1]
        else:
            text_train = array_default
            text_test = array_default

        if ('audio_train' in unimodal_activations):
            audio_train = unimodal_activations['audio_train']
            audio_test = unimodal_activations['audio_test']
            a_dim = audio_train.shape[-1]
        else:
            audio_train = array_default
            audio_test = array_default

        if ('video_train' in unimodal_activations):
            video_train = unimodal_activations['video_train']
            video_test = unimodal_activations['video_test']
            v_dim = video_train.shape[-1]
        else:
            video_train = array_default
            video_test = array_default

        train_label = unimodal_activations['train_label']
        print('train_label', train_label.shape)
        test_label = unimodal_activations['test_label']
        print('test_label', test_label.shape)

        # print(train_mask_bool)
        seqlen_train = np.sum(train_mask, axis=-1)
        print('seqlen_train', seqlen_train.shape)
        seqlen_test = np.sum(test_mask, axis=-1)
        print('seqlen_test', seqlen_test.shape)

    if attn_fusion:
        print('With attention fusion')
    allow_soft_placement = True
    log_device_placement = False

    # Multimodal model
    session_conf = tf.ConfigProto(
        # device_count={'GPU': gpu_count},
        allow_soft_placement=allow_soft_placement,
        log_device_placement=log_device_placement,
        gpu_options=tf.GPUOptions(allow_growth=True))
    gpu_device = 0
    best_acc = 0
    best_loss_accuracy = 0
    best_loss = 10000000.0
    best_epoch = 0
    best_epoch_loss = 0
    with tf.device('/device:GPU:%d' % gpu_device):
        print('Using GPU - ', '/device:GPU:%d' % gpu_device)
        with tf.Graph().as_default():
            tf.set_random_seed(seed)
            sess = tf.Session(config=session_conf)
            with sess.as_default():
                model = LSTM_Model(text_train.shape[1:], 0.0001, a_dim=a_dim, v_dim=v_dim, t_dim=t_dim,
                                   emotions=classes, attn_fusion=attn_fusion,
                                   unimodal=False, enable_attn_2=enable_attn_2,
                                   seed=seed)
                sess.run(tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()))

                test_feed_dict = {
                    model.y: test_label,
                    model.seq_len: seqlen_test,
                    model.mask: test_mask,
                    model.lstm_dropout: 0.0,
                    model.lstm_inp_dropout: 0.0,
                    model.dropout: 0.0,
                    model.dropout_lstm_out: 0.0
                }

                if t_dim: test_feed_dict.update({model.t_input: text_test})
                if a_dim: test_feed_dict.update({model.a_input: audio_test})
                if v_dim: test_feed_dict.update({model.v_input: video_test})

                # print('\n\nDataset: %s' % (data))
                print("\nEvaluation before training:")
                # Evaluation after epoch
                step, loss, accuracy = sess.run(
                    [model.global_step, model.loss, model.accuracy],
                    #[model.global_step, model.loss, model.f1],
                    test_feed_dict)
                print("EVAL: epoch {}: step {}, loss {:g}, acc {:g}".format(0, step, loss, accuracy))

                for epoch in range(epochs):
                    epoch += 1

                    batches = batch_iter(list(
                        zip(text_train, audio_train, video_train, train_mask, seqlen_train, train_label)),
                        batch_size)

                    # Training loop. For each batch...
                    print('\nTraining epoch {}'.format(epoch))
                    l = []
                    a = []
                    for i, batch in tqdm(enumerate(batches)):
                        b_text_train, b_audio_train, b_video_train, b_train_mask, b_seqlen_train, b_train_label = zip(
                            *batch)
                        # print('batch_hist_v', len(batch_utt_v))
                        feed_dict = {}

                        if t_dim: feed_dict.update({model.t_input: b_text_train})
                        if a_dim: feed_dict.update({model.a_input: b_audio_train})
                        if v_dim: feed_dict.update({model.v_input: b_video_train})

                        feed_dict.update({
                            model.y: b_train_label,
                            model.seq_len: b_seqlen_train,
                            model.mask: b_train_mask,
                            model.lstm_dropout: 0.4,
                            model.lstm_inp_dropout: 0.0,
                            model.dropout: 0.2,
                            model.dropout_lstm_out: 0.2
                        })

                        _, step, loss, accuracy = sess.run(
                            [model.train_op, model.global_step, model.loss, model.accuracy],
                            feed_dict)
                        l.append(loss)
                        a.append(accuracy)

                    print("\t \tEpoch {}:, loss {:g}, accuracy {:g}".format(epoch, np.average(l), np.average(a)))
                    # Evaluation after epoch
                    step, loss, accuracy, preds, y, mask = sess.run(
                        [model.global_step, model.loss, model.accuracy, model.preds, model.y, model.mask],
                        test_feed_dict)
                    f1 = f1_score(np.ndarray.flatten(tf.argmax(y, -1, output_type=tf.int32).eval()),
                                  np.ndarray.flatten(tf.argmax(preds, -1, output_type=tf.int32).eval()),
                                  sample_weight=np.ndarray.flatten(tf.cast(mask, tf.int32).eval()), average="weighted")

                    print("EVAL: After epoch {}: step {}, loss {:g}, acc {:g}, f1 {:g}}".format(epoch, step,
                                                                                loss / test_label.shape[
                                                                                0],
                                                                                accuracy, f1))
                    if accuracy > best_acc:
                        best_epoch = epoch
                        best_acc = accuracy
                    if loss < best_loss:
                        best_loss = loss
                        best_loss_accuracy = accuracy
                        best_epoch_loss = epoch

                print(
                    "\n\nBest epoch: {}\nBest test accuracy: {}\nBest epoch loss: {}\nBest test accuracy when loss is least: {}".format(
                        best_epoch, best_acc, best_epoch_loss, best_loss_accuracy))

def unimodal(mode, data, classes):
    print(('starting unimodal ', mode))

    # with open('./mosei/text_glove_average.pickle', 'rb') as handle:
    if data == 'mosei' or data == 'mosi':
        with open('./dataset/{0}/raw/{1}_{2}way.pickle'.format(data, mode, classes), 'rb') as handle:
            u = pickle._Unpickler(handle)
            u.encoding = 'latin1'
            # (train_data, train_label, test_data, test_label, maxlen, train_length, test_length) = u.load()
            if data == 'mosei':
                (train_data, train_label, _, _, test_data, test_label, _, train_length, _, test_length, _, _,
                 _) = u.load()
                if classes == '2':
                    train_label, test_label = createOneHotMosei2way(train_label, test_label)
            elif data == 'mosi':
                (train_data, train_label, test_data, test_label, maxlen, train_length, test_length) = u.load()
                train_label = train_label.astype('int')
                test_label = test_label.astype('int')
                train_label, test_label = createOneHot(train_label, test_label)
            else:
                raise NotImplementedError('Unknown dataset...')

            train_label = train_label.astype('int')
            test_label = test_label.astype('int')

            train_mask = np.zeros((train_data.shape[0], train_data.shape[1]), dtype='float')
            for i in range(len(train_length)):
                train_mask[i, :train_length[i]] = 1.0

            test_mask = np.zeros((test_data.shape[0], test_data.shape[1]), dtype='float')
            for i in range(len(test_length)):
                test_mask[i, :test_length[i]] = 1.0
    elif data == 'iemocap':
        train_data, test_data, audio_train, audio_test, text_train, text_test, video_train, video_test, train_label, test_label, seqlen_train, seqlen_test, train_mask, test_mask = get_raw_data(
            data, classes)
        if mode == 'text':
            train_data = text_train
            test_data = text_test
        elif mode == 'audio':
            train_data = audio_train
            test_data = audio_test
        elif mode == 'video':
            train_data = video_train
            test_data = video_test
    elif data =='cognimuse':
        maxlen=368
        train_data= pickle.load(open("X_train_np.pkl", "rb")).numpy()
        train_label = pickle.load(open("Y_train_np.pkl", "rb")).numpy()
        test_data = pickle.load(open("X_test_np.pkl", "rb")).numpy()
        test_label = pickle.load(open("Y_test_np.pkl", "rb")).numpy()
        train_length= pickle.load(open("train_length.pkl", "rb"))
        test_length= pickle.load(open("test_length.pkl", "rb"))
        #train_label = train_label.astype('int')
        #test_label = test_label.astype('int')
        train_mask = np.zeros((train_data.shape[0], train_data.shape[1]), dtype='float')
        for i in range(len(train_length)):
            train_mask[i, :train_length[i]] = 1.0

        test_mask = np.zeros((test_data.shape[0], test_data.shape[1]), dtype='float')
        for i in range(len(test_length)):
            test_mask[i, :test_length[i]] = 1.0

    # train_label, test_label = createOneHotMosei3way(train_label, test_label)

    attn_fusion = False

    print('train_mask', train_mask.shape)

    # print(train_mask_bool)
    seqlen_train = np.sum(train_mask, axis=-1)
    print('seqlen_train', seqlen_train.shape)
    seqlen_test = np.sum(test_mask, axis=-1)
    print('seqlen_test', seqlen_test.shape)

    allow_soft_placement = True
    log_device_placement = False

    # Multimodal model
    session_conf = tf.ConfigProto(
        # device_count={'GPU': gpu_count},
        allow_soft_placement=allow_soft_placement,
        log_device_placement=log_device_placement,
        gpu_options=tf.GPUOptions(allow_growth=True))
    gpu_device = 0
    best_acc = 0
    best_f1=0
    best_epoch = 0
    best_loss = 1000000.0
    best_epoch_loss = 0
    is_unimodal = True
    with tf.device('/device:GPU:%d' % gpu_device):
        print('Using GPU - ', '/device:GPU:%d' % gpu_device)
        with tf.Graph().as_default():
            tf.set_random_seed(seed)
            sess = tf.Session(config=session_conf)
            with sess.as_default():
                model = LSTM_Model(train_data.shape[1:], 0.003, a_dim=0, v_dim=0, t_dim=0, emotions=classes,
                                   attn_fusion=attn_fusion, unimodal=is_unimodal, seed=seed)
                sess.run(tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()))

                test_feed_dict = {
                    model.input: test_data,
                    model.y: test_label,
                    model.seq_len: seqlen_test,
                    model.mask: test_mask,
                    model.lstm_dropout: 0.0,
                    model.lstm_inp_dropout: 0.0,
                    model.dropout: 0.0,
                    model.dropout_lstm_out: 0.0

                }
                train_feed_dict = {
                    model.input: train_data,
                    model.y: train_label,
                    model.seq_len: seqlen_train,
                    model.mask: train_mask,
                    model.lstm_dropout: 0.4,
                    model.lstm_inp_dropout: 0.0,
                    model.dropout: 0.2,
                    model.dropout_lstm_out: 0.2}

                # print('\n\nDataset: %s' % (data))
                print("\nEvaluation before training:")
                # Evaluation after epoch

                step, loss, accuracy = sess.run(
                    [model.global_step, model.loss, model.accuracy],
                    #model.global_step, model.loss, model.f1],
                    test_feed_dict)
                print(accuracy)
                #accuracy=accuracy[0]
                #print("EVAL: epoch {}: step {}, loss {:g}, acc {:g}".format(0, step, loss, accuracy))

                for epoch in range(epochs):
                    epoch += 1
                    print(len(train_data))
                    batches = batch_iter(list(
                        zip(train_data, train_mask, seqlen_train, train_label)),
                        batch_size)

                    # Training loop. For each batch...
                    print('\nTraining epoch {}'.format(epoch))
                    l = []
                    a = []
                    for i, batch in tqdm(enumerate(batches)):
                        b_train_data, b_train_mask, b_seqlen_train, b_train_label = zip(
                            *batch)
                        feed_dict = {
                            model.input: b_train_data,
                            model.y: b_train_label,
                            model.seq_len: b_seqlen_train,
                            model.mask: b_train_mask,
                            model.lstm_dropout: 0.4,
                            model.lstm_inp_dropout: 0.0,
                            model.dropout: 0.2,
                            model.dropout_lstm_out: 0.2

                        }

                        _, step, loss, accuracy = sess.run(
                            [model.train_op, model.global_step, model.loss, model.accuracy],
                            #[model.train_op, model.global_step, model.loss, model.f1],
                            feed_dict)
                        #accuracy=accuracy[0]
                        l.append(loss)
                        a.append(accuracy)

                    print("\t \tEpoch {}:, loss {:g}, accuracy {:g}".format(epoch, np.average(l), np.average(a)))
                    # Evaluation after epoch
                    #step, loss, accuracy, preds, y, mask = sess.run(
                        #[model.global_step, model.loss, model.accuracy, model.preds, model.y, model.mask],
                        #test_feed_dict)
                    step, loss, accuracy, test_activations,preds,y,mask = sess.run(
                        [model.global_step, model.loss, model.accuracy, model.inter1,model.preds,model.y,model.mask],
                        #[model.global_step, model.loss, model.f1, model.inter1],
                        test_feed_dict)
                    #accuracy=accuracy[0]
                    loss = loss / test_label.shape[0]
                    f1 = f1_score(np.ndarray.flatten(tf.argmax(y, -1, output_type=tf.int32).eval()),
                                  np.ndarray.flatten(tf.argmax(preds, -1, output_type=tf.int32).eval()),
                                  sample_weight=np.ndarray.flatten(tf.cast(mask, tf.int32).eval()), average="weighted")
                    auc=roc_auc_score(np.ndarray.flatten(tf.argmax(y, -1, output_type=tf.int32).eval()),
                                  np.ndarray.flatten(tf.argmax(preds, -1, output_type=tf.int32).eval()),
                                  sample_weight=np.ndarray.flatten(tf.cast(mask, tf.int32).eval()), average="weighted")
                    print("EVAL: After epoch {}: step {}, loss {:g}, acc {:g},f1 {:g},auc {:g}".format(epoch, step, loss,
                                                                                accuracy, f1,auc))

                    if accuracy > best_acc:
                        #best_epoch = epoch
                        best_acc = accuracy

                    if f1 > best_f1:
                        best_epoch = epoch
                        best_f1 = f1

                    if epoch == 30:
                        step, loss, accuracy, train_activations = sess.run(
                            [model.global_step, model.loss, model.accuracy, model.inter1],
                            #[model.global_step, model.loss, model.f1, model.inter1],
                            train_feed_dict)
                        #accuracy=accuracy[0]
                        unimodal_activations[mode + '_train'] = train_activations
                        unimodal_activations[mode + '_test'] = test_activations

                        unimodal_activations['train_mask'] = train_mask
                        unimodal_activations['test_mask'] = test_mask
                        unimodal_activations['train_label'] = train_label
                        unimodal_activations['test_label'] = test_label

                    if loss < best_loss:
                        best_epoch_loss = epoch
                        best_loss = loss
                        # step, loss, accuracy, train_activations = sess.run(
                        # [model.global_step, model.loss, model.accuracy, model.inter1],
                        # train_feed_dict)
                        # unimodal_activations[mode + '_train'] = train_activations
                        # unimodal_activations[mode + '_test'] = test_activations

                        # unimodal_activations['train_mask'] = train_mask
                        # unimodal_activations['test_mask'] = test_mask
                        # unimodal_activations['train_label'] = train_label
                        # unimodal_activations['test_label'] = test_label

                print("\n\nBest epoch: {}\nBest test accuracy: {}".format(best_epoch, best_acc))
                print("\n\nBest epoch: {}\nBest test loss: {}".format(best_epoch_loss, best_loss))
                print("\n\nBest epoch: {}\nBest test f1: {}".format(best_epoch, best_f1))

def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

if __name__ == "__main__":
    argv = sys.argv[1:]
    parser = argparse.ArgumentParser()
    parser.add_argument("--unimodal", type=str2bool, nargs='?', const=True, default=True)
    parser.add_argument("--mode", type=str,nargs='*',default=['text', 'audio', 'video'])
    parser.add_argument("--fusion", type=str2bool, nargs='?', const=True, default=False)
    parser.add_argument("--attention_2", type=str2bool, nargs='?', const=True, default=False)
    parser.add_argument("--use_raw", type=str2bool, nargs='?', const=True, default=False)
    parser.add_argument("--data", type=str, default='mosi')
    parser.add_argument("--classes", type=str, default='2')
    args, _ = parser.parse_known_args(argv)

    print(args)

    batch_size = 50
    epochs = 100
    emotions = args.classes
    assert args.data in ['mosi', 'mosei', 'iemocap','cognimuse']

    if args.unimodal:
        print("Training unimodals first")
        # modality = ['text', 'audio', 'video']
        modality = args.mode
        for mode in modality:
            unimodal(mode, args.data, args.classes)

        print("Saving unimodal activations")
        with open('unimodal_{0}_{1}way.pickle'.format(args.data, args.classes), 'wb') as handle:
            pickle.dump(unimodal_activations, handle, protocol=pickle.HIGHEST_PROTOCOL)

    if not args.use_raw:
        with open('unimodal_{0}_{1}way.pickle'.format(args.data, args.classes), 'rb') as handle:
            u = pickle._Unpickler(handle)
            u.encoding = 'latin1'
            unimodal_activations = u.load()

    #epochs = 50
    #multimodal(unimodal_activations, args.data, args.classes, args.fusion, args.attention_2, use_raw=args.use_raw)


"""
resp_optimizer.py
Implementation of open-source vrsion of RESP method

"""
import os, sys
import numpy as np
from shutil import copyfile
from respyte.readinp_esp import *
from respyte.readmol import * # read molecules and assign (what?)
from respyte.molecule import * # copied from forcebalance package and modified
from respyte.gen_grid_pts import * # take either rdkit ol or oechem mol to generate grid points around molecules (modified CIB's code)
from respyte.engine import *

def main():
    # cwd = current working directory in which input folder exists
    cwd = os.getcwd()
    # check if the current working idr contains input folder
    print('\n')
    print('---------------------------------------------------------')
    print(' 1. Reading  input files and folders.                    ')
    print('---------------------------------------------------------')
    if os.path.isdir("%s/input" % cwd):
        print(' Found the input folder. Now read input/input.yml')
    else:
        print(' Failed to find input folder. Should have input(folder) containing input.yml and molecules(folder)')
    # read input.yml and generate gridOption.
    inp = Input('%s/input/input.yml' % (cwd))
    gridOptions = inp.gridOptions

    # Visit all subfolder molecule/ moli/confj and generate grid.dat
    for idx, i in enumerate(inp.nmols):
        molN = 'mol%d' % (idx+1)
        wkd = '%s/input/molecules/%s' % (cwd, molN)
        for j in range(i):
            confN = 'conf%d' % (j+1)
            path = wkd + '/%s' % (confN)
            # gen molecule object from pdb file
            pdbfile = path + '/%s_%s.pdb' % (molN, confN)
            mol2file = path + '/%s_%s.mol2' % (molN, confN)
            xyzfile = path + '/%s_%s.xyz' % (molN, confN)
            if os.path.isfile(pdbfile):
                coordfile = pdbfile
                #mol = Molecule(pdbfile)
            elif os.path.isfile(mol2file):
                coordfile = mol2file
                #mol = Molecule(mol2file)
            elif os.path.isfile(xyzfile):
                print(" Warning: Are you using xyz files? xyz format is not recommendable.")
                print("          But no worries! It will convert your xyz files to pdb with single residue named 'MOL'.")
                xyztopdb(xyzfile)
                coordfile = pdbfile
            else:
                raise RuntimeError(" Coordinate file does not exist! ")
            # Assign radii
            radii = inp.settings['radii']
            if inp.cheminformatics == 'openeye':
                mol = ReadOEMolFromFile(coordfile)
                mol = assignRadii(mol, 'OEMol', radii)
                moltype = 'OEMol'
            elif inp.cheminformatics == 'rdkit':
                mol = ReadRdMolFromFile(coordfile)
                mol = assignRadii(mol, 'RDMol', radii)
                moltype = 'RDMol'
            else:
                mol = ReadMolFromFile(coordfile)
                mol = assignRadii(mol, 'FBMol', radii)
                moltype = 'FBMol'
            # Generate points
            print('\n')
            print('---------------------------------------------------------')
            print(' 2. Generating grid points around each molecules.        ')
            print('---------------------------------------------------------')

            if moltype is 'FBMol':
                raise NotImplementedError('Grid gen using fbmol not implemented yet! Sorry!')
            if gridOptions['gridType'] == 'MSK' or gridOptions['gridType']=='extendedMsk':
                # generate Mers-Singh-Kollman Connolly surfaces at 1.4, 1.6, 1.8 and 2.0 * vdW radii
                allPts = GenerateMSKShellPts(mol, gridOptions, moltype)
            else:
                # Generate a face-centered cubic grid shell around the molecule using gridOptions
                allPts = [GenerateGridPointSetAroundOEMol(mol, gridOptions, moltype)]
            print('Total points:', np.sum([len(i) for i in allPts]))

            # Write grid.dat
            print('\n')
            print('---------------------------------------------------------')
            print(' Psi calculations are done. Will create espf files.      ')
            print('---------------------------------------------------------')
            tmppath = '%s/tmp' % path
            if inp.settings['forcegen'] == 'N' and os.path.isfile('%s/grid.dat' % path):
                if os.path.isdir(tmppath):
                    print('Rename pre-existing tmp dir into tmp.0 and re-make tmp dir.')
                    os.rename(tmppath, '%s/tmp.0' % path)
                    os.mkdir(tmppath)
                    copyfile('%s/grid.dat' % path, '%s/grid.dat' % tmppath)
                else:
                    os.mkdir(tmppath)
                    # copy grid.dat into tmp dir
                    copyfile('%s/grid.dat' % path, '%s/grid.dat' % tmppath)
                print(' Pre-existing grid.dat has been copied into tmp for Psi4 calculation.')
            else:
                if os.path.isdir(tmppath):
                    print(' Rename pre-existing tmp dir into tmp.0 and re-make tmp dir.')
                    os.rename(tmppath, '%s/tmp.0' % path)
                os.mkdir(tmppath)
                ofs = open('%s/grid.dat' % tmppath,'w')

                for pts in allPts: #changed (for msk-based scheme, allPts is a list of lists)
                    for pt in pts:
                        ofs.write( '{0:10.6f} {1:10.6f} {2:10.6f}\n'.format(pt[0], pt[1], pt[2]) )
                ofs.close()
                print(' grid.dat is successfully generated in %s.' % tmppath)
            # Run engine to run Psi4 calculation inside tmp directory.
            engine = EnginePsi4()
            if 'method' in inp.settings:
                mthd = inp.settings['method']
            else:
                mthd = 'hf'
            if 'basis' in inp.settings:
                bss = inp.settings['basis']
            else:
                bss = '6-31g*'

            chg = inp.charges[idx]

            ### Need to add 'solvation' setting!
            if 'pcm' in inp.settings:
                if inp.settings['pcm'] is 'N':
                    solv = None
                elif inp.settings['pcm'] is 'Y':
                    solv = inp.settings['solvent']
            else:
                solv = None

            engine.write_input(coordfile, method = mthd, basis = bss, charge = chg, solvent=solv, job_path = tmppath)
            engine.espcal(job_path= tmppath)
            griddat = tmppath + '/grid.dat'
            espdat = tmppath + '/grid_esp.dat'
            efdat = tmppath + '/grid_field.dat'
            espfoutput = tmppath + '/%s_%s.espf' % (molN, confN)
            engine.genespf(griddat, espdat, efdat, espfoutput)
            # copy espf file generated in moli/confj/tmp to moli/confj/
            copyfile(espfoutput, path + '/%s_%s.espf' % (molN, confN))
            copyfile(griddat, path + '/grid.dat')
    print('\n')
    print('---------------------------------------------------------')
    print('       Done! Hope it helps, Will see you again:)         ')
    print('---------------------------------------------------------')

if __name__ == "__main__":
    main()


#!/usr/bin/env python3
import argparse
import glob
import sys
import numpy as np
import matplotlib.pyplot as plt
import pysam
import math
import pandas
import haplotagging_stats
import os
import seaborn as sns
import collections

FILE_NAME = "file_name"
BLOCK_N50 = "block_n50"
ALL_SWITCH_RATE = "all_switch_rate"
ALL_HAMMING_RATE = "blockwise_hamming_rate"
UNCLASSIFIED_RATIO = 'unclassified_ratio'
CHROM="#chrom"
POS="position"

CLASSIFICATION_KEY = "classifications"

def parse_args(args = None):
    parser = argparse.ArgumentParser("Plots information from haplotagging_stats tsv")
    parser.add_argument('--input', '-i', dest='input_stats_tsvs', default=None, required=True, type=str, action='append',
                        help='TSV file generated from "haplotagging_stats.py".  Can set multiple times, all values will be plotted')
    parser.add_argument('--figure_name', '-f', dest='figure_name', default=None, required=False, type=str,
                       help='Figure name (will save if set)')
    parser.add_argument('--figure_name_input', '-F', dest='figure_name_input', default=False, required=False, action='store_true',
                       help='Figure name should be based off of input file (-f overrides this)')
    parser.add_argument('--plot_classified_not_accuracy', dest='plot_classified_not_accuracy', default=False, required=False, action='store_true',
                       help='Will plot the classified ratio instead of accuracy')
    parser.add_argument('--classification', '-c', dest='classification', default=None, required=True, type=str, action='append',
                        help='Specification of the form "Classification Name:/path/to/classification.bed".  If set, will not plot based on depth profiles')

    return parser.parse_args() if args is None else parser.parse_args(args)

def get_color(filename):
    if "margin" in filename.lower():
        return "red"
    if "whatshap" in filename.lower():
        return "blue"
    return "grey"

def get_label(filename):
    if "margin" in filename.lower():
        return "Margin"
    if "whatshap" in filename.lower():
        return "Whatshap"
    return "Unknown"

def get_classifications(classification_list):
    classification_position_map=collections.defaultdict(lambda : collections.defaultdict(lambda : list()))
    classification_position_map[CLASSIFICATION_KEY] = list()

    for classification_specification in classification_list:
        parts = classification_specification.split(":")
        assert(len(parts) == 2)
        classification_name = parts[0]
        classification_bed = parts[1]
        classification_position_map[CLASSIFICATION_KEY].append(classification_name)
        with open(classification_bed) as cin:
            for line in cin:
                if line.startswith("#"): continue
                bed_parts = line.split("\t")
                assert(len(bed_parts) >= 3)
                chr = bed_parts[0]
                start = int(bed_parts[1])
                end = int(bed_parts[2])
                if (end - start) < 1000: continue
                for i in range(start//1000, end//1000 + 1):
                    classification_position_map[chr][i*1000].append(classification_name)

    return classification_position_map

def plot_hist(dataframes, args):

    fig, ((ax1)) = plt.subplots(nrows=1,ncols=1)

    for i, filename in enumerate(dataframes.keys()):
        dataframe = dataframes[filename]
        accuracies = list()
        for x in dataframe.iterrows():
            if x[1][haplotagging_stats.TOTAL_DEPTH] < 10: continue
            if x[1][haplotagging_stats.CORRECT_RATIO] <= 0: continue
            # accuracies.append(min(.99, x[1][haplotagging_stats.CORRECT_RATIO]))
            accuracies.append(x[1][haplotagging_stats.CORRECT_RATIO])

        n, bins, patches = ax1.hist(accuracies, 17, density=True, histtype='stepfilled', cumulative=False, label=get_label(filename), alpha=.5,
                 color=get_color(filename))

    # ax1.set_yscale('log')
    # ax2.set_yscale('log')
    plt.title("Haplotagging Accuracy")
    ax1.set_xlabel("Haplotagging Accuracy (1kb granularity)")
    ax1.set_ylabel("Relative Occurrence")
    plt.legend()

    fig_name = args.figure_name
    if fig_name is None and args.figure_name_input:
        fig_name = ".".join(list(map(os.path.basename, args.input_stats_tsvs)))
        fig_name += ".png"
    if fig_name is not None:
        plt.savefig(fig_name)
    plt.show()

def plot_violin_with_depths(dataframes, args):

    merged_dataframe = []
    columns = ["Depth", "Accuracy", "Tool", "Classified Ratio"]
    total_depths = collections.defaultdict(lambda : 0)

    for i, filename in enumerate(dataframes.keys()):
        dataframe = dataframes[filename]
        for x in dataframe.iterrows():
            if x[1][haplotagging_stats.TOTAL_DEPTH] == 0: continue

            # get depths
            td = x[1][haplotagging_stats.TOTAL_DEPTH]
            if td < 20:
                depth = "<20"
            elif 20 <= td < 40:
                depth = "20-40"
            elif 40 <= td < 60:
                depth = "40-60"
            elif 60 <= td < 80:
                depth = "60-80"
            elif 80 <= td < 100:
                depth = "80-100"
            elif td >= 100:
                depth = ">100"
            else:
                assert(False)

            # quit for non-classified ratio
            cr = x[1][haplotagging_stats.CORRECT_RATIO]
            if not args.plot_classified_not_accuracy and cr <= 0:
                continue

            # save rows
            row = [depth, cr, get_label(filename), 1.0 - x[1][UNCLASSIFIED_RATIO]]
            merged_dataframe.append(row)
            row = ["All", cr, get_label(filename), 1.0 - x[1][UNCLASSIFIED_RATIO]]
            merged_dataframe.append(row)

            # track total depths
            total_depths[depth] += 1

    total_rows = sum(total_depths.values())

    col_order = ["<20", "20-40", "40-60", "60-80", "80-100", ">100", "All"]
    df = pandas.DataFrame(merged_dataframe, columns=columns)
    ax = sns.violinplot(x="Depth", y="Classified Ratio" if args.plot_classified_not_accuracy else "Accuracy",
                        hue="Tool", data=df, split=True, scale="count", inner="quartile", order=col_order, cut=0)

    ax.set_xticklabels(list(map(lambda x: "{}\n{:.1f}%".format(x, 100 if x == "All" else 100.0 * total_depths[x] / total_rows), col_order)))

    plt.legend()
    plt.tight_layout()

    fig_name = args.figure_name
    if fig_name is None and args.figure_name_input:
        fig_name = ".".join(list(map(os.path.basename, args.input_stats_tsvs)))
        fig_name += ".png"
    if fig_name is not None:
        plt.savefig(fig_name)
    plt.show()

def plot_violin_with_classifications(dataframes, args, classifications):

    merged_dataframe = []
    columns = ["Stratification", "Accuracy", "Tool", "Classified Ratio"]
    total_depths = collections.defaultdict(lambda : 0)

    plt.figure(figsize=[12,8])
    for i, filename in enumerate(dataframes.keys()):
        dataframe = dataframes[filename]
        for x in dataframe.iterrows():
            if x[1][haplotagging_stats.TOTAL_DEPTH] == 0: continue

            # get position
            chr = x[1][CHROM]
            pos = x[1][POS]

            # quit for non-classified ratio
            cr = x[1][haplotagging_stats.CORRECT_RATIO]
            if not args.plot_classified_not_accuracy and cr <= 0:
                continue

            # save rows
            row = ["All", cr, get_label(filename), 1.0 - x[1][UNCLASSIFIED_RATIO]]
            merged_dataframe.append(row)

            for classification in classifications[chr][pos]:
                row = [classification, cr, get_label(filename), 1.0 - x[1][UNCLASSIFIED_RATIO]]
                merged_dataframe.append(row)

    classifications[CLASSIFICATION_KEY].append("All")
    df = pandas.DataFrame(merged_dataframe, columns=columns)
    ax = sns.violinplot(x="Stratification", y="Classified Ratio" if args.plot_classified_not_accuracy else "Accuracy",
                        hue="Tool", data=df, split=True, scale="count",
                        order=classifications[CLASSIFICATION_KEY], cut=0)
    # ax = sns.violinplot(x="Stratification", y="Classified Ratio" if args.plot_classified_not_accuracy else "Accuracy",
    #                     hue="Tool", data=df, split=True, scale="count", inner="quartile",
    #                     order=classifications[CLASSIFICATION_KEY], cut=0)
    ax.set_ylim(.79, 1.01)

    plt.legend()
    plt.tight_layout()

    fig_name = args.figure_name
    if fig_name is None and args.figure_name_input:
        fig_name = ".".join(list(map(os.path.basename, args.input_stats_tsvs)))
        fig_name += ".png"
    if fig_name is not None:
        plt.savefig(fig_name)
    plt.show()

def main():
    args = parse_args()
    dataframes = dict()
    for tsv in args.input_stats_tsvs:
        dataframes[os.path.basename(tsv)] = pandas.read_csv(tsv, sep='\t')

    # plot_hist(dataframes, args)
    if args.classification is None or len(args.classification) == 0:
        plot_violin_with_depths(dataframes, args)
    else:
        classification_data = get_classifications(args.classification)
        plot_violin_with_classifications(dataframes, args, classification_data)

if __name__ == "__main__":
    main()


# coding: utf-8
# pylint: disable=invalid-name, protected-access, too-many-arguments, too-many-lines
# pylint: disable=import-error, no-name-in-module
"""Symbolic configuration API of mxnet."""
from __future__ import absolute_import as _abs

import ctypes
import warnings
from numbers import Number

import os as _os
import sys as _sys
import numpy as _numpy

from .base import _LIB, numeric_types
from .base import c_array, c_str, mx_uint, py_str, string_types, mx_real_t
from .base import NDArrayHandle, ExecutorHandle, SymbolHandle
from .base import check_call, MXNetError
from .context import Context, cpu
from .ndarray import NDArray, zeros as _nd_zeros, _DTYPE_NP_TO_MX, _DTYPE_MX_TO_NP
from .executor import Executor
from . import _symbol_internal as _internal
from .attribute import AttrScope

# Use different version of SymbolBase
# When possible, use cython to speedup part of computation.
try:
    if int(_os.environ.get("MXNET_ENABLE_CYTHON", True)) == 0:
        from ._ctypes.symbol import SymbolBase, _init_symbol_module
    elif _sys.version_info >= (3, 0):
        from ._cy3.symbol import SymbolBase, _init_symbol_module
    else:
        from ._cy2.symbol import SymbolBase, _init_symbol_module
except ImportError:
    if int(_os.environ.get("MXNET_ENFORCE_CYTHON", False)) != 0:
        raise ImportError("Cython Module cannot be loaded but MXNET_ENFORCE_CYTHON=1")
    from ._ctypes.symbol import SymbolBase, _init_symbol_module

class Symbol(SymbolBase):
    """Symbol is symbolic graph of the mxnet."""
    # disable dictionary storage, also do not have parent type.
    # pylint: disable=no-member
    __slots__ = []

    def __repr__(self):
        """Get a string representation of the symbol."""
        name = self.name
        return '<%s %s>' % (self.__class__.__name__,
                            'Grouped' if name is None else name)

    def __iter__(self):
        """Return all outputs in a list"""
        return (self[i] for i in self.list_outputs())

    def __add__(self, other):
        """x.__add__(y) <=> x+y """
        if isinstance(other, Symbol):
            return _internal._Plus(self, other)
        if isinstance(other, Number):
            return _internal._PlusScalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        """x.__sub__(y) <=> x-y """
        if isinstance(other, Symbol):
            return _internal._Minus(self, other)
        if isinstance(other, Number):
            return _internal._MinusScalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __rsub__(self, other):
        """x.__rsub__(y) <=> y-x """
        if isinstance(other, Number):
            return _internal._RMinusScalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __mul__(self, other):
        """x.__mul__(y) <=> x*y """
        if isinstance(other, Symbol):
            return _internal._Mul(self, other)
        if isinstance(other, Number):
            return _internal._MulScalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __rmul__(self, other):
        return self.__mul__(other)

    def __div__(self, other):
        """x.__div__(y) <=> x/y """
        if isinstance(other, Symbol):
            return _internal._Div(self, other)
        if isinstance(other, Number):
            return _internal._DivScalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __rdiv__(self, other):
        """x.__rdiv__(y) <=> y/x """
        if isinstance(other, Number):
            return _internal._RDivScalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __truediv__(self, other):
        return self.__div__(other)

    def __rtruediv__(self, other):
        return self.__rdiv__(other)

    def __pow__(self, other):
        """x.__pow__(y) <=> x**y """
        if isinstance(other, Symbol):
            return _internal._Power(self, other)
        if isinstance(other, Number):
            return _internal._PowerScalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __neg__(self):
        """x.__neg__(y) <=> -x """
        return self.__mul__(-1.0)

    def __copy__(self):
        return self.__deepcopy__(None)

    def __deepcopy__(self, _):
        handle = SymbolHandle()
        check_call(_LIB.MXSymbolCopy(self.handle,
                                     ctypes.byref(handle)))
        return Symbol(handle)

    def __eq__(self, other):
        """x.__eq__(y) <=> x==y """
        if isinstance(other, Symbol):
            return _internal._equal(self, other)
        if isinstance(other, numeric_types):
            return _internal._equal_scalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __ne__(self, other):
        """x.__ne__(y) <=> x!=y """
        if isinstance(other, Symbol):
            return _internal._not_equal(self, other)
        if isinstance(other, numeric_types):
            return _internal._not_equal_scalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __gt__(self, other):
        """x.__gt__(y) <=> x>y """
        if isinstance(other, Symbol):
            return _internal._greater(self, other)
        if isinstance(other, numeric_types):
            return _internal._greater_scalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __ge__(self, other):
        """x.__ge__(y) <=> x>=y """
        if isinstance(other, Symbol):
            return _internal._greater_equal(self, other)
        if isinstance(other, numeric_types):
            return _internal._greater_equal_scalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __lt__(self, other):
        """x.__lt__(y) <=> x<y """
        if isinstance(other, Symbol):
            return _internal._lesser(self, other)
        if isinstance(other, numeric_types):
            return _internal._lesser_scalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __le__(self, other):
        """x.__le__(y) <=> x<=y """
        if isinstance(other, Symbol):
            return _internal._lesser_equal(self, other)
        if isinstance(other, numeric_types):
            return _internal._lesser_equal_scalar(self, scalar=other)
        else:
            raise TypeError('type %s not supported' % str(type(other)))

    def __getstate__(self):
        handle = self.handle
        if handle is not None:
            return {'handle': self.tojson()}
        else:
            return {'handle': None}

    def __setstate__(self, state):
        # pylint: disable=assigning-non-slot
        handle = state['handle']
        if handle is not None:
            json_str = handle
            handle = SymbolHandle()
            check_call(_LIB.MXSymbolCreateFromJSON(c_str(json_str), ctypes.byref(handle)))
            self.handle = handle
        else:
            self.handle = None

    def __call__(self, *args, **kwargs):
        """Compose symbol on inputs.

        x.__call__(y, z) <=> x(y,z)

        Parameters
        ----------
        args:
            provide positional arguments

        kwargs:
            provide keyword arguments
        Returns
        -------
        the resulting symbol
        """
        s = self.__copy__()
        s._compose(*args, **kwargs)
        return s

    def _compose(self, *args, **kwargs):
        """Compose symbol on inputs.

        This call mutates the current symbol.

        Parameters
        ----------
        args:
            provide positional arguments

        kwargs:
            provide keyword arguments
        Returns
        -------
        the resulting symbol
        """
        name = kwargs.pop('name', None)

        if name:
            name = c_str(name)
        if len(args) != 0 and len(kwargs) != 0:
            raise TypeError('compose only accept input Symbols \
                either as positional or keyword arguments, not both')

        for arg in args:
            if not isinstance(arg, Symbol):
                raise TypeError('Compose expect `Symbol` as arguments')
        for val in kwargs.values():
            if not isinstance(val, Symbol):
                raise TypeError('Compose expect `Symbol` as arguments')

        num_args = len(args) + len(kwargs)
        if len(kwargs) != 0:
            keys = c_array(ctypes.c_char_p, [c_str(key) for key in kwargs.keys()])
            args = c_array(SymbolHandle, [s.handle for s in kwargs.values()])
        else:
            keys = None
            args = c_array(SymbolHandle, [s.handle for s in args])
        check_call(_LIB.MXSymbolCompose(
            self.handle, name, num_args, keys, args))

    def __getitem__(self, index):
        """x.__getitem__(i) <=> x[i]

        Get an output of this symbol

        Parameters
        ----------
        index : int or str
            indexing key

        """
        if isinstance(index, string_types):
            idx = None
            for i, name in enumerate(self.list_outputs()):
                if name == index:
                    if idx is not None:
                        raise ValueError('There are multiple outputs with name \"%s\"' % index)
                    idx = i
            if idx is None:
                raise ValueError('Cannot find output that matches name \"%s\"' % index)
            index = idx
        if not isinstance(index, int):
            raise TypeError('Symbol only support integer index to fetch i-th output')
        if index >= (len(self.list_outputs())):
            # Important, python determines the end by this exception
            raise IndexError
        handle = SymbolHandle()
        check_call(_LIB.MXSymbolGetOutput(
            self.handle, mx_uint(index), ctypes.byref(handle)))
        return Symbol(handle=handle)

    @property
    def name(self):
        """Get name string from the symbol, this function only works for non-grouped symbol.

        Returns
        -------
        value : str
            The name of this symbol, returns None for grouped symbol.
        """
        ret = ctypes.c_char_p()
        success = ctypes.c_int()
        check_call(_LIB.MXSymbolGetName(
            self.handle, ctypes.byref(ret), ctypes.byref(success)))
        if success.value != 0:
            return py_str(ret.value)
        else:
            return None

    def attr(self, key):
        """Get attribute string from the symbol, this function only works for non-grouped symbol.

        Parameters
        ----------
        key : str
            The key to get attribute from.

        Returns
        -------
        value : str
            The attribute value of the key, returns None if attribute do not exist.
        """
        ret = ctypes.c_char_p()
        success = ctypes.c_int()
        check_call(_LIB.MXSymbolGetAttr(
            self.handle, c_str(key), ctypes.byref(ret), ctypes.byref(success)))
        if success.value != 0:
            return py_str(ret.value)
        else:
            return None

    def list_attr(self, recursive=False):
        """Get all attributes from the symbol.

        Returns
        -------
        ret : dict of str to str
            a dicitonary mapping attribute keys to values
        """
        if recursive:
            raise DeprecationWarning("Symbol.list_attr with recursive=True has been deprecated. "
                                     "Please use attr_dict instead.")
        size = mx_uint()
        pairs = ctypes.POINTER(ctypes.c_char_p)()
        f_handle = _LIB.MXSymbolListAttrShallow
        check_call(f_handle(self.handle, ctypes.byref(size), ctypes.byref(pairs)))
        return {py_str(pairs[i*2]): py_str(pairs[i*2+1]) for i in range(size.value)}

    def attr_dict(self):
        """Recursively get all attributes from the symbol and its childrens

        Returns
        -------
        ret : dict of str to dict
            Returns a dict whose keys are names of the symbol and its children.
            Values of the returned dict are dictionaries that map attribute keys to values
        """
        size = mx_uint()
        pairs = ctypes.POINTER(ctypes.c_char_p)()
        f_handle = _LIB.MXSymbolListAttr
        check_call(f_handle(self.handle, ctypes.byref(size), ctypes.byref(pairs)))
        ret = {}
        for i in range(size.value):
            name, key = py_str(pairs[i*2]).split('$')
            val = py_str(pairs[i*2+1])
            if name not in ret:
                ret[name] = {}
            ret[name][key] = val
        return ret

    def _set_attr(self, **kwargs):
        """Set the attribute of the symbol.

        Parameters
        ----------
        **kwargs
            The attributes to set
        """
        for key, value in kwargs.items():
            if not isinstance(value, string_types):
                raise ValueError("Set Attr only accepts string values")
            check_call(_LIB.MXSymbolSetAttr(
                self.handle, c_str(key), c_str(str(value))))

    def get_internals(self):
        """Get a new grouped symbol whose output contains
        internal outputs of this symbol.

        Returns
        -------
        sgroup : Symbol
            The internal of the symbol.
        """
        handle = SymbolHandle()
        check_call(_LIB.MXSymbolGetInternals(
            self.handle, ctypes.byref(handle)))
        return Symbol(handle=handle)

    def get_children(self):
        """Get a new grouped symbol whose output contains
        inputs to output nodes of the original symbol

        Returns
        -------
        sgroup : Symbol or None
            The children of the head node. If the symbol has no
            inputs None will be returned.
        """
        handle = SymbolHandle()
        check_call(_LIB.MXSymbolGetChildren(
            self.handle, ctypes.byref(handle)))
        ret = Symbol(handle=handle)
        if len(ret.list_outputs()) == 0:
            return None
        return ret

    def list_arguments(self):
        """List all the arguments in the symbol.

        Returns
        -------
        args : list of string
            List of all the arguments.
        """
        size = ctypes.c_uint()
        sarr = ctypes.POINTER(ctypes.c_char_p)()
        check_call(_LIB.MXSymbolListArguments(
            self.handle, ctypes.byref(size), ctypes.byref(sarr)))
        return [py_str(sarr[i]) for i in range(size.value)]

    def list_outputs(self):
        """List all outputs in the symbol.

        Returns
        -------
        returns : list of string
            List of all the outputs.
        """
        size = ctypes.c_uint()
        sarr = ctypes.POINTER(ctypes.c_char_p)()
        check_call(_LIB.MXSymbolListOutputs(
            self.handle, ctypes.byref(size), ctypes.byref(sarr)))
        return [py_str(sarr[i]) for i in range(size.value)]

    def list_auxiliary_states(self):
        """List all auxiliary states in the symbol.

        Returns
        -------
        aux_states : list of string
            List the names of the auxiliary states.

        Notes
        -----
        Auxiliary states are special states of symbols that do not corresponds to an argument,
        and do not have gradient. But still be useful for the specific operations.
        A common example of auxiliary state is the moving_mean and moving_variance in BatchNorm.
        Most operators do not have Auxiliary states.
        """
        size = ctypes.c_uint()
        sarr = ctypes.POINTER(ctypes.c_char_p)()
        check_call(_LIB.MXSymbolListAuxiliaryStates(
            self.handle, ctypes.byref(size), ctypes.byref(sarr)))
        return [py_str(sarr[i]) for i in range(size.value)]

    def infer_type(self, *args, **kwargs):
        """Infer the type of outputs and arguments of given known types of arguments.

        User can either pass in the known types in positional way or keyword argument way.
        Tuple of Nones is returned if there is not enough information passed in.
        An error will be raised if there is inconsistency found in the known types passed in.

        Parameters
        ----------
        *args :
            Provide type of arguments in a positional way.
            Unknown type can be marked as None

        **kwargs :
            Provide keyword arguments of known types.

        Returns
        -------
        arg_types : list of numpy.dtype or None
            List of types of arguments.
            The order is in the same order as list_arguments()
        out_types : list of numpy.dtype or None
            List of types of outputs.
            The order is in the same order as list_outputs()
        aux_types : list of numpy.dtype or None
            List of types of outputs.
            The order is in the same order as list_auxiliary_states()
        """
        # pylint: disable=too-many-locals
        if len(args) != 0 and len(kwargs) != 0:
            raise ValueError('Can only specify known argument \
                    types either by positional or kwargs way.')
        sdata = []
        if len(args) != 0:
            keys = None
            for s in args:
                if s is not None:
                    s = _numpy.dtype(s).type
                    if s not in _DTYPE_NP_TO_MX:
                        raise TypeError('Argument need to be one of '+str(_DTYPE_NP_TO_MX))
                    sdata.append(_DTYPE_NP_TO_MX[s])
                else:
                    sdata.append(-1)
        else:
            keys = []
            for k, v in kwargs.items():
                v = _numpy.dtype(v).type
                if v in _DTYPE_NP_TO_MX:
                    keys.append(c_str(k))
                    sdata.append(_DTYPE_NP_TO_MX[v])
        arg_type_size = mx_uint()
        arg_type_data = ctypes.POINTER(ctypes.c_int)()
        out_type_size = mx_uint()
        out_type_data = ctypes.POINTER(ctypes.c_int)()
        aux_type_size = mx_uint()
        aux_type_data = ctypes.POINTER(ctypes.c_int)()
        complete = ctypes.c_int()
        check_call(_LIB.MXSymbolInferType(
            self.handle,
            mx_uint(len(sdata)),
            c_array(ctypes.c_char_p, keys),
            c_array(ctypes.c_int, sdata),
            ctypes.byref(arg_type_size),
            ctypes.byref(arg_type_data),
            ctypes.byref(out_type_size),
            ctypes.byref(out_type_data),
            ctypes.byref(aux_type_size),
            ctypes.byref(aux_type_data),
            ctypes.byref(complete)))
        if complete.value != 0:
            arg_types = [
                _DTYPE_MX_TO_NP[arg_type_data[i]] for i in range(arg_type_size.value)]
            out_types = [
                _DTYPE_MX_TO_NP[out_type_data[i]] for i in range(out_type_size.value)]
            aux_types = [
                _DTYPE_MX_TO_NP[aux_type_data[i]] for i in range(aux_type_size.value)]
            return (arg_types, out_types, aux_types)
        else:
            return (None, None, None)
        # pylint: enable=too-many-locals

    def infer_shape(self, *args, **kwargs):
        """Infer the shape of outputs and arguments of given known shapes of arguments.

        User can either pass in the known shapes in positional way or keyword argument way.
        Tuple of Nones is returned if there is not enough information passed in.
        An error will be raised if there is inconsistency found in the known shapes passed in.

        Parameters
        ----------
        *args :
            Provide shape of arguments in a positional way.
            Unknown shape can be marked as None

        **kwargs :
            Provide keyword arguments of known shapes.

        Returns
        -------
        arg_shapes : list of tuple or None
            List of shapes of arguments.
            The order is in the same order as list_arguments()
        out_shapes : list of tuple or None
            List of shapes of outputs.
            The order is in the same order as list_outputs()
        aux_shapes : list of tuple or None
            List of shapes of outputs.
            The order is in the same order as list_auxiliary_states()
        """
        try:
            res = self._infer_shape_impl(False, *args, **kwargs)
            if res[1] is None:
                arg_shapes, _, _ = self._infer_shape_impl(True, *args, **kwargs)
                arg_names = self.list_arguments()
                unknowns = []
                for name, shape in zip(arg_names, arg_shapes):
                    if not shape or not _numpy.prod(shape):
                        if len(unknowns) >= 10:
                            unknowns.append('...')
                            break
                        unknowns.append('%s: %s'%(name, str(shape)))
                warnings.warn(
                    "Cannot decide shape for the following arguments " +
                    "(0s in shape means unknown dimensions). " +
                    "Consider providing them as input:\n\t" +
                    "\n\t".join(unknowns), stacklevel=2)
            return res
        except MXNetError:
            print("infer_shape error. Arguments:")
            for i, arg in enumerate(args):
                print("  #%d: %s" % (i, arg))
            for k, v in kwargs.items():
                print("  %s: %s" % (k, v))
            raise

    def infer_shape_partial(self, *args, **kwargs):
        """Partially infer the shape. The same as infer_shape, except that the partial
        results can be returned.
        """
        return self._infer_shape_impl(True, *args, **kwargs)

    def _infer_shape_impl(self, partial, *args, **kwargs):
        """The actual implementation for calling shape inference API."""
        # pylint: disable=too-many-locals
        if len(args) != 0 and len(kwargs) != 0:
            raise ValueError('Can only specify known argument \
                    shapes either by positional or kwargs way.')
        sdata = []
        indptr = [0]
        if len(args) != 0:
            keys = None
            for s in args:
                if s is not None:
                    if not isinstance(s, tuple):
                        raise TypeError('Argument need to be shapes(tuple)')
                    sdata.extend(s)
                indptr.append(len(sdata))
        else:
            keys = []
            for k, v in kwargs.items():
                if isinstance(v, tuple):
                    keys.append(c_str(k))
                    sdata.extend(v)
                    indptr.append(len(sdata))
        arg_shape_size = mx_uint()
        arg_shape_ndim = ctypes.POINTER(mx_uint)()
        arg_shape_data = ctypes.POINTER(ctypes.POINTER(mx_uint))()
        out_shape_size = mx_uint()
        out_shape_ndim = ctypes.POINTER(mx_uint)()
        out_shape_data = ctypes.POINTER(ctypes.POINTER(mx_uint))()
        aux_shape_size = mx_uint()
        aux_shape_ndim = ctypes.POINTER(mx_uint)()
        aux_shape_data = ctypes.POINTER(ctypes.POINTER(mx_uint))()
        complete = ctypes.c_int()
        if partial:
            infer_func = _LIB.MXSymbolInferShapePartial
        else:
            infer_func = _LIB.MXSymbolInferShape
        check_call(infer_func(
            self.handle,
            mx_uint(len(indptr) - 1),
            c_array(ctypes.c_char_p, keys),
            c_array(mx_uint, indptr),
            c_array(mx_uint, sdata),
            ctypes.byref(arg_shape_size),
            ctypes.byref(arg_shape_ndim),
            ctypes.byref(arg_shape_data),
            ctypes.byref(out_shape_size),
            ctypes.byref(out_shape_ndim),
            ctypes.byref(out_shape_data),
            ctypes.byref(aux_shape_size),
            ctypes.byref(aux_shape_ndim),
            ctypes.byref(aux_shape_data),
            ctypes.byref(complete)))
        if complete.value != 0:
            arg_shapes = [
                tuple(arg_shape_data[i][:arg_shape_ndim[i]]) for i in range(arg_shape_size.value)]
            out_shapes = [
                tuple(out_shape_data[i][:out_shape_ndim[i]]) for i in range(out_shape_size.value)]
            aux_shapes = [
                tuple(aux_shape_data[i][:aux_shape_ndim[i]]) for i in range(aux_shape_size.value)]
            return (arg_shapes, out_shapes, aux_shapes)
        else:
            return (None, None, None)
        # pylint: enable=too-many-locals

    def debug_str(self):
        """Get a debug string.

        Returns
        -------
        debug_str : string
            Debug string of the symbol.
        """
        debug_str = ctypes.c_char_p()
        check_call(_LIB.MXSymbolPrint(
            self.handle, ctypes.byref(debug_str)))
        return py_str(debug_str.value)

    def save(self, fname):
        """Save symbol into file.

        You can also use pickle to do the job if you only work on python.
        The advantage of load/save is the file is language agnostic.
        This means the file saved using save can be loaded by other language binding of mxnet.
        You also get the benefit being able to directly load/save from cloud storage(S3, HDFS)

        Parameters
        ----------
        fname : str
            The name of the file
            - s3://my-bucket/path/my-s3-symbol
            - hdfs://my-bucket/path/my-hdfs-symbol
            - /path-to/my-local-symbol

        See Also
        --------
        symbol.load : Used to load symbol from file.
        """
        if not isinstance(fname, string_types):
            raise TypeError('fname need to be string')
        check_call(_LIB.MXSymbolSaveToFile(self.handle, c_str(fname)))

    def tojson(self):
        """Save symbol into a JSON string.

        See Also
        --------
        symbol.load_json : Used to load symbol from JSON string.
        """
        json_str = ctypes.c_char_p()
        check_call(_LIB.MXSymbolSaveToJSON(self.handle, ctypes.byref(json_str)))
        return py_str(json_str.value)

    @staticmethod
    def _get_ndarray_inputs(arg_key, args, arg_names, allow_missing):
        """Helper function to get ndarray lists handles from various inputs.

        Parameters
        ----------
        arg_key : str
            The name of argument, used for error message.

        args : list of NDArray or dict of str to NDArray
            Input arguments to the symbols.
            If type is list of NDArray, the position is in the same order of arg_names.
            If type is dict of str to NDArray, then it maps the name of arguments
            to the corresponding NDArray,

        args_names : list of string
            List of argument names.

        allow_missing : boolean
            Whether missing argument is allowed.
            When allowed, the missing handle will be set to None(null)

        Returns
        -------
        handles : list of NDArrayHandle
            The positional list of NDArrayHandles generated from input.
        """
        # setup args
        arg_handles = []
        arg_arrays = []
        if isinstance(args, list):
            if len(args) != len(arg_names):
                raise ValueError('Length of %s do not match number of arguments' % arg_key)
            for narr in args:
                if not isinstance(narr, NDArray):
                    raise TypeError('Only Accept list of NDArrays or dict of str to NDArray')
                arg_handles.append(narr.handle)
            arg_arrays = args
        elif isinstance(args, dict):
            for name in arg_names:
                if name in args:
                    narr = args[name]
                    if not isinstance(narr, NDArray):
                        raise TypeError('Only Accept list of NDArrays or dict of str to NDArray')
                    arg_handles.append(narr.handle)
                    arg_arrays.append(narr)
                else:
                    if allow_missing:
                        arg_handles.append(None)
                        arg_arrays.append(None)
                    else:
                        raise ValueError('key `%s` is missing in `%s`' % (name, arg_key))
        else:
            raise TypeError('Only Accept list of NDArrays or dict of str to NDArray')
        return c_array(NDArrayHandle, arg_handles), arg_arrays

    def simple_bind(self, ctx,
                    grad_req='write',
                    type_dict=None,
                    group2ctx=None,
                    **kwargs):
        """Bind current symbol to get an executor, allocate all the ndarrays needed.
        Allows specifying data types.

        This function will ask user to pass in ndarray of position
        they like to bind to, and it will automatically allocate the ndarray
        for arguments and auxiliary states that user did not specify explicitly.

        Parameters
        ----------
        ctx : Context
            The device context the generated executor to run on.

        grad_req: string
            {'write', 'add', 'null'}, or list of str or dict of str to str, optional
            Specifies how we should update the gradient to the args_grad.
            - 'write' means everytime gradient is write to specified args_grad NDArray.
            - 'add' means everytime gradient is add to the specified NDArray.
            - 'null' means no action is taken, the gradient may not be calculated.

        type_dict  : dict of str->numpy.dtype
            Input type dictionary, name->dtype

        group2ctx : dict of string to mx.Context
            The dict mapping the ``ctx_group`` attribute to the context assignment.

        kwargs : dict of str->shape
            Input shape dictionary, name->shape

        Returns
        -------
        executor : mxnet.Executor
            The generated Executor
        """
        # pylint: disable=too-many-locals
        if type_dict is None:
            attrs = self.attr_dict()
            type_dict = {k: mx_real_t for k in self.list_arguments()
                         if k not in attrs or '__dtype__' not in attrs[k]}
        arg_shapes, _, aux_shapes = self.infer_shape(**kwargs)
        arg_types, _, aux_types = self.infer_type(**type_dict)

        if arg_shapes is None or arg_types is None:
            raise ValueError("Input node is not complete")

        if group2ctx is not None:
            attr_dict = self.attr_dict()
            arg_ctx = [group2ctx.get(attr_dict[name]['__ctx_group__'], ctx) \
                         if name in attr_dict and '__ctx_group__' in attr_dict[name] \
                         else ctx for name in self.list_arguments()]
            aux_ctx = [group2ctx.get(attr_dict[name]['__ctx_group__'], ctx) \
                         if name in attr_dict and '__ctx_group__' in attr_dict[name] \
                         else ctx for name in self.list_auxiliary_states()]
        else:
            arg_ctx = [ctx] * len(arg_shapes)
            aux_ctx = [ctx] * len(aux_shapes)

        # alloc space
        arg_ndarrays = [
            _nd_zeros(shape, dev, dtype=dtype)
            for dtype, dev, shape in zip(arg_types, arg_ctx, arg_shapes)]
        if grad_req != 'null':
            grad_ndarrays = {}
            for name, shape, dev, dtype in zip(
                    self.list_arguments(), arg_shapes, arg_ctx, arg_types):
                if not isinstance(grad_req, dict) or grad_req[name] != 'null':
                    grad_ndarrays[name] = _nd_zeros(shape, dev, dtype=dtype)
        else:
            grad_ndarrays = None

        aux_ndarrays = [_nd_zeros(shape, dev, dtype=dtype)
                        for shape, dev, dtype in zip(aux_shapes, aux_ctx, aux_types)]
        executor = self.bind(ctx, arg_ndarrays,
                             grad_ndarrays, grad_req, aux_ndarrays,
                             group2ctx=group2ctx)
        return executor

    def bind(self, ctx, args, args_grad=None, grad_req='write',
             aux_states=None, group2ctx=None, shared_exec=None):
        """Bind current symbol to get an executor.

        Parameters
        ----------
        ctx : Context
            The device context the generated executor to run on.

        args : list of NDArray or dict of str to NDArray
            Input arguments to the symbol.

            - If type is list of NDArray, the position is in the same order of list_arguments.
            - If type is dict of str to NDArray, then it maps the name of arguments
              to the corresponding NDArray.
            - In either case, all the arguments must be provided.

        args_grad : list of NDArray or dict of str to NDArray, optional
            When specified, args_grad provide NDArrays to hold
            the result of gradient value in backward.

            - If type is list of NDArray, the position is in the same order of list_arguments.
            - If type is dict of str to NDArray, then it maps the name of arguments
              to the corresponding NDArray.
            - When the type is dict of str to NDArray, users only need to provide the dict
              for needed argument gradient.
              Only the specified argument gradient will be calculated.

        grad_req : {'write', 'add', 'null'}, or list of str or dict of str to str, optional
            Specifies how we should update the gradient to the args_grad.

            - 'write' means everytime gradient is write to specified args_grad NDArray.
            - 'add' means everytime gradient is add to the specified NDArray.
            - 'null' means no action is taken, the gradient may not be calculated.

        aux_states : list of NDArray, or dict of str to NDArray, optional
            Input auxiliary states to the symbol, only need to specify when
            list_auxiliary_states is not empty.

            - If type is list of NDArray, the position is in the same order of list_auxiliary_states
            - If type is dict of str to NDArray, then it maps the name of auxiliary_states
              to the corresponding NDArray,
            - In either case, all the auxiliary_states need to be provided.

        group2ctx : dict of string to mx.Context
            The dict mapping the ``ctx_group`` attribute to the context assignment.

        shared_exec : mx.executor.Executor
            Executor to share memory with. This is intended for runtime reshaping, variable length
            sequences, etc. The returned executor shares state with shared_exec, and should not be
            used in parallel with it.

        Returns
        -------
        executor : Executor
            The generated Executor

        Notes
        -----
        Auxiliary states are special states of symbols that do not corresponds to an argument,
        and do not have gradient. But still be useful for the specific operations.
        A common example of auxiliary state is the moving_mean and moving_variance in BatchNorm.
        Most operators do not have auxiliary states and this parameter can be safely ignored.

        User can give up gradient by using a dict in args_grad and only specify
        gradient they interested in.
        """
        # pylint: disable=too-many-locals, too-many-branches
        if not isinstance(ctx, Context):
            raise TypeError("Context type error")

        listed_arguments = self.list_arguments()
        args_handle, args = self._get_ndarray_inputs('args', args, listed_arguments, False)
        # setup args gradient
        if args_grad is None:
            args_grad_handle = c_array(NDArrayHandle, [None] * len(args))
        else:
            args_grad_handle, args_grad = self._get_ndarray_inputs(
                'args_grad', args_grad, listed_arguments, True)

        if aux_states is None:
            aux_states = []
        aux_args_handle, aux_states = self._get_ndarray_inputs(
            'aux_states', aux_states, self.list_auxiliary_states(), False)

        # setup requirements
        req_map = {'null': 0, 'write': 1, 'add': 3}
        if isinstance(grad_req, string_types):
            if grad_req not in req_map:
                raise ValueError('grad_req must be in %s' % str(req_map))
            reqs_array = c_array(mx_uint, [mx_uint(req_map[grad_req])] * len(listed_arguments))
        elif isinstance(grad_req, list):
            reqs_array = c_array(mx_uint, [mx_uint(req_map[item]) for item in grad_req])
        elif isinstance(grad_req, dict):
            req_array = []
            for name in listed_arguments:
                if name in grad_req:
                    req_array.append(mx_uint(req_map[grad_req[name]]))
                else:
                    req_array.append(mx_uint(0))
            reqs_array = c_array(mx_uint, req_array)

        ctx_map_keys = []
        ctx_map_dev_types = []
        ctx_map_dev_ids = []

        if group2ctx:
            for key, val in group2ctx.items():
                ctx_map_keys.append(c_str(key))
                ctx_map_dev_types.append(ctypes.c_int(val.device_typeid))
                ctx_map_dev_ids.append(ctypes.c_int(val.device_id))

        handle = ExecutorHandle()
        shared_handle = shared_exec.handle if shared_exec is not None else ExecutorHandle()
        check_call(_LIB.MXExecutorBindEX(self.handle,
                                         ctypes.c_int(ctx.device_typeid),
                                         ctypes.c_int(ctx.device_id),
                                         mx_uint(len(ctx_map_keys)),
                                         c_array(ctypes.c_char_p, ctx_map_keys),
                                         c_array(ctypes.c_int, ctx_map_dev_types),
                                         c_array(ctypes.c_int, ctx_map_dev_ids),
                                         mx_uint(len(args)),
                                         args_handle,
                                         args_grad_handle,
                                         reqs_array,
                                         mx_uint(len(aux_states)),
                                         aux_args_handle,
                                         shared_handle,
                                         ctypes.byref(handle)))
        executor = Executor(handle, self, ctx, grad_req, group2ctx)
        executor.arg_arrays = args
        executor.grad_arrays = args_grad
        executor.aux_arrays = aux_states
        return executor

    def grad(self, wrt):
        """Get the autodiff of current symbol.

        This function can only be used if current symbol is a loss function.

        Parameters
        ----------
        wrt : Array of String
            keyword arguments of the symbol that the gradients are taken.

        Returns
        -------
        grad : Symbol
            A gradient Symbol with returns to be the corresponding gradients.
        """
        handle = SymbolHandle()
        c_wrt = c_array(ctypes.c_char_p, [c_str(key) for key in wrt])
        check_call(_LIB.MXSymbolGrad(self.handle,
                                     mx_uint(len(wrt)),
                                     c_wrt,
                                     ctypes.byref(handle)))
        return Symbol(handle)
    # pylint: enable= no-member

    def eval(self, ctx=cpu(), **kwargs):
        """Evaluate a symbol given arguments

        The `eval` method combines a call to `bind` (which returns an executor)
        with a call to `forward` (executor method).
        For the common use case, where you might repeatedly evaluate with same arguments,
        eval is slow.
        In that case, you should call `bind` once and then repeatedly call forward.
        Eval allows simpler syntax for less cumbersome introspection.

        Parameters
        ----------
        ctx : Context
            The device context the generated executor to run on.

        kwargs : list of NDArray or dict of str to NDArray
            Input arguments to the symbol.

            - If type is list of NDArray, the position is in the same order of list_arguments.
            - If type is dict of str to NDArray, then it maps the name of arguments
              to the corresponding NDArray.
            - In either case, all the arguments must be provided.

        Returns
        ----------
        result :  a list of NDArrays corresponding to the values
        taken by each symbol when evaluated on given args.
        When called on a single symbol (not a group),
        the result will be a list with one element.
        """
        return self.bind(ctx, kwargs).forward()

def var(name, attr=None, shape=None, lr_mult=None, wd_mult=None, dtype=None, init=None):
    """Create a symbolic variable with specified name.

    Parameters
    ----------
    name : str
        Name of the variable.
    attr : dict of string -> string
        Additional attributes to set on the variable.
    shape : tuple
        Optionally, one can specify the shape of a variable. This will be used during
        shape inference. If user specified a different shape for this variable using
        keyword argument when calling shape inference, this shape information will be ignored.
    lr_mult : float
        Specify learning rate muliplier for this variable.
    wd_mult : float
        Specify weight decay muliplier for this variable.
    dtype : str or numpy.dtype
        Similar to shape, we can specify dtype for this variable.
    init : initializer (mxnet.init.*)
        Specify initializer for this variable to override the default initializer

    Returns
    -------
    variable : Symbol
        The created variable symbol.
    """
    if not isinstance(name, string_types):
        raise TypeError('Expect a string for variable `name`')
    handle = SymbolHandle()
    check_call(_LIB.MXSymbolCreateVariable(c_str(name), ctypes.byref(handle)))
    ret = Symbol(handle)
    attr = AttrScope.current.get(attr)
    attr = {} if attr is None else attr
    if shape is not None:
        attr['__shape__'] = str(shape)
    if lr_mult is not None:
        attr['__lr_mult__'] = str(lr_mult)
    if wd_mult is not None:
        attr['__wd_mult__'] = str(wd_mult)
    if dtype is not None:
        attr['__dtype__'] = str(_DTYPE_NP_TO_MX[_numpy.dtype(dtype).type])
    if init is not None:
        attr['__init__'] = init.dumps()
    ret._set_attr(**attr)
    return ret

# for back compatibility
Variable = var

def Group(symbols):
    """Create a symbol that groups symbols together.

    Parameters
    ----------
    symbols : list
        List of symbols to be grouped.

    Returns
    -------
    sym : Symbol
        The created group symbol.
     """
    ihandles = []
    for sym in symbols:
        if not isinstance(sym, Symbol):
            raise TypeError('Expect Symbols in the list input')
        ihandles.append(sym.handle)
    handle = SymbolHandle()
    check_call(_LIB.MXSymbolCreateGroup(
        mx_uint(len(ihandles)),
        c_array(SymbolHandle, ihandles), ctypes.byref(handle)))
    return Symbol(handle)

def load(fname):
    """Load symbol from a JSON file.

    You can also use pickle to do the job if you only work on python.
    The advantage of load/save is the file is language agnostic.
    This means the file saved using save can be loaded by other language binding of mxnet.
    You also get the benefit being able to directly load/save from cloud storage(S3, HDFS)

    Parameters
    ----------
    fname : str
        The name of the file, examples:

        - `s3://my-bucket/path/my-s3-symbol`
        - `hdfs://my-bucket/path/my-hdfs-symbol`
        - `/path-to/my-local-symbol`

    Returns
    -------
    sym : Symbol
        The loaded symbol.

    See Also
    --------
    Symbol.save : Used to save symbol into file.
    """
    if not isinstance(fname, string_types):
        raise TypeError('fname need to be string')
    handle = SymbolHandle()
    check_call(_LIB.MXSymbolCreateFromFile(c_str(fname), ctypes.byref(handle)))
    return Symbol(handle)

def load_json(json_str):
    """Load symbol from json string.

    Parameters
    ----------
    json_str : str
        A json string.

    Returns
    -------
    sym : Symbol
        The loaded symbol.

    See Also
    --------
    Symbol.tojson : Used to save symbol into json string.
    """
    if not isinstance(json_str, string_types):
        raise TypeError('fname need to be string')
    handle = SymbolHandle()
    check_call(_LIB.MXSymbolCreateFromJSON(c_str(json_str), ctypes.byref(handle)))
    return Symbol(handle)

# Initialize the atomic symbol in startups
_init_symbol_module(Symbol, "mxnet")

# pylint: disable=no-member
# pylint: disable=redefined-builtin
def pow(base, exp):
    """ Raise base to an exp.

    Parameters
    ---------
    base: Symbol or Number
    exp: Symbol or Number

    Returns
    -------
    result: Symbol or Number
    """
    if isinstance(base, Symbol) and isinstance(exp, Symbol):
        return _internal._Power(base, exp)
    if isinstance(base, Symbol) and isinstance(exp, Number):
        return _internal._PowerScalar(base, scalar=exp)
    if isinstance(base, Number) and isinstance(exp, Symbol):
        return _internal._RPowerScalar(exp, scalar=base)
    if isinstance(base, Number) and isinstance(exp, Number):
        return base**exp
    else:
        raise TypeError('types (%s, %s) not supported' % (str(type(base)), str(type(exp))))

# pylint: disable=no-member
# pylint: disable=redefined-builtin
def maximum(left, right):
    """ maximum left and right

    Parameters
    ---------
    left: Symbol or Number
    right: Symbol or Number

    Returns
    -------
    result: Symbol or Number
    """
    if isinstance(left, Symbol) and isinstance(right, Symbol):
        return _internal._Maximum(left, right)
    if isinstance(left, Symbol) and isinstance(right, Number):
        return _internal._MaximumScalar(left, scalar=right)
    if isinstance(left, Number) and isinstance(right, Symbol):
        return _internal._MaximumScalar(right, scalar=left)
    if isinstance(left, Number) and isinstance(right, Number):
        return left if left > right else right
    else:
        raise TypeError('types (%s, %s) not supported' % (str(type(left)), str(type(right))))

# pylint: disable=no-member
# pylint: disable=redefined-builtin
def minimum(left, right):
    """ minimum left and right

    Parameters
    ---------
    left: Symbol or Number
    right: Symbol or Number

    Returns
    -------
    result: Symbol or Number
    """
    if isinstance(left, Symbol) and isinstance(right, Symbol):
        return _internal._Minimum(left, right)
    if isinstance(left, Symbol) and isinstance(right, Number):
        return _internal._MinimumScalar(left, scalar=right)
    if isinstance(left, Number) and isinstance(right, Symbol):
        return _internal._MinimumScalar(right, scalar=left)
    if isinstance(left, Number) and isinstance(right, Number):
        return left if left > right else right
    else:
        raise TypeError('types (%s, %s) not supported' % (str(type(left)), str(type(right))))

# pylint: disable=no-member
# pylint: disable=redefined-builtin
def hypot(left, right):
    """ minimum left and right

    Parameters
    ---------
    left: Symbol or Number
    right: Symbol or Number

    Returns
    -------
    result: Symbol or Number
    """
    if isinstance(left, Symbol) and isinstance(right, Symbol):
        return _internal._Hypot(left, right)
    if isinstance(left, Symbol) and isinstance(right, Number):
        return _internal._HypotScalar(left, scalar=right)
    if isinstance(left, Number) and isinstance(right, Symbol):
        return _internal._HypotScalar(right, scalar=left)
    if isinstance(left, Number) and isinstance(right, Number):
        return _numpy.hypot(left, right)
    else:
        raise TypeError('types (%s, %s) not supported' % (str(type(left)), str(type(right))))

def zeros(shape, dtype=None, **kwargs):
    """Return a new symbol of given shape and type, filled with zeros.

    Parameters
    ----------
    shape :  int or sequence of ints
        Shape of the new array.
    dtype : str or numpy.dtype, optional
        The value type of the inner value, default to np.float32

    Returns
    -------
    out : Symbol
        The created Symbol
    """
    if dtype is None:
        dtype = _numpy.float32
    return _internal._zeros(shape=shape, dtype=dtype, **kwargs)

def ones(shape, dtype=None, **kwargs):
    """Return a new symbol of given shape and type, filled with ones.

    Parameters
    ----------
    shape :  int or sequence of ints
        Shape of the new array.
    dtype : str or numpy.dtype, optional
        The value type of the inner value, default to np.float32

    Returns
    -------
    out : Symbol
        The created Symbol
    """
    if dtype is None:
        dtype = _numpy.float32
    return _internal._ones(shape=shape, dtype=dtype, **kwargs)

def arange(start, stop=None, step=1.0, repeat=1, name=None, dtype=None):
    """Return evenly spaced values within a given interval.

    Parameters
    ----------
    start : number
        Start of interval. The interval includes this value. The default start value is 0.
    stop : number, optional
        End of interval. The interval does not include this value.
    step : number, optional
        Spacing between values
    repeat : int, optional
        "The repeating time of all elements.
        E.g repeat=3, the element a will be repeated three times --> a, a, a.
    dtype : str or numpy.dtype, optional
        The value type of the inner value, default to np.float32

    Returns
    -------
    out : Symbol
        The created Symbol
    """
    if dtype is None:
        dtype = _numpy.float32
    return _internal._arange(start=start, stop=stop, step=step, repeat=repeat,
                             name=name, dtype=dtype)


"""Code to read Peter Behroozi's Rockstar/consistent-trees output"""
from __future__ import print_function

import os
import numpy as np
from six.moves import xrange
from ..log import logger

class ConsistentTrees(object):
    def __init__(self, path):
        self._path = self._infer_subpath(path)
        self._load_raw_trees()
        self._load_scale_to_snap_number()
        self._setup_map_to_original_finder_catalogues()

    def _infer_subpath(self, rootpath):
        if os.path.exists(os.path.join(rootpath,'trees')):
            return rootpath
        elif os.path.exists(os.path.join(rootpath, 'halos', 'trees')):
            return os.path.join(rootpath, 'halos')
        else:
            raise IOError("Cannot find the consistent-trees output")

    def _load_raw_trees(self):
        filename = os.path.join(self._path, "trees", "tree_0_0_0.dat")
        f = open(filename)
        while True:
            l = f.readline()
            if l.startswith("#tree"):
                break
        read_type = np.dtype([
                  ('id_this', np.int64),
                  ('id_desc', np.int64),
                  ('phantom', np.int8),
                  ('Mvir', np.float32)])
        results = np.loadtxt(f, usecols=(1,3,8,10), dtype=read_type)
        self.links = results

    def _load_scale_to_snap_number(self):
        filename = os.path.join(self._path, "outputs", "scales.txt")
        read_type = np.dtype([('snap_number', np.int32), ('scale', np.float32)])
        self.scale_to_snap = np.loadtxt(filename, dtype=read_type)
        self._snap_min = self.scale_to_snap['snap_number'].min()
        self._snap_max = self.scale_to_snap['snap_number'].max()

    def _load_original_catalogue(self, snapnum):
        filename = os.path.join(self._path, "outputs", "really_consistent_%d.list"%snapnum)
        read_cols = (0, 49)
        return np.loadtxt(filename, dtype=np.int64, usecols=read_cols, unpack=True)

    def _setup_map_to_original_finder_catalogues(self):
        maxval = self.links['id_this'].max()
        self._id_to_finder_id = np.zeros(maxval + 1, dtype=np.int64)
        self._id_to_snap_num = np.zeros(maxval + 1, dtype=np.int32) - 1
        for snapnum in xrange(self._snap_min, self._snap_max + 1):
            ctid, original_id = self._load_original_catalogue(snapnum)
            self._id_to_finder_id[ctid] = original_id
            self._id_to_snap_num[ctid] = snapnum

        self._sanity_check_snap_num_assignment()
        self._snap_nums = self._get_snapshot_nums(self.links['id_this'])

        ids = self.links['id_this']
        phantom = self.links['phantom']
        # any remaining IDs correspond to phantom halos
        for snapnum in range(self._snap_min, self._snap_max + 1):
            phantom_mask = ((self._snap_nums == snapnum) & (phantom != 0))
            num_phantoms = phantom_mask.sum()
            phantom_ids = -np.arange(1, num_phantoms + 1)
            self._id_to_finder_id[ids[phantom_mask]] = phantom_ids

    def _sanity_check_snap_num_assignment(self):
        snap_nums_this = self._get_snapshot_nums(self.links['id_this'])
        snap_nums_next = self._get_snapshot_nums(self.links['id_desc'])
        snap_nums_next[snap_nums_next == -1] = self._snap_max + 1
        snap_nums_diff = snap_nums_next-snap_nums_this
        assert (snap_nums_diff == 1).all()

    def _get_finder_ids(self, consistent_trees_ids):
        return self._id_to_finder_id[consistent_trees_ids]

    def _get_snapshot_nums(self, consistent_trees_ids):
        snapnums = self._id_to_snap_num[consistent_trees_ids]
        snapnums[consistent_trees_ids<0]=-1
        return snapnums

    def get_num_phantoms_in_snapshot(self, snapnum):
        finder_ids = self._id_to_finder_id[self._id_to_snap_num==snapnum]
        if len(finder_ids)==0:
            return 0
        else:
            return -finder_ids.min()

    def get_links_for_snapshot(self, snapnum):
        """Get the links from snapshot snapnum to its immediate successor.

        Returns a dictionary; keys are the finder IDs at the given snapnum, whereas the values are
        a tuple containing the ID at the subsequent snapshot and the merger ratio (or 1.0 for no merger).

        Negative values indicate phantom halos, i.e. halos that were not present in the finder output"""

        this_snap_mask = self._snap_nums == snapnum
        ids_this_snap = self._id_to_finder_id[self.links['id_this'][this_snap_mask]]
        ids_next_snap = self._id_to_finder_id[self.links['id_desc'][this_snap_mask]]

        merger_ratios = self._get_merger_ratio_array(ids_next_snap, snapnum)

        return dict(zip(ids_this_snap, zip(ids_next_snap, merger_ratios)))

    def get_finder_id_to_tree_id_for_snapshot(self, snapnum):
        """Get the internal consistent-trees ids for each original halo-finder ID"""
        this_snap_mask = (self._snap_nums == snapnum)
        internal_ids_this_snap = self.links['id_this'][this_snap_mask]
        finder_ids_this_snap = self._id_to_finder_id[internal_ids_this_snap]
        finder_id_to_id = dict(zip(finder_ids_this_snap, internal_ids_this_snap))
        return finder_id_to_id

    def _get_merger_ratio_array(self, ids_next_snap, snapnum):
        this_snap_mask = self._snap_nums == snapnum
        ratio = np.ones(len(ids_next_snap))
        num_occurences_next_snap = np.bincount(ids_next_snap[ids_next_snap >= 0])
        mergers_next_snap = np.where(num_occurences_next_snap > 1)[0]
        logger.info("Identified %d mergers between snapshot %d and %d", len(mergers_next_snap), snapnum, snapnum + 1)
        for merger in mergers_next_snap:
            contributor_offsets = np.where(ids_next_snap == merger)[0]
            contributing_masses = self.links['Mvir'][this_snap_mask][contributor_offsets]
            ratio[contributor_offsets] = contributing_masses / contributing_masses.sum()
        return ratio


import numpy as np
from gripper_module.gripper_base import GripperBase
import time
import threading

class GripperRobotiq2F140(GripperBase):
    def __init__(self, bullet_client, gripper_size):
        r""" Initialization of robotiq-2f-140 gripper
        specific args for robotiq-2f-140:
            - gripper_size: global scaling of the gripper when loading URDF
        """
        super().__init__()

        self._bullet_client = bullet_client
        self._gripper_size = gripper_size

        # offset the gripper to a down facing pose for grasping
        self._pos_offset = np.array([0, 0, 0.235 * self._gripper_size]) # offset from base to center of grasping
        self._orn_offset = self._bullet_client.getQuaternionFromEuler([np.pi, 0, np.pi / 2])

        # define force and speed (grasping)
        self._force = 50
        self._grasp_speed = 0.8

        # define driver joint; the follower joints need to satisfy constraints when grasping
        self._driver_joint_id = 0
        self._follower_joint_ids = [5, 2, 7, 4, 9]
        self._follower_joint_sign = [-1, 1, 1, -1, -1]
        self._driver_joint_lower = 0.0
        self._driver_joint_upper = 0.7

    def load(self, basePosition):
        gripper_urdf = "assets/gripper/robotiq_2f_140/model.urdf"
        body_id = self._bullet_client.loadURDF(
            gripper_urdf,
            flags=self._bullet_client.URDF_USE_SELF_COLLISION,
            globalScaling=self._gripper_size,
            basePosition=basePosition
        )
        return body_id

    def configure(self, mount_gripper_id, n_links_before):
        # Set friction coefficients for gripper fingers
        for i in range(n_links_before, self._bullet_client.getNumJoints(mount_gripper_id)):
            self._bullet_client.changeDynamics(mount_gripper_id,i,lateralFriction=1.0,spinningFriction=1.0,rollingFriction=0.0001,frictionAnchor=True)

    def step_constraints(self, mount_gripper_id, n_joints_before):
        pos = self._bullet_client.getJointState(mount_gripper_id, self._driver_joint_id+n_joints_before)[0]
        targets = pos * np.array(self._follower_joint_sign)
        self._bullet_client.setJointMotorControlArray(
            mount_gripper_id,
            [(joint_id+n_joints_before) for joint_id in self._follower_joint_ids],
            self._bullet_client.POSITION_CONTROL,
            targetPositions=targets,
            forces=[self._force] * len(self._follower_joint_ids),
            positionGains=[1.5] * len(self._follower_joint_ids)
        )
        return pos

    def open(self, mount_gripper_id, n_joints_before, open_scale):
        target_pos = open_scale*self._driver_joint_lower + (1-open_scale)*self._driver_joint_upper  # recalculate scale because larger joint position corresponds to smaller open width
        self._bullet_client.setJointMotorControl2(
            mount_gripper_id,
            self._driver_joint_id+n_joints_before,
            self._bullet_client.POSITION_CONTROL,
            targetPosition=target_pos,
            force=self._force
        )
        for _ in range(240 * 2):
            driver_pos = self.step_constraints(mount_gripper_id, n_joints_before)
            states_diff = np.abs(np.array(driver_pos) - self._driver_joint_lower)
            if np.abs(driver_pos - target_pos)<1e-5:
                break
            self._bullet_client.stepSimulation()

    def close(self, mount_gripper_id, n_joints_before):
        self._bullet_client.setJointMotorControl2(
            mount_gripper_id,
            self._driver_joint_id+n_joints_before,
            self._bullet_client.VELOCITY_CONTROL,
            targetVelocity=self._grasp_speed,
            force=self._force
        )
        for i in range(240 * 3):
            pos = self.step_constraints(mount_gripper_id, n_joints_before)
            if pos>self._driver_joint_upper:
                break
            self._bullet_client.stepSimulation()

    def get_pos_offset(self):
        return self._pos_offset

    def get_orn_offset(self):
        return self._orn_offset

    def get_vis_pts(self, open_scale):
        width = 0.075 * np.sin(open_scale)
        return self._gripper_size * np.array([
            [-width, 0],
            [width, 0]
        ])