#!/usr/bin/python3

########################################################################
# This file is part of the Honeyris project made by the Astar Company: #
# https://github.com/astar-security/Honeyris                           #
# The project is published under GPLv3 license                         #
# Author: David Soria (@Sibwara)                                       #
########################################################################

import pyshark
import argparse
import logging
import logging.handlers
import signal
from scapy.all import *

blacklist = None
log = None

##################
# Whitelist part #
##################

# Get the local IPv4 addresses which will be whitelisted
def getHoneyIPAddresses(iface):
    global log
    ips = set()
    try:
        ips.add("127.0.0.1")
        ips.add(get_if_addr(iface))
        return ips
    except Exception as e:
        log.error(f"Error during local IP collection: {e}")
        return ips

# Get the local GW which will be whitelisted
def getHoneyGateway(iface):
    global log
    gws, macs = set(), set()
    try:
        # get default gw for chosen iface
        res = conf.route.route()
        if res[0] == iface :
            gws.add(res[2])
        # get the associated known MAC address
        arp = open("/proc/net/arp","r")
        table = arp.read()
        for gw in gws :
            mac = table.split(gw)[1].split("\n")[0].split()[2]
            macs.add(mac)
        return gws, macs
    except Exception as e:
        log.error(f"Error during gateway collection: {e}")
        return gws, macs

# As the DNS will be involved for the system update, it must be whitelisted
def getHoneyDNS():
    global log
    nss = set()
    try:
        # Parsing the content of "/etc/resolv.conf"
        resol = open("/etc/resolv.conf", "r")
        dns = resol.read().split("nameserver ")[1:]
        resol.close()
        for ns in dns:
            nss.add(ns.split("\n")[0])
        return nss
    except Exception as e:
        log.error(f"Error during DNS collection: {e}")
        return nss

# As the DHCP server will be joined to update leases, it must be whitelisted
def getHoneyDHCP(iface):
    global log
    dhcp = set()
    try:
        # make a DHCP discover request to identify the DHCP server
        conf.checkIPaddr=False
        localmac = get_if_hwaddr(iface)
        localmacraw = get_if_raw_hwaddr(iface)[1]
        dhcp_discover = (Ether(src=localmac, dst='ff:ff:ff:ff:ff:ff') /
            IP(src='0.0.0.0', dst='255.255.255.255') /
            UDP(dport=67, sport=68) /
            BOOTP(chaddr=localmacraw,xid=5555) /
            DHCP(options=[('message-type', 'discover'), 'end']))
        dhcp_offer = srp1(dhcp_discover,iface=iface, verbose=0, timeout=10)
        dhcp.add(dhcp_offer['IP'].src)
        return dhcp
    except Exception as e:
        log.error(f"Error during DHCP collection: {e}")
        return dhcp

# As the honeyris server will join targets to update, they must be whitelisted
#def getHoneyUpdater():

####################
# Prepare the meal #
####################

def populate(iface):
    global log
    global blacklist

    # call whitelist functions
    ips = getHoneyIPAddresses(iface)
    gws, macs = getHoneyGateway(iface)
    nss = getHoneyDNS()
    dhcp = getHoneyDHCP(iface)

    # aggregate and set the blacklist
    whitelist = ips.union(gws.union(nss.union(dhcp)))
    blacklist = set()

    # provide information about honeyris
    info = "Information about Honeyris-"\
        f"Interface:{iface}-IP:{ips}-Gateway:{gws}{macs}-NS:{nss}-DHCP:{dhcp}"
    log.info(info)

    return ips,gws,macs,nss,dhcp,whitelist

def setLog(siem):
    global log

    # default Syslog port
    port =  514

    # Set the command line logger
    log = logging.getLogger('Honeyris')
    log.setLevel(logging.INFO)
    formatter = logging.Formatter('%(name)s--%(levelname)s--%(asctime)s--%(message)s')
    console = logging.StreamHandler()
    console.setFormatter(formatter)
    log.addHandler(console)

    try:
        # check if a specific port is provided
        if ":" in siem:
            siem, port = siem.split(":")
            port = int(port)
        # set Syslog handler
        syslog = logging.handlers.SysLogHandler(address=(siem, port))
        syslog.setFormatter(formatter)
        log.addHandler(syslog)
        log.info("Logging ready")
        return 0
    except Exception as e:
        log.error(f"Error during logging initialization: {e}")
        return 1

# Clean exit after CTRL-C
def ctrlCHandler(signum, frame):
    global log
    global blacklist
    log.info(f"quitting...")
    log.info(f"Blacklisted targets were: {blacklist}")

# send CTRL-C signal to ctrlCHandler
signal.signal(signal.SIGINT, ctrlCHandler)

####################
# Serve some honey #
####################

def blacklistIP(iface, IP, ARPPing, ARPSpoof, verbose, whitelist, ips, gws, macs):
    global blacklist
    global log

    log.info("Logging initiated")
    capture = pyshark.LiveCapture(interface=iface)
    for packet in capture.sniff_continuously():
        try:
            # if an IP request does not come from a trusted IP
            if (IP and 'IP' in packet and packet['ip'].dst in ips and
                    packet['ip'].src not in whitelist):
                blacklist.add(packet['ip'].src)
                log.warning(f"{packet['ip'].src}--IP request" + ['', f"--{packet}"][verbose])

            # if a ARP request does not come from a trusted IP
            elif (ARPPing and 'ARP' in packet and packet['arp'].opcode == '1'
            and packet['arp'].dst_proto_ipv4 in ips
            and packet['arp'].src_proto_ipv4 not in whitelist):
                blacklist.add(packet['arp'].src_proto_ipv4)
                log.warning(f"{packet['arp'].src_proto_ipv4}--ARP ping" + ['', f"--{packet}"][verbose])

            # if a ARP reply of the gateway does not come from a trusted MAC
            elif (ARPSpoof and 'ARP' in packet and packet['arp'].opcode == '2'
            and packet['arp'].src_hw_mac not in macs
            and packet['arp'].src_proto_ipv4 in gws):
                blacklist.add(packet['arp'].src_proto_ipv4)
                log.warning(f"{packet['arp'].src_hw_mac}--ARP spoof" + ['', f"--{packet}"][verbose])

        except Exception as e:
            log.error(f"Error during packet capture: {e}")
            continue

def main():
    parser = argparse.ArgumentParser(description='Detect suspucious activity from '\
            'the network', add_help=True)
    parser.add_argument('--ip', action="store_true", dest="IP", default=False,
            help='Enable IP touch blacklist')
    parser.add_argument('--arpping', action="store_true", dest="ARPPing", default=False,
            help='Enable ARP ping blacklist')
    parser.add_argument('--arpspoof', action="store_true", dest="ARPSpoof", default=False,
            help='Enable ARP spoof blacklist')
    parser.add_argument('--verbose', action="store_true", dest="verbose", default=False,
            help="Include packet content received from suspicious IP, WARNING : sensitive data could be transmitted through UDP syslog (cleartext)")
    parser.add_argument('iface',
            help="the network interface to monitor")
    parser.add_argument('siem',
            help="the server to send the UDP syslog alerts, can be 127.0.0.1:514")
    args = parser.parse_args()

    setLog(args.siem)
    ips, gws, macs, nss, dhcp, whitelist = populate(args.iface)
    blacklistIP(args.iface, args.IP, args.ARPPing, args.ARPSpoof, args.verbose, whitelist, ips, gws, macs)

if __name__ == '__main__':
    main()


import cv2
import pytesseract
from PIL import Image
import sys
import os
import pyocr
import pyocr.builders
import datetime
import argparse
import unidecode
import codecs

ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", required=True, help="Path to the image")
args = vars(ap.parse_args())

dt = datetime.datetime.now()
dt = str(dt).replace("-", "")
dt = str(dt).replace(" ", "")
dt = str(dt).replace(":", "")
dt = str(dt).replace(".", "")
dt = str(dt).replace("/", "")

# print(args["image"])

img_name = args["image"]
name, ext = os.path.splitext(str(img_name))
fname = str(name).split('/')[-1]
image_name = 'temp/jfs' + str(dt) + '.png'
file_name = 'files/ocr.txt'

#file1 = open(file_name, 'w')

tools = pyocr.get_available_tools()
tool = tools[0]

pytesseract.pytesseract.tesseract_cmd = 'C:/Program Files/Tesseract-OCR/tesseract'
TESSDATA_PREFIX = 'C:/Program Files/Tesseract-OCR'

try:
    im = cv2.imread(str(img_name), 0)
    newdata = pytesseract.image_to_osd(im)
    #image_name = img_name
    cv2.imwrite(image_name, im)
except:
    im = Image.open(img_name)
    nx, ny = im.size
    im2 = im.resize((int(nx * 2.5), int(ny * 2.5)), Image.BICUBIC)
    im2.save(image_name, dpi=(600, 600))
    im.close()

line_and_word_boxes = tool.image_to_string(
    Image.open(image_name),
    lang='best/eng',
    builder=pyocr.builders.TextBuilder()
)
print(line_and_word_boxes)
#file1.writelines(unidecode(line_and_word_boxes))
with open(file_name, "wb") as f:
   f.write(line_and_word_boxes.encode("UTF-8"))

#file1.close()

#try:
#    os.chmod(image_name, 0o777)
#    os.remove(image_name)
#except Exception as ef:
#    print(str(ef))


import importlib
import copy
import os
import re
import sys
import warnings
from collections import OrderedDict
import traceback

import h5py
import numpy as np

import deeprank
from deeprank.models.variant import PdbVariantSelection
from deeprank import config
from deeprank.config import logger
from deeprank.generate import GridTools as gt
from deeprank.operate import hdf5data
from deeprank.operate.pdb import get_pdb_path

import pdb2sql
from pdb2sql.align import align as align_along_axis
from pdb2sql.align import align_interface

try:
    from tqdm import tqdm
except ImportError:
    def tqdm(x):
        return x

try:
    from pycuda import driver, compiler, gpuarray, tools
    import pycuda.autoinit
except ImportError:
    pass

def _printif(string, cond): return print(string) if cond else None

class DataGenerator(object):

    # In the feature dataset, the first three columns are xyz.
    # The rest are feature values.
    FEATURE_POSITION_OFFSET = 3

    def __init__(self, environment, variants,
                 align=None,
                 compute_targets=None, compute_features=None,
                 data_augmentation=None, hdf5='database.h5',
                 mpi_comm=None):
        """Generate the data (features/targets/maps) required for deeprank.

        Args:
            environment (Environment): the environment settings
            variants (list(PdbVariantSelection)): the selected variants
            align (dict, optional): Dicitionary to align the compexes,
                                    e.g. align = {"selection":{"chainID":["A","B"]}, "axis":"z"}
                                    e.g. align = {"selection":{"chainID":["A","B"]}, "plane":"xy"}
                                    if "selection" is not specified the entire complex is used for alignement
            compute_targets (list(str), optional): List of python files computing the targets,
                "pdb_native" must be set if having targets to compute.
            compute_features (list(str), optional): List of python files computing the features
            data_augmentation (int, optional): Number of rotations performed one each complex
            hdf5 (str, optional): name of the hdf5 file where the data is saved, default to 'database.h5'
            mpi_comm (MPI_COMM): MPI COMMUNICATOR

        Raises:
            NotADirectoryError: if the source are not found

        Example:

            >>> from deeprank.generate import *
            >>> from deeprank.models.variant import *
            >>> from deeprank.domain.amino_acid import *
            >>> # sources to assemble the data base
            >>> variant = PdbVariantSelection(pdb_ac="1AK4",
            >>>                               chain_id="C",
            >>>                               residue_number=10,
            >>>                               wild_type_amino_acid=isoleucine,
            >>>                               variant_amino_acid=valine,
            >>>                               pssm_paths_by_chain={"C": "pssm_new/1AK4.C.pssm",
            >>>                                                    "D": "pssm_new/1AK4.D.pssm"},
            >>>                               variant_class=VariantClass.BENIGN)
            >>> h5file = '1ak4.hdf5'
            >>>
            >>> #init the data assembler
            >>> database = DataGenerator([variant],
            >>>                          data_augmentation=None,
            >>>                          compute_targets=['deeprank.targets.class'],
            >>>                          compute_features=['deeprank.features.atomic_contacts',
            >>>                                            'deeprank.features.accessibility',
            >>>                                            'deeprank.features.neighbour_profile'],
            >>>                          hdf5=h5file)
        """

        self.environment = environment

        self.variants = variants
        logger.debug("preprocess with {} variants".format(len(variants)))

        self.align = align
        logger.debug("preprocess with align={}".format(align))

        self.compute_targets = compute_targets
        self.compute_features = compute_features
        logger.debug("preprocess with targets={}, features={}".format(compute_targets, compute_features))

        self.data_augmentation = data_augmentation
        logger.debug("preprocess with data_augmentation={}".format(data_augmentation))

        self.hdf5 = hdf5

        self.mpi_comm = mpi_comm

        self.feature_error = []
        self.grid_error = []
        self.map_error = []

        self.logger = logger

# ====================================================================================
#
#       CREATE THE DATABASE ALL AT ONCE IF ALL OPTIONS ARE GIVEN
#
# ====================================================================================

    def create_database(
            self,
            verbose=False,
            remove_error=True,
            prog_bar=False,
            distance_cutoff=10.0,
            random_seed=None):
        """Create the hdf5 file architecture and compute the features/targets.

        Args:
            verbose (bool, optional): Print creation details
            remove_error (bool, optional): remove the groups that errored
            prog_bar (bool, optional): use tqdm
            distance_cutoff (float): max distance from center to include atoms, defaults to 10.0 Å
            random_seed (int): random seed for getting rotation axis and angle

        Raises:
            ValueError: If creation of the group errored.

        Example:

        >>> # sources to assemble the data base
        >>> variant = PdbVariantSelection(pdb_ac="1AK4",
        >>>                               chain_id="C",
        >>>                               residue_number=10,
        >>>                               amino_acid="T",
        >>>                               pssm_paths_by_chain={"C": "pssm_new/1AK4.C.pssm",
        >>>                                                    "D": "pssm_new/1AK4.D.pssm"},
        >>>                               variant_class=VariantClass.BENIGN)
        >>> h5file = '1ak4.hdf5'
        >>>
        >>> #init the data assembler
        >>> database = DataGenerator([variant],
        >>>                          data_augmentation=None,
        >>>                          compute_targets  = ['deeprank.targets.variant_class'],
        >>>                          compute_features = ['deeprank.features.AtomicFeature',
        >>>                                              'deeprank.features.PSSM_IC',
        >>>                                              'deeprank.features.BSA'],
        >>>                          hdf5=h5file)
        >>>
        >>> #create new files
        >>> database.create_database(prog_bar=True)
        """

        # deals with the parallelization
        self.local_variants = self.variants

        if self.mpi_comm is not None:
            rank = self.mpi_comm.Get_rank()
            size = self.mpi_comm.Get_size()
        else:
            size = 1

        if size > 1:
            if rank == 0:
                variants_divided = [self.variants[i::size] for i in range(size)]
                self.local_variants = variants_divided[0]
                # send to other procs
                for iP in range(1, size):
                    self.mpi_comm.send(variants_divided[iP], dest=iP, tag=11)
            else:
                # receive procs
                self.local_variants = self.mpi_comm.recv(source=0, tag=11)
            # change hdf5 name
            h5path, h5name = os.path.split(self.hdf5)
            self.hdf5 = os.path.join(h5path, f"{rank:03d}_{h5name}")

        # set metadata to hdf5 file
        with h5py.File(self.hdf5, 'w') as f5:

            f5.attrs['DeepRank_version'] = deeprank.__version__
            if self.compute_features is not None:
                f5.attrs['features'] = self.compute_features
            if self.compute_targets is not None:
                f5.attrs['targets'] = self.compute_targets

        ##################################################
        # Start generating HDF5 database
        ##################################################
        self.logger.info(
            f'\n# Start creating HDF5 database: {self.hdf5}')

        # get the local progress bar
        desc = '{:25s}'.format('Creating database')
        variant_tqdm = tqdm(self.local_variants, desc=desc,
                           disable=not prog_bar)

        for variant in variant_tqdm:

            pdb_path = get_pdb_path(self.environment.pdb_root, variant.pdb_ac)

            variant_tqdm.set_postfix(variant=os.path.basename(pdb_path))
            variant_name = hdf5data.get_variant_group_name(variant)
            self.logger.info(f'\nProcessing variant: {variant_name}')

            # names of the variant
            variant_aug_name_list = []

            with h5py.File(self.hdf5, 'a') as f5:

                ################################################
                #   get the pdbs of the conformation and its ref
                #   for the original data (not augmetned one)
                ################################################

                if verbose:
                    self.logger.info(
                        f'\nVariant: {variant_name}.'
                        f'\nStart generating top HDF5 group "{variant_name}"...'
                        f'\n{"":4s}Reading PDB data into database...')

                # get the bare name of the variant
                # and define the name of the native
                # i.e. 1AK4-m1324578236_100w -> 1AK4-m1324578236
                bare_variant_name = variant_name.split('_')[0]
                ref_name = bare_variant_name + '.pdb'

                # check if we have a decoy or native
                # and find the reference
                if variant_name == bare_variant_name:
                    ref = pdb_path
                else:
                    ref = None

                # create a subgroup for the variant
                variant_group = f5.require_group(variant_name)
                variant_group.attrs['type'] = 'variant'
                hdf5data.store_variant(variant_group, variant)

                # add the ref and the complex
                #self._add_pdb(variant_group, pdb_path, 'complex')
                #if ref is not None:
                #    self._add_pdb(variant_group, ref, 'native')

                if verbose:
                    self.logger.info(
                        f'{"":4s}Generated subgroup "complex"'
                        f' to store pdb data of the current model.')
                    if ref:
                        self.logger.info(
                            f'{"":4s}Generated subgroup "native"'
                            f' to store pdb data of the reference molecule.')

                ################################################
                #   add the features
                ################################################
                feature_error_flag = False  # when False: success; when True: failed

                if self.compute_features is not None:
                    if verbose:
                        self.logger.info(
                            f'{"":4s}Calculating features...')

                    variant_group.require_group('features')

                    feature_error_flag = self._compute_features(self.compute_features,
                                                                self.environment,
                                                                distance_cutoff,
                                                                variant_group['features'],
                                                                variant,
                                                                self.logger)
                    if feature_error_flag:
                        self.feature_error += [variant_name]
                        # ignore the targets/grid/augmentation computation
                        # and directly go to next variant. Remove errored
                        # variant later.
                        # Otherwise, keep computing and report errored variant.
                        if remove_error:
                            continue

                    if verbose:
                        if not feature_error_flag or not remove_error:
                            self.logger.info(
                                f'\n{"":4s}Generated subgroup "features"'
                                f' to store xyz-based feature values.'
                                f'{"":4s}Generated subgroup "features_raw"'
                                f' to store human read feature values')

                ################################################
                #   add the targets
                ################################################
                if self.compute_targets is not None:
                    if verbose:
                        self.logger.info(
                            f'{"":4s}Calculating targets...')

                    variant_group.require_group('targets')

                    self._compute_targets(self.compute_targets,
                                          variant,
                                          variant_group['targets'])

                    if verbose:
                        self.logger.info(
                            f'{"":4s}Generated subgroup "targets" '
                            f'to store targets, such as BIN_CLASS, dockQ, etc.')

                ################################################
                #   add the box center
                ################################################
                if verbose:
                    self.logger.info(
                        f'{"":4s}Calculating grid box center...')

                grid_error_flag = False
                variant_group.require_group('grid_points')

                try:
                    center = DataGenerator.get_grid_center(self.environment, variant)
                    variant_group['grid_points'].create_dataset('center', data=center, chunks=True)
                    if verbose:
                        self.logger.info(
                            f'{"":4s}Generated subgroup "grid_points"'
                            f' to store grid box center.')

                except:
                    grid_error_flag = True
                    self.grid_error += [variant_name]
                    self.logger.exception("Error while computing center for {}: {}".format(variant, traceback.format_exc()))
                    if remove_error:
                        continue

                ################################################
                #   DATA AUGMENTATION
                ################################################

                # GET ALL THE NAMES
                if self.data_augmentation is not None:
                    variant_aug_name_list = [
                        variant_name +
                        '_r%03d' %
                        (idir +
                         1) for idir in range(
                            self.data_augmentation)]
                else:
                    variant_aug_name_list = []

                if verbose and variant_aug_name_list:
                    self.logger.info(
                        f'{"":2s}Start augmenting data'
                        f' with {self.data_augmentation} times...')

                # loop over the variants
                for variant_aug_name in variant_aug_name_list:

                    # create a subgroup for the molecule
                    variant_group = f5.require_group(variant_aug_name)
                    variant_group.attrs['type'] = 'variant'
                    hdf5data.store_variant(variant_group, variant)

                    # copy the ref into it
                    #if ref is not None:
                    #    self._add_pdb(variant_group, ref, 'native')

                    # get the rotation axis and angle
                    if self.align is None:
                        rotation_axis, rotation_angle = pdb2sql.transform.get_rot_axis_angle(random_seed)
                    else:
                        rotation_axis, rotation_angle = self._get_aligned_rotation_axis_angle(random_seed, self.align)

                    # create the new pdb and get molecule center
                    # molecule center is the origin of rotation)
                    rotation_center = self._add_aug_pdb(variant_group, variant,
                                                        'complex', rotation_axis, rotation_angle)

                    # copy the targets/features
                    if 'targets' in f5[variant_name]:
                        f5.copy(variant_name + '/targets/', variant_group)
                    f5.copy(variant_name + '/features/', variant_group)

                    # rotate the feature
                    self._rotate_feature(variant_group, rotation_axis, rotation_angle, rotation_center)

                    # store the rotation axis/angl/center as attriutes
                    # in case we need them later
                    variant_group.attrs['rotation_axis'] = rotation_axis
                    variant_group.attrs['rotation_angle'] = rotation_angle
                    variant_group.attrs['rotation_center'] = rotation_center

                # cache aug variants if original variant has errored features
                if feature_error_flag:
                    self.feature_error += variant_aug_name_list
                if grid_error_flag:
                    self.grid_error += variant_aug_name_list

                if verbose and variant_aug_name_list:
                    self.logger.info(
                        f'{"":2s}Completed data augmentation'
                        f' and generated top HDF5 groups, e.g. {variant_aug_name}.')

                ################################################
                # Successul message
                ################################################
                if verbose:
                    self.logger.info(
                        f'\nSuccessfully generated top HDF5 group "{variant_name}".\n')

        ##################################################
        # Post processing
        ##################################################
        #  Remove errored variants
        errored_variants = list(set(self.feature_error + self.grid_error))
        if len(errored_variants) > 0:
            if remove_error:
                with h5py.File(self.hdf5, 'a') as f5:
                    for variant_name in errored_variants:
                        del f5[variant_name]
                if self.feature_error:
                    self.logger.info(
                        f'Molecules with errored features are removed:'
                        f'\n{self.feature_error}')
                if self.grid_error:
                    self.logger.info(
                        f'Molecules with errored grid points are removed:'
                        f'\n{self.grid_error}')
            else:
                if self.feature_error:
                    self.logger.warning(
                        f'The following variants have errored features:'
                        f'\n{self.feature_error}')
                if self.grid_error:
                    self.logger.warning(
                        f'The following variants have errored grid points:'
                        f'\n{self.grid_error}')

        # close the file
        self.logger.info(
            f'\n# Successfully created database: {self.hdf5}\n')

    def aug_data(self, augmentation, keep_existing_aug=True, random_seed=None):
        """Augment exiting original PDB data and features.

        Args:
            augmentation(int): Times of augmentation
            keep_existing_aug (bool, optional): Keep existing augmentated data.
                If False, existing aug will be removed. Defaults to True.

        Examples:
            >>> database = DataGenerator(h5='database.h5')
            >>> database.aug_data(augmentation=3, append=True)
            >>> grid_info = {
            >>>     'number_of_points': [20,20,20],
            >>>     'resolution': [1.,1.,1.],
            >>>     'atomic_densities': {'C':1.7, 'N':1.55, 'O':1.52, 'S':1.8},
            >>>     }
            >>> database.map_features(grid_info)
        """

        # check if file exists
        if not os.path.isfile(self.hdf5):
            raise FileNotFoundError(
                'File %s does not exists' % self.hdf5)

        # get the folder names
        f5 = h5py.File(self.hdf5, 'a')
        fnames = f5.keys()

        # get the non rotated ones
        fnames_original = list(
            filter(lambda x: not re.search(r'_r\d+$', x), fnames))

        # get the rotated ones
        fnames_augmented = list(
            filter(lambda x: re.search(r'_r\d+$', x), fnames))

        aug_id_start = 0
        if keep_existing_aug:
            exiting_augs = list(
                filter(lambda x: re.search(fnames_original[0] + r'_r\d+$', x), fnames_augmented))
            aug_id_start += len(exiting_augs)
        else:
            for i in fnames_augmented:
                del f5[i]

        self.logger.info(
            f'{"":s}\n# Start augmenting data'
            f' with {augmentation} times...')

        # GET ALL THE NAMES
        for variant_name in fnames_original:
            variant_aug_name_list = [
                variant_name + '_r%03d' % (idir + 1) for idir in
                range(aug_id_start, aug_id_start + augmentation)]

            variant = hdf5data.load_variant(f5[variant_name])

            # loop over the complexes
            for variant_aug_name in variant_aug_name_list:

                # create a subgroup for the variant
                variant_group = f5.require_group(variant_aug_name)
                variant_group.attrs['type'] = 'variant'
                hdf5data.store_variant(variant_group, variant)

                # copy the ref into it
                if 'native' in f5[variant_name]:
                    f5.copy(variant_name + '/native', variant_group)

                # get the rotation axis and angle
                if self.align is None:
                    rotation_axis, rotation_angle = pdb2sql.transform.get_rot_axis_angle(random_seed)
                else:
                    rotation_axis, rotation_angle = self._get_aligned_rotation_axis_angle(random_seed, self.align)

                # create the new pdb and get molecule center
                # molecule center is the origin of rotation)
                rotation_center = self._add_aug_pdb(variant_group, variant, 'complex',
                                                    rotation_axis, rotation_angle)

                # copy the targets/features
                if 'targets' in f5[variant_name]:
                    f5.copy(variant_name + '/targets/', variant_name)
                f5.copy(variant_name + '/features/', variant_group)

                # rotate the feature
                self._rotate_feature(variant_group, rotation_axis, rotation_angle, rotation_center)

                # store the rotation axis/angl/center as attriutes
                # in case we need them later
                variant_group.attrs['rotation_axis'] = rotation_axis
                variant_group.attrs['rotation_angle'] = rotation_angle
                variant_group.attrs['rotation_center'] = rotation_center
        f5.close()
        self.logger.info(
            f'\n# Successfully augmented data in {self.hdf5}')

# ====================================================================================
#
#       ADD FEATURES TO AN EXISTING DATASET
#
# ====================================================================================

    def add_feature(self, remove_error=True, prog_bar=True, distance_cutoff=10.0):
        """Add a feature to an existing hdf5 file.

        Args:
            remove_error (bool): remove errored variant
            prog_bar (bool, optional): use tqdm
            distance_cutoff (float, optional): max distance from center to include atoms, default 10.0 Å

        Example:

        >>> h5file = '1ak4.hdf5'
        >>>
        >>> #init the data assembler
        >>> database = DataGenerator(compute_features  = ['deeprank.features.ResidueDensity'],
        >>>                          hdf5=h5file)
        >>>
        >>> database.add_feature(remove_error=True, prog_bar=True)
        """

        # check if file exists
        if not os.path.isfile(self.hdf5):
            raise FileNotFoundError(
                'File %s does not exists' % self.hdf5)

        # get the folder names
        f5 = h5py.File(self.hdf5, 'a')
        fnames = f5.keys()

        # get the non rotated ones
        fnames_original = list(
            filter(lambda x: not re.search(r'_r\d+$', x), fnames))

        # get the rotated ones
        fnames_augmented = list(
            filter(lambda x: re.search(r'_r\d+$', x), fnames))

        # check feature_error
        if not self.feature_error:
            self.feature_error = []

        # computes the features of the original
        desc = '{:25s}'.format('Add features')
        for cplx_name in tqdm(
                fnames_original,
                desc=desc,
                ncols=100,
                disable=not prog_bar):

            # variant group
            variant_group = f5[cplx_name]

            variant = hdf5data.load_variant(variant_group)

            error_flag = False
            if self.compute_features is not None:

                # the internal features
                variant_group.require_group('features')
                variant_group.require_group('features_raw')

                error_flag = self._compute_features(self.compute_features,
                                                    self.environment,
                                                    distance_cutoff,
                                                    variant_group['features'],
                                                    variant,
                                                    self.logger)

                if error_flag:
                    self.feature_error += [cplx_name]

        # copy the data from the original to the augmented
        for cplx_name in fnames_augmented:

            # group of the variant
            aug_variant_group = f5[cplx_name]

            # get the source group
            variant_name = re.split(r'_r\d+', variant_group.name)[0]
            src_variant_group = f5[variant_name]

            # get the rotation parameters
            axis = aug_variant_group.attrs['axis']
            angle = aug_variant_group.attrs['angle']
            center = aug_variant_group.attrs['center']

            # copy the features to the augmented
            for k in variant_group['features']:
                if k not in aug_variant_group['features']:

                    # copy
                    data = src_variant_group['features/' + k][()]
                    aug_variant_group.require_group('features')

                    aug_variant_group.create_dataset("features/" + k, data=data, compression='lzf', chunks=True)

                    # rotate
                    self._rotate_feature(aug_variant_group, axis, angle, center, feat_name=[k])

        # find errored augmented variants
        tmp_aug_error = []
        for variant_name in self.feature_error:
            tmp_aug_error += list(filter(lambda x: variant_name in x,
                                         fnames_augmented))
        self.feature_error += tmp_aug_error

        #  Remove errored variants
        if self.feature_error:
            if remove_error:
                for variant_name in self.feature_error:
                    del f5[variant_name]
                self.logger.info(
                    f'Molecules with errored features are removed:\n'
                    f'{self.feature_error}')
            else:
                self.logger.warning(
                    f"The following variants have errored features:\n"
                    f'{self.feature_error}')

        # close the file
        f5.close()

# ====================================================================================
#
#       ADD TARGETS TO AN EXISTING DATASET
#
# ====================================================================================

    def add_unique_target(self, targdict):
        """Add identical targets for all the complexes in the datafile.

        This is usefull if you want to add the binary class of all the complexes
        created from decoys or natives

        Args:
            targdict (dict): Example: {'DOCKQ':1.0}

        >>> database = DataGenerator(hdf5='1ak4.hdf5')
        >>> database.add_unique_target({'DOCKQ':1.0})
        """

        # check if file exists
        if not os.path.isfile(self.hdf5):
            raise FileNotFoundError(
                'File %s does not exists' % self.hdf5)

        f5 = h5py.File(self.hdf5, 'a')
        for variant_name in list(f5.keys()):
            targrp = f5[variant_name].require_group('targets')
            for name, value in targdict.items():
                targrp.create_dataset(name, data=np.array([value]), chunks=True)
        f5.close()

    def add_target(self, prog_bar=False):
        """Add a target to an existing hdf5 file.

        Args:
            prog_bar (bool, optional): Use tqdm

        Example:

        >>> h5file = '1ak4.hdf5'
        >>>
        >>> #init the data assembler
        >>> database = DataGenerator(compute_targets =['deeprank.targets.binary_class'],
        >>>                          hdf5=h5file)
        >>>
        >>> database.add_target(prog_bar=True)
        """

        # check if file exists
        if not os.path.isfile(self.hdf5):
            raise FileNotFoundError(
                'File %s does not exists' % self.hdf5)

        # name of the hdf5 file
        f5 = h5py.File(self.hdf5, 'a')

        # get the folder names
        fnames = f5.keys()

        # get the non rotated ones
        fnames_original = list(
            filter(lambda x: not re.search(r'_r\d+$', x), fnames))
        fnames_augmented = list(
            filter(lambda x: re.search(r'_r\d+$', x), fnames))

        # compute the targets  of the original
        desc = '{:25s}'.format('Add targets')

        for variant_name in tqdm(fnames_original, desc=desc,
                              ncols=100, disable=not prog_bar):

            # group of the variant
            variant_group = f5[variant_name]

            # add the targets
            if self.compute_targets is not None:

                variant = hdf5data.load_variant(f5[variant_name])

                variant_group.require_group('targets')
                self._compute_targets(self.compute_targets,
                                      variant,
                                      variant_group['targets'])

        # copy the targets of the original to the rotated
        for cplx_name in fnames_augmented:

            # group of the variant
            aug_variant_group = f5[cplx_name]

            # get the source group
            variant_name = re.split(r'_r\d+', variant_group.name)[0]
            src_variant_group = f5[variant_name]

            # copy the targets to the augmented
            for k in variant_group['targets']:
                if k not in aug_variant_group['targets']:
                    data = src_variant_group['targets/' + k][()]
                    aug_variant_group.require_group('targets')
                    aug_variant_group.create_dataset("targets/" + k, data=data, chunks=True)

        # close the file
        f5.close()

    def realign_complexes(self, align, compute_features=None, pssm_source=None, distance_cutoff=10.0):
        """Align all the complexes already present in the HDF5.

        Arguments:
            align {dict} -- alignement dictionary (see __init__)

        Keyword Arguments:
            compute_features {list} -- list of features to be computed
                                       if None computes the features specified in
                                       the attrs['features'] of the file (if present)
            pssm_source {str} -- path of the pssm files. If None the source specfied in
                                 the attrs['pssm_source'] will be used (if present) (default: {None})
            distance_cutoff {float} -- max distance from the center to include atoms (default 10.0 Å)

        Raises:
            ValueError: If no PSSM detected

        Example:

        >>> database = DataGenerator(hdf5='1ak4.hdf5')
        >>> # if comute_features and pssm_source are not specified
        >>> # the values in hdf5.attrs['features'] and hdf5.attrs['pssm_source'] will be used
        >>> database.realign_complex(align={'axis':'x'},
        >>>                          compute_features['deeprank.features.X'],
        >>>                           pssm_source='./1ak4_pssm/')
        """

        f5 = h5py.File(self.hdf5, 'a')

        variant_names = f5.keys()
        self.logger.info(
            f'\n# Start aligning the HDF5 database: {self.hdf5}')

        # deal with the features
        if self.compute_features is None:
            if compute_features is None:
                if 'features' in f5.attrs:
                    self.compute_features = list(f5.attrs['features'])
            else:
                self.compute_features = compute_features

        # deal with the pssm source
        if self.pssm_source is not None:
            config.PATH_PSSM_SOURCE = self.pssm_source

        elif pssm_source is not None:
            config.PATH_PSSM_SOURCE = pssm_source

        elif 'pssm_source' in f5.attrs:
            config.PATH_PSSM_SOURCE = f5.attrs['pssm_source']
        else:
            raise ValueError('No pssm source detected')

        # loop over the complexes
        desc = '{:25s}'.format('Add features')
        for variant_name in tqdm(variant_names, desc=desc, ncols=100):

            variant = hdf5data.load_variant(f5[variant_name])

            # align the pdb
            variant_group = f5[variant_name]
            pdb = variant_group['complex'][()]

            sqldb = self._get_aligned_sqldb(pdb, align)
            data = sqldb.sql2pdb()

            data = np.array(data).astype('|S78')
            variant_group['complex'][...] = data

            # remove prexisting features
            old_dir = ['features', 'features_raw', 'mapped_features']
            for od in old_dir:
                if od in variant_group:
                    del variant_group[od]

            # the internal features
            variant_group.require_group('features')
            variant_group.require_group('features_raw')

            # compute features
            error_flag = self._compute_features(self.compute_features,
                                                self.environment,
                                                distance_cutoff,
                                                variant_group['features'],
                                                variant,
                                                self.logger)

        f5.close()

# ====================================================================================
#
#       PRECOMPUTE THE GRID POINTS
#
# ====================================================================================

    @staticmethod
    def get_grid_center(environment, variant):
        "gets the C-alpha position of the variant residue"

        pdb_path = get_pdb_path(environment.pdb_root, variant.pdb_ac)

        pdb = pdb2sql.pdb2sql(pdb_path)
        try:
            c_alpha_positions = pdb.get_xyz(chainID=variant.chain_id,
                                            resSeq=variant.residue_number,
                                            name="CA")
        finally:
            pdb._close()

        if len(c_alpha_positions) == 0:
            raise ValueError("C-alpha of chain {} residue {} not found in {}"
                             .format(variant.chain_id, variant.residue_number, pdb_path))

        return c_alpha_positions[0]

    def precompute_grid(self,
                        grid_info,
                        contact_distance=10.0,
                        prog_bar=False,
                        time=False,
                        try_sparse=True):

        # name of the hdf5 file
        f5 = h5py.File(self.hdf5, 'a')

        # check all the input PDB files
        variant_names = f5.keys()

        # get the local progress bar
        desc = '{:25s}'.format('Precompute grid points')
        variant_tqdm = tqdm(variant_names, desc=desc, disable=not prog_bar)

        if not prog_bar:
            print(desc, ':', self.hdf5)
            sys.stdout.flush()

        # loop over the data files
        for variant_name in variant_tqdm:

            variant_tqdm.set_postfix(variant=variant_name)

            variant = hdf5data.load_variant(f5[variant_name])

            # compute the data we want on the grid
            gt.GridTools(environment=self.environment,
                         variant_group=f5[variant_name], variant=variant,
                         grid_center=DataGenerator.get_grid_center(self.environment, variant),
                         number_of_points=grid_info['number_of_points'],
                         resolution=grid_info['resolution'],
                         contact_distance=contact_distance,
                         time=time,
                         prog_bar=prog_bar,
                         try_sparse=try_sparse)

        f5.close()

# ====================================================================================
#
#       MAP THE FEATURES TO THE GRID
#
# ====================================================================================

    def map_features(self, grid_info={},
                     cuda=False, gpu_block=None,
                     cuda_kernel='kernel_map.c',
                     cuda_func_name='gaussian',
                     try_sparse=True,
                     reset=False, use_tmpdir=False,
                     time=False,
                     prog_bar=True, grid_prog_bar=False,
                     remove_error=True):
        """Map the feature on a grid of points centered at the interface.

        If features to map are not given, they will be are automatically
        determined for each variant. Otherwise, given features will be mapped
        for all variants (i.e. existing mapped features will be recalculated).

        Args:
            grid_info (dict): Informaton for the grid.
                See deeprank.generate.GridTools.py for details.
            cuda (bool, optional): Use CUDA
            gpu_block (None, optional): GPU block size to be used
            cuda_kernel (str, optional): filename containing CUDA kernel
            cuda_func_name (str, optional): The name of the function in the kernel
            try_sparse (bool, optional): Try to save the grids as sparse format
            reset (bool, optional): remove grids if some are already present
            use_tmpdir (bool, optional): use a scratch directory
            time (bool, optional): time the mapping process
            prog_bar (bool, optional): use tqdm for each variant
            grid_prog_bar (bool, optional): use tqdm for each grid
            remove_error (bool, optional): remove the data that errored

        Example:

        >>> #init the data assembler
        >>> database = DataGenerator(hdf5='1ak4.hdf5')
        >>>
        >>> # map the features
        >>> grid_info = {
        >>>     'number_of_points': [20,20,20],
        >>>     'resolution': [1.,1.,1.],
        >>>     'atomic_densities': {'C':1.7, 'N':1.55, 'O':1.52, 'S':1.8},
        >>> }
        >>>
        >>> database.map_features(grid_info,try_sparse=True,time=False,prog_bar=True)
        """

        logger.debug("preprocess, map features with grid_info={}".format(grid_info))
        logger.debug("preprocess, map features with try_sparse={}".format(try_sparse))

        # default CUDA
        cuda_func = None
        cuda_atomic = None

        # disable CUDA when using MPI
        if self.mpi_comm is not None:
            if self.mpi_comm.Get_size() > 1:
                if cuda:
                    self.logger.warning(
                        'CUDA mapping disabled when using MPI')
                    cuda = False

        # check all the input PDB files
        with h5py.File(self.hdf5, 'r') as f5:
            variant_names = list(f5.keys())

        if len(variant_names) == 0:
            raise ValueError(f'No variants found in {self.hdf5}.')

        ################################################################
        # Check grid_info
        ################################################################
        # fills in the grid data if not provided: default = NONE
        grid_info_ref = copy.deepcopy(grid_info)
        grinfo = ['number_of_points', 'resolution']
        for gr in grinfo:
            if gr not in grid_info:
                grid_info[gr] = None

        # by default we do not map atomic densities
        if 'atomic_densities' not in grid_info:
            grid_info['atomic_densities'] = None

        # fills in the features mode if somes are missing: default = IND
        modes = ['atomic_densities_mode', 'feature_mode']
        for m in modes:
            if m not in grid_info:
                grid_info[m] = 'ind'

        ################################################################
        #
        ################################################################
        # sanity check for cuda
        if cuda and gpu_block is None:  # pragma: no cover
            self.logger.info(
                f'GPU block automatically set to 8 x 8 x 8. '
                f'You can set block size with gpu_block=[n,m,k]')
            gpu_block = [8, 8, 8]

        # initialize cuda
        if cuda:  # pragma: no cover

            # compile cuda module
            npts = grid_info['number_of_points']
            res = grid_info['resolution']
            module = self._compile_cuda_kernel(cuda_kernel, npts, res)

            # get the cuda function for the atomic/residue feature
            cuda_func = self._get_cuda_function(
                module, cuda_func_name)

            # get the cuda function for the atomic densties
            cuda_atomic_name = 'atomic_densities'
            cuda_atomic = self._get_cuda_function(
                module, cuda_atomic_name)

        # get the local progress bar
        desc = '{:25s}'.format('Map Features')
        variant_tqdm = tqdm(variant_names, desc=desc, disable=not prog_bar)

        if not prog_bar:
            self.logger.info(f'{desc}: {self.hdf5}')

        # loop over the data files
        for variant_name in variant_tqdm:

            with h5py.File(self.hdf5, 'a') as f5:

                variant_tqdm.set_postfix(variant=variant_name)

                variant = hdf5data.load_variant(f5[variant_name])
                grid_center = self.get_grid_center(self.environment, variant)

                # Determine which feature to map
                # if feature not given, then determine it for each variant
                if 'feature' not in grid_info_ref:
                    # if we havent mapped anything yet or if we reset
                    if 'mapped_features' not in list(f5[variant_name].keys()) or reset:
                        grid_info['feature'] = list(
                            f5[variant_name + '/features'].keys())

                    # if we have already mapped stuff
                    elif 'mapped_features' in list(f5[variant_name].keys()):

                        # feature name
                        all_feat = list(f5[variant_name + '/features'].keys())

                        # feature already mapped
                        mapped_feat = list(
                            f5[variant_name + '/mapped_features/Feature_ind'].keys())

                        # we select only the feture that were not mapped yet
                        grid_info['feature'] = []
                        for feat_name in all_feat:
                            if not any(map(lambda x: x.startswith(feat_name + '_'),
                                           mapped_feat)):
                                grid_info['feature'].append(feat_name)

                try:
                    # compute the data we want on the grid
                    gt.GridTools(
                        environment=self.environment,
                        variant_group=f5[variant_name],
                        variant=variant,
                        grid_center=grid_center,
                        number_of_points=grid_info['number_of_points'],
                        resolution=grid_info['resolution'],
                        atomic_densities=grid_info['atomic_densities'],
                        atomic_densities_mode=grid_info['atomic_densities_mode'],
                        feature=grid_info['feature'],
                        feature_mode=grid_info['feature_mode'],
                        cuda=cuda,
                        gpu_block=gpu_block,
                        cuda_func=cuda_func,
                        cuda_atomic=cuda_atomic,
                        time=time,
                        prog_bar=grid_prog_bar,
                        try_sparse=try_sparse)

                except:
                    self.map_error.append(variant_name)
                    self.logger.exception("Error during the mapping of {}: {}".format(variant, traceback.format_exc()))

        # remove the variants with issues
        if self.map_error:
            if remove_error:
                with h5py.File(self.hdf5, 'a') as f5:
                    for variant_name in self.map_error:
                        del f5[variant_name]
                self.logger.warning(
                    f"Variants with errored feature mapping are removed:\n"
                    f"{self.map_error}")
            else:
                self.logger.warning(
                    f"The following variants have errored feature mapping:\n"
                    f"{self.map_error}")

# ====================================================================================
#
#       REMOVE DATA FROM THE DATA SET
#
# ====================================================================================

    def remove(self, feature=True, pdb=True, points=True, grid=False):
        """Remove data from the data set.

        Equivalent to the cleandata command line tool. Once the data has been
        removed from the file it is impossible to add new features/targets

        Args:
            feature (bool, optional): Remove the features
            pdb (bool, optional): Remove the pdbs
            points (bool, optional): remove teh grid points
            grid (bool, optional): remove the maps
        """

        self.logger.debug('Remove features')

        # name of the hdf5 file
        f5 = h5py.File(self.hdf5, 'a')

        # get the folder names
        variant_names = f5.keys()

        for variant_name in variant_names:

            variant_group = f5[variant_name]

            if feature and 'features' in variant_group:
                del variant_group['features']
                del variant_group['features_raw']
            if pdb and 'complex' in variant_group and 'native' in variant_group:
                del variant_group['complex']
                del variant_group['native']
            if points and 'grid_points' in variant_group:
                del variant_group['grid_points']
            if grid and 'mapped_features' in variant_group:
                del variant_group['mapped_features']

        f5.close()

        # reclaim the space
        os.system('h5repack %s _tmp.h5py' % self.hdf5)
        os.system('mv _tmp.h5py %s' % self.hdf5)

# ====================================================================================
#
#       Simply tune or test the kernel
#
# ====================================================================================

    def _tune_cuda_kernel(self, grid_info, cuda_kernel='kernel_map.c', func='gaussian'):  # pragma: no cover
        """Tune the CUDA kernel using the kernel tuner
        http://benvanwerkhoven.github.io/kernel_tuner/

        Args:
            grid_info (dict): information for the grid definition
            cuda_kernel (str, optional): file containing the kernel
            func (str, optional): function in the kernel to be used

        Raises:
            ValueError: If the tuner has not been used
        """

        try:
            from kernel_tuner import tune_kernel
        except BaseException:
            print(
                'Install the Kernel Tuner: \n \t\t pip install kernel_tuner')
            print('http://benvanwerkhoven.github.io/kernel_tuner/')

        # fills in the grid data if not provided: default = NONE
        grinfo = ['number_of_points', 'resolution']
        for gr in grinfo:
            if gr not in grid_info:
                raise ValueError(
                    '%s must be specified to tune the kernel')

        # define the grid
        nx, ny, nz = grid_info['number_of_points']
        dx, dy, dz = grid_info['resolution']
        lx, ly, lz = nx * dx, ny * dy, nz * dz

        x = np.linspace(0, lx, nx)
        y = np.linspace(0, ly, ny)
        z = np.linspace(0, lz, nz)

        # create the dictionary containing the tune parameters
        tune_params = OrderedDict()
        tune_params['block_size_x'] = [2, 4, 8, 16, 32]
        tune_params['block_size_y'] = [2, 4, 8, 16, 32]
        tune_params['block_size_z'] = [2, 4, 8, 16, 32]

        # define the final grid
        grid = np.zeros(grid_info['number_of_points'])

        # arguments of the CUDA function
        x0, y0, z0 = np.float32(0), np.float32(0), np.float32(0)
        alpha = np.float32(0)
        args = [alpha, x0, y0, z0, x, y, z, grid]

        # dimensionality
        problem_size = grid_info['number_of_points']

        # get the kernel
        kernel = os.path.dirname(
            os.path.abspath(__file__)) + '/' + cuda_kernel
        kernel_code_template = open(kernel, 'r').read()

        npts = grid_info['number_of_points']
        res = grid_info['resolution']
        kernel_code = kernel_code_template % {
            'nx': npts[0], 'ny': npts[1], 'nz': npts[2], 'RES': np.max(res)}
        tunable_kernel = self._tunable_kernel(kernel_code)

        # tune
        tune_kernel(func, tunable_kernel,
                    problem_size, args, tune_params)

# ====================================================================================
#
#       Simply test the kernel
#
# ====================================================================================

    def _test_cuda(self, grid_info, gpu_block=8, cuda_kernel='kernel_map.c', func='gaussian'):  # pragma: no cover
        """Test the CUDA kernel.

        Args:
            grid_info (dict): Information for the grid definition
            gpu_block (int, optional): GPU block size to be used
            cuda_kernel (str, optional): File containing the kernel
            func (str, optional): function in the kernel to be used

        Raises:
            ValueError: If the kernel has not been installed
        """

        from time import time

        # fills in the grid data if not provided: default = NONE
        grinfo = ['number_of_points', 'resolution']
        for gr in grinfo:
            if gr not in grid_info:
                raise ValueError(
                    '%s must be specified to tune the kernel')

        # get the cuda function
        npts = grid_info['number_of_points']
        res = grid_info['resolution']
        module = self._compile_cuda_kernel(cuda_kernel, npts, res)
        cuda_func = self._get_cuda_function(module, func)

        # define the grid
        nx, ny, nz = grid_info['number_of_points']
        dx, dy, dz = grid_info['resolution']
        lx, ly, lz = nx * dx, ny * dy, nz * dz

        # create the coordinate
        x = np.linspace(0, lx, nx)
        y = np.linspace(0, ly, ny)
        z = np.linspace(0, lz, nz)

        # book memp on the gpu
        x_gpu = gpuarray.to_gpu(x.astype(np.float32))
        y_gpu = gpuarray.to_gpu(y.astype(np.float32))
        z_gpu = gpuarray.to_gpu(z.astype(np.float32))
        grid_gpu = gpuarray.zeros(
            grid_info['number_of_points'], np.float32)

        #  make sure we have three block value
        if not isinstance(gpu_block, list):
            gpu_block = [gpu_block] * 3

        #  get the grid
        gpu_grid = [int(np.ceil(n / b))
                    for b, n in zip(gpu_block, grid_info['number_of_points'])]
        print('GPU BLOCK:', gpu_block)
        print('GPU GRID :', gpu_grid)

        xyz_center = np.random.rand(500, 3).astype(np.float32)
        alpha = np.float32(1)
        t0 = time()
        for xyz in xyz_center:
            x0, y0, z0 = xyz
            cuda_func(alpha, x0, y0, z0, x_gpu, y_gpu, z_gpu, grid_gpu,
                      block=tuple(gpu_block), grid=tuple(gpu_grid))

        print('Done in: %f ms' % ((time() - t0) * 1000))

# ====================================================================================
#
#       Routines needed to handle CUDA
#
# ====================================================================================

    @staticmethod
    def _compile_cuda_kernel(cuda_kernel, npts, res):  # pragma: no cover
        """Compile the cuda kernel.

        Args:
            cuda_kernel (str): filename
            npts (tuple(int)): number of grid points in each direction
            res (tuple(float)): resolution in each direction

        Returns:
            compiler.SourceModule: compiled kernel
        """
        # get the cuda kernel path
        kernel = os.path.dirname(
            os.path.abspath(__file__)) + '/' + cuda_kernel
        kernel_code_template = open(kernel, 'r').read()
        kernel_code = kernel_code_template % {
            'nx': npts[0], 'ny': npts[1], 'nz': npts[2], 'RES': np.max(res)}

        # compile the kernel
        mod = compiler.SourceModule(kernel_code)
        return mod

    @staticmethod
    def _get_cuda_function(module, func_name):  # pragma: no cover
        """Get a single function from the compiled kernel.

        Args:
            module (compiler.SourceModule): compiled kernel module
            func_name (str): Name of the funtion

        Returns:
            func: cuda function
        """
        cuda_func = module.get_function(func_name)
        return cuda_func

    # tranform the kernel to a tunable one
    @staticmethod
    def _tunable_kernel(kernel):  # pragma: no cover
        """Make a tunale kernel.

        Args:
            kernel (str): String of the kernel

        Returns:
            TYPE: tunable kernel
        """
        switch_name = {
            'blockDim.x': 'block_size_x',
            'blockDim.y': 'block_size_y',
            'blockDim.z': 'block_size_z'}
        for old, new in switch_name.items():
            kernel = kernel.replace(old, new)
        return kernel

# ====================================================================================
#
#       FILTER DATASET
#
# ===================================================================================

    def _filter_cplx(self):
        """Filter the name of the complexes."""

        # read the class ID
        with open(self.pdb_select) as f:
            pdb_name = f.readlines()
        pdb_name = [name.split()[0] + '.pdb' for name in pdb_name]

        # create the filters
        tmp_path = []
        for name in pdb_name:
            tmp_path += list(filter(lambda x: name in x,
                                    self.pdb_path))

        # update the pdb_path
        self.pdb_path = tmp_path

# ====================================================================================
#
#       FEATURES ROUTINES
#
# ====================================================================================

    @staticmethod
    def _compute_features(feat_list, environment, distance_cutoff, featgrp, variant, logger):
        """Compute the features.

        Args:
            feat_list (list(str)): list of function name, e.g.,
                ['deeprank.features.ResidueDensity',
                'deeprank.features.PSSM_IC']
            environment (Environment): the environment settings
            distance_cutoff (float): max distance from center to include atoms, default 10.0 Å
            featgrp (str): name of the group where to store the xyz feature
            variant (PdbVariantSelection): the selected variant
            logger (logger): name of logger object

        Return:
            bool: error happened or not
        """
        error_flag = False  # when False: success; when True: failed
        for feat in feat_list:
            try:
                feat_module = importlib.import_module(feat, package=None)
                feat_module.__compute_feature__(environment, distance_cutoff, featgrp, variant)

                for feature_key in featgrp:
                    if np.any(np.isnan(featgrp[feature_key][()])):
                        logger.exception("Got NaN output for feature {} for {}".format(feature_key, variant))
                        error_flag = True
                        break
            except:
                logger.exception("Error while computing {} for {}: {}".format(feat, variant, traceback.format_exc()))
                error_flag = True

        return error_flag

# ====================================================================================
#
#       TARGETS ROUTINES
#
# ====================================================================================

    @staticmethod
    def _compute_targets(targ_list, variant, targrp):
        """Compute the targets.

        Args:
            targ_list (list(str)): list of function name
            pdb_data (bytes): PDB translated in btes
            targrp (str): name of the group where to store the targets
            logger (logger): name of logger object
        """
        for targ in targ_list:
            targ_module = importlib.import_module(targ, package=None)
            targ_module.__compute_target__(variant, targrp)

# ====================================================================================
#
#       ADD PDB FILE
#
# ====================================================================================

    def _add_pdb(self, variant_group, pdbfile, name):
        """Add a pdb to a variant group

        Args:
            variant_group (str): variant group where to add the pdb
            pdbfile (str): psb file to add
            name (str): dataset name in the hdf5 variant group
        """

        # no alignement
        if self.align is None:
            # read the pdb and extract the ATOM lines
            with open(pdbfile, 'r') as fi:
                data = [line.split('\n')[0]
                        for line in fi if line.startswith('ATOM')]

        # some alignement
        elif isinstance(self.align, dict):

            sqldb = self._get_aligned_sqldb(pdbfile, self.align)
            data = sqldb.sql2pdb()

        #  PDB default line length is 80
        #  http://www.wwpdb.org/documentation/file-format
        data = np.array(data).astype('|S78')
        variant_group.create_dataset(name, data=data, compression='lzf', chunks=True)

    # @staticmethod
    def _get_aligned_sqldb(self, pdbfile, dict_align):
        """return a sqldb of the pdb that is aligned as specified in the dict

        Arguments:
            pdbfile {str} -- path ot the pdb
            dict_align {dict} -- dictionanry of options to align the pdb
        """
        if 'selection' not in dict_align.keys():
            dict_align['selection'] = {}

        if 'export' not in dict_align.keys():
            dict_align['export'] = False

        if dict_align['selection'] == 'interface':
            raise ValueError("interface alignment is not supported in this version of deeprank")
        else:

            sqldb = align_along_axis(pdbfile, axis=dict_align['axis'],
                                     export=dict_align['export'],
                                     **dict_align['selection'])

        return sqldb

# ====================================================================================
#
#       AUGMENTED DATA
#
# ====================================================================================

    @staticmethod
    def _get_aligned_rotation_axis_angle(random_seed, dict_align):
        """Returns the axis and angle of rotation for data
           augmentation with aligned complexes

        Arguments:
            random_seed {int} -- random seed of rotation
            dict_align {dict} -- the dict describing the alignement

        Returns:
            list(float): axis of rotation
            float: angle of rotation
        """

        if random_seed is not None:
            np.random.seed(random_seed)

        angle = 2 * np.pi * np.random.rand()

        if 'plane' in dict_align.keys():
            if dict_align['plane'] == 'xy':
                axis = [0., 0., 1.]
            elif dict_align['plane'] == 'xz':
                axis = [0., 1., 0.]
            elif dict_align['plane'] == 'yz':
                axis = [1., 0., 0.]
            else:
                raise ValueError("plane must be xy, xz or yz")

        elif 'axis' in dict_align.keys():
            if dict_align['axis'] == 'x':
                axis = [1., 0., 0.]
            elif dict_align['axis'] == 'y':
                axis = [0., 1., 0.]
            elif dict_align['axis'] == 'z':
                axis = [0., 0., 1.]
            else:
                raise ValueError("axis must be x, y or z")
        else:
            raise ValueError('dict_align must contains plane or axis')

        return axis, angle

    # add a rotated pdb structure to the database
    def _add_aug_pdb(self, variant_group, variant, name, axis, angle):
        """Add augmented pdbs to the dataset.

        Args:
            variant_group (h5py.Group): name of the variant group
            pdbfile (str): pdb file name
            name (str): name of the dataset
            axis (list(float)): axis of rotation
            angle (float): angle of rotation
            dict_align (dict): dict for alignement of the original pdb

        Returns:
            list(float): center of the variant
        """

        pdb_path = get_pdb_path(self.environment.pdb_root, variant.pdb_ac)

        # create the sqldb and extract positions
        if self.align is None:
            sqldb = pdb2sql.pdb2sql(pdb_path)
        else:
            sqldb = self._get_aligned_sqldb(pdb_path, self.align)

        # rotate the positions
        pdb2sql.transform.rot_axis(sqldb, axis, angle)

        # get molecule center
        center = self.get_grid_center(self.environment, variant)

        # get the pdb-format data
        data = sqldb.sql2pdb()
        data = np.array(data).astype('|S78')

        variant_group.create_dataset(name, data=data, compression='lzf', chunks=True)

        # close the db
        sqldb._close()

        return center

    # rotate th xyz-formatted feature in the database

    @staticmethod
    def _rotate_feature(variant_group, axis, angle, center, feat_name='all'):
        """Rotate the raw feature values.

        Args:
            variant group (str): name pf the variant group
            axis (list(float)): axis of rotation
            angle (float): angle of rotation
            center (list(float)): center of rotation
            feat_name (str): name of the feature to rotate or 'all'
        """
        if feat_name == 'all':
            feat = list(variant_group['features'].keys())
        else:
            feat = feat_name
            if not isinstance(feat, list):
                feat = list(feat)

        for fn in feat:

            # extract the data
            data = variant_group['features/' + fn][()]

            # if data not empty
            if data.shape[0] != 0:

                # xyz
                xyz = data[:, :DataGenerator.FEATURE_POSITION_OFFSET]

                # get rotated xyz
                xyz_rot = pdb2sql.transform.rot_xyz_around_axis(
                    xyz, axis, angle, center)

                # put back the data
                variant_group['features/' + fn][:, :DataGenerator.FEATURE_POSITION_OFFSET] = xyz_rot

{'tune_kernel': 'kernel_tuner.tune_kernel', 'time': 'time.time'}.FEATURE_POSITION_OFFSET


import math, json
import bpy, mathutils

# ExportHelper is a helper class, defines filename and
# invoke() function which calls the file selector.
from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty, BoolProperty, EnumProperty

class TransverseMercator:
    radius = 6378137

    def __init__(self, **kwargs):
        # setting default values
        self.lat = 0 # in degrees
        self.lon = 0 # in degrees
        self.k = 1 # scale factor

        for attr in kwargs:
            setattr(self, attr, kwargs[attr])
        self.latInRadians = math.radians(self.lat)

    def fromGeographic(self, lat, lon):
        lat = math.radians(lat)
        lon = math.radians(lon-self.lon)
        B = math.sin(lon) * math.cos(lat)
        x = 0.5 * self.k * self.radius * math.log((1+B)/(1-B))
        y = self.k * self.radius * ( math.atan(math.tan(lat)/math.cos(lon)) - self.latInRadians )
        return (x,y)

    def toGeographic(self, x, y):
        x = x/(self.k * self.radius)
        y = y/(self.k * self.radius)
        D = y + self.latInRadians
        lon = math.atan(math.sinh(x)/math.cos(D))
        lat = math.asin(math.sin(D)/math.cosh(x))

        lon = self.lon + math.degrees(lon)
        lat = math.degrees(lat)
        return (lat, lon)

class GeoJson(bpy.types.Operator, ExportHelper):
    bl_idname = "prk.export_geojson"  # important since its how bpy.ops.import_test.some_data is constructed
    bl_label = "Export GeoJSON"

    # ExportHelper mixin class uses this
    filename_ext = ".json"

    filter_glob = StringProperty(default="*.json", options={'HIDDEN'})

    # List of operator properties, the attributes will be assigned
    # to the class instance from the operator settings before calling.
    use_setting = BoolProperty(
            name="Example Boolean",
            description="Example Tooltip",
            default=True,
            )

    type = EnumProperty(
            name="Example Enum",
            description="Choose between two items",
            items=(('OPT_A', "First Option", "Description one"),
                   ('OPT_B', "Second Option", "Description two")),
            default='OPT_A',
            )

    def execute(self, context):
        features = []
        data = {
            "type": "FeatureCollection",
            "features": features
        }
        rotationMatrix = mathutils.Matrix.Rotation(math.radians(context.scene["heading"]), 4, "Z")
        projection = TransverseMercator(lat=context.scene["latitude"], lon=context.scene["longitude"])
        # iterate through all rooms
        for o in context.scene.objects:
            if not ("t" in o and o["t"] == "room"):
                continue
            coords = []
            # iterate through EMPTYs that control the vertices of polygon of the area
            for m in o.modifiers:
                e = m.object
                p = rotationMatrix * e.parent.matrix_world * e.location
                p = projection.toGeographic(p.x, p.y)
                coords.append( (p[1], p[0]) )
            features.append({
                "type": "Feature",
                "geometry": {
                    "type": "Polygon",
                    "coordinates": (coords,)
                },
                "properties": {
                    "name": o.name
                }
            })
        f = open(self.filepath, 'w', encoding="utf-8")
        f.write(json.dumps(data))
        f.close()

        return {'FINISHED'}


import sys
sys.path.insert(0, './src')
# import os
# os.getcwd()
# os.chdir('./DriverActionEstimators')
from importlib import reload
import numpy as np
np.set_printoptions(suppress=True)
import pickle
import os

def pickle_this(item, data_files_dir, item_name):
    data_files_dir += '/'+item_name+'.pickle'
    if not os.path.exists(data_files_dir):
        with open(data_files_dir, 'wb') as handle:
            pickle.dump(item, handle)
    else:
        print('This data id exists')

def prep_data(training_data):
    all_epis = np.unique(training_data[0][:, 0, 0])
    np.random.seed(2021)
    np.random.shuffle(all_epis)

    train_epis = all_epis[:int(len(all_epis)*0.7)]
    val_epis = np.setdiff1d(all_epis, train_epis)
    train_samples = np.where(training_data[0][:, 0:1, 0] == train_epis)[0]
    val_samples = np.where(training_data[0][:, 0:1, 0] == val_epis)[0]

    _, history_sca, future_idm_s, future_m_veh_c, history_e_veh_a = training_data

    train_input = [history_sca[train_samples, :, 2:],
                history_e_veh_a[train_samples, :, 2:3]]

    test_input = [history_sca[val_samples, :, 2:],
                history_e_veh_a[val_samples, :, 2:3]]

    return train_input, test_input
# %%
"""
Load generated data.
"""
data_id = '049'
dataset_name = 'sim_data_'+data_id
data_files_dir = './src/datasets/'+dataset_name
with open(data_files_dir+'/sim_data.pickle', 'rb') as handle:
    sim_data = pickle.load(handle)
sim_data.shape

# %%

plt.hist(size)
np.array(size)[:, 1].mean()
np.array(size)[:, 1].max()
np.array(size).max()
# %%
"""
Prepare data
"""
history_len = 30 # steps
rollout_len = 50
from data.latent_mlp import data_prep
reload(data_prep)
from data.latent_mlp.data_prep import DataPrepLatentMLP
data_prep = DataPrepLatentMLP()
sim_data, dummy_value_set = data_prep.fill_missing_values(sim_data)
sim_data_scaled, env_scaler, m_scaler = data_prep.scale_data(sim_data)
history_future_seqs = data_prep.sequence(sim_data, history_len, rollout_len)
history_future_seqs_scaled = data_prep.sequence(sim_data_scaled, history_len, rollout_len)
data_arrays = data_prep.split_data(history_future_seqs, history_future_seqs_scaled)

history_future_usc, history_sca, \
                future_idm_s, future_m_veh_c, history_e_veh_a = data_arrays
# %%
train_input, test_input = prep_data(data_arrays)
train_input[0].shape
train_input[1].shape
# %%
data_files_dir += '/latent_mlp'
if not os.path.exists(data_files_dir):
    os.makedirs(data_files_dir)

item_name = 'data_arrays_h{history_len}_f{rollout_len}'.format(\
                                history_len=history_len, rollout_len=rollout_len)

pickle_this(data_arrays, data_files_dir, item_name)
# (841,502,560 bytes)

item_name = 'train_input{history_len}_f{rollout_len}'.format(\
                                history_len=history_len, rollout_len=rollout_len)

pickle_this(train_input, data_files_dir, item_name)

item_name = 'test_input{history_len}_f{rollout_len}'.format(\
                                history_len=history_len, rollout_len=rollout_len)

pickle_this(test_input, data_files_dir, item_name)


# var.py
#
# Read the var files.
# NB: the f array returned is C-ordered: f[nvar, nz, ny, nx]
#     NOT Fortran as in Pencil (& IDL):  f[nx, ny, nz, nvar]
#
# Authors:
# J. Oishi (joishi@amnh.org)
# T. Gastine (tgastine@ast.obs-mip.fr)
# S. Candelaresi (iomsn1@gmail.com).
"""
Contains the read class for the VAR file reading,
some simulation attributes and the data cube.
"""

def var(*args, **kwargs):
    """
    Read VAR files from Pencil Code. If proc < 0, then load all data
    and assemble, otherwise load VAR file from specified processor.

    The file format written by output() (and used, e.g. in var.dat)
    consists of the followinig Fortran records:
    1. data(mx, my, mz, nvar)
    2. t(1), x(mx), y(my), z(mz), dx(1), dy(1), dz(1), deltay(1)
    Here nvar denotes the number of slots, i.e. 1 for one scalar field, 3
    for one vector field, 8 for var.dat in the case of MHD with entropy.
    but, deltay(1) is only there if lshear is on! need to know parameters.

    call signature:

    var(var_file='', datadir='data', proc=-1, ivar=-1, quiet=True,
        trimall=False, magic=None, sim=None, precision='f')

    Keyword arguments:
        var_file:   Name of the VAR file.
        datadir:    Directory where the data is stored.
        proc:       Processor to be read. If -1 read all and assemble to one array.
        ivar:       Index of the VAR file, if var_file is not specified.
        quiet:      Flag for switching off output.
        trimall:    Trim the data cube to exclude ghost zones.
        magic:      Values to be computed from the data, e.g. B = curl(A).
        sim:        Simulation sim object.
        precision:  Float (f) or double (d).
    """

    from ..sim import __Simulation__

    started = None

    for a in args:
        if isinstance(a, __Simulation__):
            started = a.started()
            break

    if 'sim' in kwargs.keys():
        started = kwargs['sim'].started()
    elif 'datadir' in kwargs.keys():
        from os.path import join, exists
        if exists(join(kwargs['datadir'], 'time_series.dat')):
            started = True
    else:
        from os.path import join, exists
        if exists(join('data', 'time_series.dat')):
            started = True

        #else:
        #    print('!! ERROR: No simulation of path specified..')

    if not started:
        print('ERROR: Simulation has not jet started. There are no var files.')
        return False

    var_tmp = DataCube()
    var_tmp.read(*args, **kwargs)
    return var_tmp

class DataCube(object):
    """
    DataCube -- holds Pencil Code VAR file data.
    """

    def __init__(self):
        """
        Fill members with default values.
        """

        self.t = 0.
        self.dx = 1.
        self.dy = 1.
        self.dz = 1.

    def read(self, var_file='', datadir='data', proc=-1, ivar=-1, quiet=True,
             trimall=False, magic=None, sim=None, precision='f'):
        """
        Read VAR files from Pencil Code. If proc < 0, then load all data
        and assemble, otherwise load VAR file from specified processor.

        The file format written by output() (and used, e.g. in var.dat)
        consists of the followinig Fortran records:
        1. data(mx, my, mz, nvar)
        2. t(1), x(mx), y(my), z(mz), dx(1), dy(1), dz(1), deltay(1)
        Here nvar denotes the number of slots, i.e. 1 for one scalar field, 3
        for one vector field, 8 for var.dat in the case of MHD with entropy.
        but, deltay(1) is only there if lshear is on! need to know parameters.

        call signature:

        var(var_file='', datadir='data', proc=-1, ivar=-1, quiet=True,
            trimall=False, magic=None, sim=None, precision='f')

        Keyword arguments:
            var_file:   Name of the VAR file.
            datadir:    Directory where the data is stored.
            proc:       Processor to be read. If -1 read all and assemble to one array.
            ivar:       Index of the VAR file, if var_file is not specified.
            quiet:      Flag for switching off output.
            trimall:    Trim the data cube to exclude ghost zones.
            magic:      Values to be computed from the data, e.g. B = curl(A).
            sim:        Simulation sim object.
            precision:  Float (f) or double (d).
        """

        import numpy as np
        import os
        from scipy.io import FortranFile
        from ..math.derivatives import curl, curl2
        from .. import read
        from ..sim import __Simulation__

        dim = None
        param = None
        index = None

        if isinstance(sim, __Simulation__):
            datadir = os.path.expanduser(sim.datadir)
            dim = sim.dim
            param = read.param(datadir=sim.datadir, quiet=True)
            index = read.index(datadir=sim.datadir)
        else:
            datadir = os.path.expanduser(datadir)
            if dim is None:
                if var_file[0:2].lower() == 'og':
                    dim = read.ogdim(datadir, proc)
                else:
                    dim = read.dim(datadir, proc)
            if param is None:
                param = read.param(datadir=datadir, quiet=quiet)
            if index is None:
                index = read.index(datadir=datadir)

        run2D = param.lwrite_2d

        if dim.precision == 'D':
            precision = 'd'
        else:
            precision = 'f'

        if param.lwrite_aux:
            total_vars = dim.mvar + dim.maux
        else:
            total_vars = dim.mvar

        if not var_file:
            if ivar < 0:
                var_file = 'var.dat'
            else:
                var_file = 'VAR' + str(ivar)

        if proc < 0:
            proc_dirs = self.__natural_sort(filter(lambda s: s.startswith('proc'),
                                                   os.listdir(datadir)))
        else:
            proc_dirs = ['proc' + str(proc)]

        # Set up the global array.
        if not run2D:
            f = np.zeros((total_vars, dim.mz, dim.my, dim.mx),
                         dtype=precision)
        else:
            if dim.ny == 1:
                f = np.zeros((total_vars, dim.mz, dim.mx), dtype=precision)
            else:
                f = np.zeros((total_vars, dim.my, dim.mx), dtype=precision)

        x = np.zeros(dim.mx, dtype=precision)
        y = np.zeros(dim.my, dtype=precision)
        z = np.zeros(dim.mz, dtype=precision)

        for directory in proc_dirs:
            proc = int(directory[4:])
            if var_file[0:2].lower() == 'og':
                procdim = read.ogdim(datadir, proc)
            else:
                procdim = read.dim(datadir, proc)
            if not quiet:
                print("Reading data from processor {0} of {1} ...".format( \
                      proc, len(proc_dirs)))

            mxloc = procdim.mx
            myloc = procdim.my
            mzloc = procdim.mz

            # Read the data.
            file_name = os.path.join(datadir, directory, var_file)
            infile = FortranFile(file_name)
            if not run2D:
                f_loc = infile.read_record(dtype=precision)
                f_loc = f_loc.reshape((-1, mzloc, myloc, mxloc))
            else:
                if dim.ny == 1:
                    f_loc = infile.read_record(dtype=precision)
                    f_loc = f_loc.reshape((-1, mzloc, mxloc))
                else:
                    f_loc = infile.read_record(dtype=precision)
                    f_loc = f_loc.reshape((-1, myloc, mxloc))
            raw_etc = infile.read_record(precision)
            infile.close()

            t = raw_etc[0]
            x_loc = raw_etc[1:mxloc+1]
            y_loc = raw_etc[mxloc+1:mxloc+myloc+1]
            z_loc = raw_etc[mxloc+myloc+1:mxloc+myloc+mzloc+1]
            if param.lshear:
                shear_offset = 1
                deltay = raw_etc[-1]
            else:
                shear_offset = 0

            dx = raw_etc[-3-shear_offset]
            dy = raw_etc[-2-shear_offset]
            dz = raw_etc[-1-shear_offset]

            if len(proc_dirs) > 1:
                # Calculate where the local processor will go in
                # the global array.
                #
                # Don't overwrite ghost zones of processor to the left (and
                # accordingly in y and z direction -- makes a difference on the
                # diagonals)
                #
                # Recall that in NumPy, slicing is NON-INCLUSIVE on the right end
                # ie, x[0:4] will slice all of a 4-digit array, not produce
                # an error like in idl.

                if procdim.ipx == 0:
                    i0x = 0
                    i1x = i0x + procdim.mx
                    i0xloc = 0
                    i1xloc = procdim.mx
                else:
                    i0x = procdim.ipx*procdim.nx + procdim.nghostx
                    i1x = i0x + procdim.mx - procdim.nghostx
                    i0xloc = procdim.nghostx
                    i1xloc = procdim.mx

                if procdim.ipy == 0:
                    i0y = 0
                    i1y = i0y + procdim.my
                    i0yloc = 0
                    i1yloc = procdim.my
                else:
                    i0y = procdim.ipy*procdim.ny + procdim.nghosty
                    i1y = i0y + procdim.my - procdim.nghosty
                    i0yloc = procdim.nghosty
                    i1yloc = procdim.my

                if procdim.ipz == 0:
                    i0z = 0
                    i1z = i0z+procdim.mz
                    i0zloc = 0
                    i1zloc = procdim.mz
                else:
                    i0z = procdim.ipz*procdim.nz + procdim.nghostz
                    i1z = i0z + procdim.mz - procdim.nghostz
                    i0zloc = procdim.nghostz
                    i1zloc = procdim.mz

                x[i0x:i1x] = x_loc[i0xloc:i1xloc]
                y[i0y:i1y] = y_loc[i0yloc:i1yloc]
                z[i0z:i1z] = z_loc[i0zloc:i1zloc]

                if not run2D:
                    f[:, i0z:i1z, i0y:i1y, i0x:i1x] = \
                        f_loc[:, i0zloc:i1zloc, i0yloc:i1yloc, i0xloc:i1xloc]
                else:
                    if dim.ny == 1:
                        f[:, i0z:i1z, i0x:i1x] = \
                            f_loc[:, i0zloc:i1zloc, i0xloc:i1xloc]
                    else:
                        f[:, i0y:i1y, i0x:i1x] = \
                            f_loc[:, i0yloc:i1yloc, i0xloc:i1xloc]
            else:
                f = f_loc
                x = x_loc
                y = y_loc
                z = z_loc

        if magic is not None:
            if 'bb' in magic:
                # Compute the magnetic field before doing trimall.
                aa = f[index.ax-1:index.az, ...]
                # TODO: Specify coordinate system.
                self.bb = curl(aa, dx, dy, dz, run2D=run2D)
                if trimall:
                    self.bb = self.bb[:, dim.n1:dim.n2+1,
                                      dim.m1:dim.m2+1, dim.l1:dim.l2+1]
            if 'jj' in magic:
                # Compute the electric current field before doing trimall.
                aa = f[index.ax-1:index.az, ...]
                # TODO: Specify coordinate system.
                self.jj = curl2(aa, dx, dy, dz)
                if trimall:
                    self.jj = self.jj[:, dim.n1:dim.n2+1,
                                      dim.m1:dim.m2+1, dim.l1:dim.l2+1]
            if 'vort' in magic:
                # Compute the vorticity field before doing trimall.
                uu = f[index.ux-1:index.uz, ...]
                # TODO: Specify coordinate system.
                # WL: The curl subroutine should take care of it.
                self.vort = curl(uu, dx, dy, dz, run2D=run2D)
                if trimall:
                    if run2D:
                        if dim.nz == 1:
                            self.vort = self.vort[:, dim.m1:dim.m2+1,
                                                  dim.l1:dim.l2+1]
                        else:
                            self.vort = self.vort[:, dim.n1:dim.n2+1,
                                                  dim.l1:dim.l2+1]
                    else:
                        self.vort = self.vort[:, dim.n1:dim.n2+1,
                                              dim.m1:dim.m2+1,
                                              dim.l1:dim.l2+1]

        # Trim the ghost zones of the global f-array if asked.
        if trimall:
            self.x = x[dim.l1:dim.l2+1]
            self.y = y[dim.m1:dim.m2+1]
            self.z = z[dim.n1:dim.n2+1]
            if not run2D:
                self.f = f[:, dim.n1:dim.n2+1, dim.m1:dim.m2+1, dim.l1:dim.l2+1]
            else:
                if dim.ny == 1:
                    self.f = f[:, dim.n1:dim.n2+1, dim.l1:dim.l2+1]
                else:
                    self.f = f[:, dim.m1:dim.m2+1, dim.l1:dim.l2+1]
        else:
            self.x = x
            self.y = y
            self.z = z
            self.f = f
            self.l1 = dim.l1
            self.l2 = dim.l2 + 1
            self.m1 = dim.m1
            self.m2 = dim.m2 + 1
            self.n1 = dim.n1
            self.n2 = dim.n2 + 1

        # Assign an attribute to self for each variable defined in
        # 'data/index.pro' so that e.g. self.ux is the x-velocity
        for key in index.__dict__.keys():
            if key != 'global_gg' and key != 'keys':
                value = index.__dict__[key]
                setattr(self, key, self.f[value-1, ...])
        # Special treatment for vector quantities.
        if hasattr(index, 'uu'):
            self.uu = self.f[index.ux-1:index.uz, ...]
        if hasattr(index, 'aa'):
            self.aa = self.f[index.ax-1:index.az, ...]

        self.t = t
        self.dx = dx
        self.dy = dy
        self.dz = dz
        if param.lshear:
            self.deltay = deltay

        # Do the rest of magic after the trimall (i.e. no additional curl...).
        self.magic = magic
        if self.magic is not None:
            self.magic_attributes(param)

    def __natural_sort(self, procs_list):
        """
        Sort array in a more natural way, e.g. 9VAR < 10VAR
        """

        import re

        convert = lambda text: int(text) if text.isdigit() else text.lower()
        alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
        return sorted(procs_list, key=alphanum_key)

    def magic_attributes(self, param):
        """
        Compute some additional 'magic' quantities.
        """

        import numpy as np
        import sys

        for field in self.magic:
            if field == 'rho' and not hasattr(self, 'rho'):
                if hasattr(self, 'lnrho'):
                    setattr(self, 'rho', np.exp(self.lnrho))
                else:
                    sys.exit("pb in magic!")

            if field == 'tt' and not hasattr(self, 'tt'):
                if hasattr(self, 'lnTT'):
                    tt = np.exp(self.lnTT)
                    setattr(self, 'tt', tt)
                else:
                    if hasattr(self, 'ss'):
                        if hasattr(self, 'lnrho'):
                            lnrho = self.lnrho
                        elif hasattr(self, 'rho'):
                            lnrho = np.log(self.rho)
                        else:
                            sys.exit("pb in magic: missing rho or lnrho variable")
                        cp = param.cp
                        gamma = param.gamma
                        cs20 = param.cs0**2
                        lnrho0 = np.log(param.rho0)
                        lnTT0 = np.log(cs20/(cp*(gamma-1.)))
                        lnTT = lnTT0+gamma/cp*self.ss+(gamma-1.)* \
                               (lnrho-lnrho0)
                        setattr(self, 'tt', np.exp(lnTT))
                    else:
                        sys.exit("pb in magic!")

            if field == 'ss' and not hasattr(self, 'ss'):
                cp = param.cp
                gamma = param.gamma
                cs20 = param.cs0**2
                lnrho0 = np.log(param.rho0)
                lnTT0 = np.log(cs20/(cp*(gamma-1.)))
                if hasattr(self, 'lnTT'):
                    setattr(self, 'ss', cp/gamma*(self.lnTT-lnTT0- \
                            (gamma-1.)*(self.lnrho-lnrho0)))
                elif hasattr(self, 'tt'):
                    setattr(self, 'ss', cp/gamma*(np.log(self.tt)- \
                            lnTT0-(gamma-1.)*(self.lnrho-lnrho0)))
                else:
                    sys.exit("pb in magic!")

            if field == 'pp' and not hasattr(self, 'pp'):
                cp = param.cp
                gamma = param.gamma
                cv = cp/gamma
                if hasattr(self, 'lnrho'):
                    lnrho = self.lnrho
                elif hasattr(self, 'rho'):
                    lnrho = np.log(self.rho)
                else:
                    sys.exit("pb in magic: missing rho or lnrho variable")
                if hasattr(self, 'ss'):
                    setattr(self, 'pp', np.exp(gamma*(self.ss+lnrho)))
                elif hasattr(self, 'lntt'):
                    setattr(self, 'pp', (cp-cv)*np.exp(self.lntt+lnrho))
                elif hasattr(self, 'tt'):
                    setattr(self, 'pp', (cp-cv)*self.tt*np.exp(lnrho))
                else:
                    sys.exit("pb in magic!")


import jwt
from rest_framework import authentication, exceptions
from django.conf import settings
from django.contrib.auth.models import User
from jwt.algorithms import get_default_algorithms

class JWTAuthentication(authentication.BaseAuthentication):

    def authenticate(self, request):
        auth_data = authentication.get_authorization_header(request)

        if not auth_data:
            return None

        prefix, token = auth_data.decode('utf-8').split(' ')
        # print(settings.JWT_SECRET_KEY)
        # print(jwt.decode(token, settings.JWT_SECRET_KEY, algorithms=[
        #       "HS256"], options={"verify_signature": False}))
        try:
            payload = jwt.decode(
                token, settings.JWT_SECRET_KEY, algorithms=["HS256"], options={
                    "verify_signature": False})

            user = User.objects.get(username=payload['username'])
            print(user)
            return (user, token)

        except jwt.DecodeError as identifier:
            raise exceptions.AuthenticationFailed(
                'Your token is invalid, login')
        except jwt.ExpiredSignatureError as identifier:
            raise exceptions.AuthenticationFailed(
                'Your token is expired, login')

        return super().authenticate(request)