# -*- coding: utf-8 -*-
#
# Base class for all computational classes in Syncopy
#

# Builtin/3rd party package imports
import os
import sys
import psutil
import h5py
import time
import numpy as np
from abc import ABC, abstractmethod
from copy import copy
from glob import glob
from numpy.lib.format import open_memmap
from tqdm.auto import tqdm
if sys.platform == "win32":
    # tqdm breaks term colors on Windows - fix that (tqdm issue #446)
    import colorama
    colorama.deinit()
    colorama.init(strip=False)

# Local imports
from .tools import get_defaults
from syncopy import __storage__, __dask__, __path__
from syncopy.shared.errors import SPYIOError, SPYValueError, SPYParallelError, SPYWarning
if __dask__:
    import dask.distributed as dd
    import dask.bag as db
    # # In case of problems w/worker-stealing, uncomment the following lines
    # import dask
    # dask.config.set(distributed__scheduler__work_stealing=False)

__all__ = []

class ComputationalRoutine(ABC):
    """Abstract class for encapsulating sequential/parallel algorithms

    A Syncopy compute class consists of a
    :class:`ComputationalRoutine`-subclass that binds a static
    :func:`computeFunction` and provides the class method
    :meth:`process_metadata`.

    Requirements for :meth:`computeFunction`:

    * First positional argument is a :class:`numpy.ndarray`, the keywords
      `chunkShape` and `noCompute` are supported
    * Returns a :class:`numpy.ndarray` if `noCompute` is `False` and expected
      shape and numerical type of output array otherwise.

    Requirements for :class:`ComputationalRoutine`:

    * Child of :class:`ComputationalRoutine`, binds :func:`computeFunction`
      as static method
    * Provides class method :func:`process_metadata`

    For details on writing compute classes and metafunctions for Syncopy, please
    refer to :doc:`/developer/compute_kernels`.
    """

    # Placeholder: the actual workhorse
    @staticmethod
    def computeFunction(arr, *argv, chunkShape=None, noCompute=None, **kwargs):
        """Computational core routine

        Parameters
        ----------
        arr : :class:`numpy.ndarray`
           Numerical data from a single trial
        *argv : tuple
           Arbitrary tuple of positional arguments
        chunkShape : None or tuple
           Mandatory keyword. If not `None`, represents global block-size of
           processed trial.
        noCompute : None or bool
           Preprocessing flag. If `True`, do not perform actual calculation but
           instead return expected shape and :class:`numpy.dtype` of output
           array.
        **kwargs: dict
           Other keyword arguments.

        Returns
        -------
        out Shape : tuple, if ``noCompute == True``
           expected shape of output array
        outDtype : :class:`numpy.dtype`, if ``noCompute == True``
           expected numerical type of output array
        res : :class:`numpy.ndarray`, if ``noCompute == False``
           Result of processing input `arr`

        Notes
        -----
        This concrete method is a placeholder that is intended to be
        overloaded.

        See also
        --------
        ComputationalRoutine : Developer documentation: :doc:`/developer/compute_kernels`.
        """
        return None

    def __init__(self, *argv, **kwargs):
        """
        Instantiate a :class:`ComputationalRoutine` subclass

        Parameters
        ----------
        *argv : tuple
           Tuple of positional arguments passed on to :meth:`computeFunction`
        **kwargs : dict
           Keyword arguments passed on to :meth:`computeFunction`

        Returns
        -------
        obj : instance of :class:`ComputationalRoutine`-subclass
           Usable class instance for processing Syncopy data objects.
        """

        # list of positional arguments to `computeFunction` for all workers, format:
        # ``self.argv = [3, [0, 1, 1], ('a', 'b', 'c')]`` (compare to `self.ArgV` below)
        self.argv = list(argv)

        # list of positional keyword arguments split up for each worker w/format:
        # ``self.ArgV = [(3,0,'a'), (3,1,'b'), (3,1,'c')`` (compare `self.argv` above)
        self.ArgV = None

        # dict of default keyword values accepted by `computeFunction`
        self.defaultCfg = get_defaults(self.computeFunction)

        # dict of actual keyword argument values to `computeFunction` provided by user
        self.cfg = copy(self.defaultCfg)
        for key in set(self.cfg.keys()).intersection(kwargs.keys()):
            self.cfg[key] = kwargs[key]

        # binary flag: if `True`, average across trials, do nothing otherwise
        self.keeptrials = None

        # full shape of final output dataset (all trials, all chunks, etc.)
        self.outputShape = None

        # numerical type of output dataset
        self.dtype = None

        # list of dicts encoding header info of raw binary input files (experimental!)
        self.hdr = None

        # list of trial numbers to process (either `data.trials` or `data._selection.trials`)
        self.trialList = None

        # list of index-tuples for extracting trial-chunks from input HDF5 dataset
        # >>> MUST be ordered, no repetitions! <<<
        # indices are ABSOLUTE, i.e., wrt entire dataset, not just current trial!
        self.sourceLayout = None

        # list of index-tuples for re-ordering NumPy arrays extracted w/`self.sourceLayout`
        # >>> can be unordered w/repetitions <<<
        # indices are RELATIVE, i.e., wrt current trial!
        self.sourceSelectors = None

        # list of index-tuples for storing trial-chunk result in output dataset
        # >>> MUST be ordered, no repetitions! <<<
        # indices are ABSOLUTE, i.e., wrt entire dataset, not just current trial
        self.targetLayout = None

        # list of shape-tuples of trial-chunk results
        self.targetShapes = None

        # binary flag: if `True`, use fancy array indexing via `np.ix_` to extract
        # data from input via `self.sourceLayout` + `self.sourceSelectors`; if `False`,
        # only use `self.sourceLayout` (selections ordered, no reps)
        self.useFancyIdx = None

        # integer, max. memory footprint of largest input array piece (in bytes)
        self.chunkMem = None

        # directory for storing source-HDF5 files making up virtual output dataset
        self.virtualDatasetDir = None

        # h5py layout encoding shape/geometry of file sources within virtual output dataset
        self.VirtualDatasetLayout = None

        # name of output dataset
        self.datasetName = None

        # tmp holding var for preserving original access mode of `data`
        self.dataMode = None

        # time (in seconds) b/w querying state of futures ('pending' -> 'finished')
        self.sleepTime = 0.1

        # if `True`, enforces use of single-threaded scheduler in `compute_parallel`
        self.parallelDebug = False

        # format string for tqdm progress bars in sequential and parallel computations
        self.tqdmFormat = "{desc}: {percentage:3.0f}% |{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]"

        # maximal acceptable size (in MB) of any provided positional argument
        self._maxArgSize = 100

        # counter and maximal recursion depth for calling `self._sizeof`
        self._callMax = 10000
        self._callCount = 0

    def initialize(self, data, chan_per_worker=None, keeptrials=True):
        """
        Perform dry-run of calculation to determine output shape

        Parameters
        ----------
        data : syncopy data object
           Syncopy data object to be processed (has to be the same object
           that is passed to :meth:`compute` for the actual calculation).
        chan_per_worker : None or int
           Number of channels to be processed by each worker (only relevant in
           case of concurrent processing). If `chan_per_worker` is `None` (default)
           by-trial parallelism is used, i.e., each worker processes
           data corresponding to a full trial. If `chan_per_worker > 0`, trials
           are split into channel-groups of size `chan_per_worker` (+ rest if the
           number of channels is not divisible by `chan_per_worker` without
           remainder) and workers are assigned by-trial channel-groups for
           processing.
        keeptrials : bool
            Flag indicating whether to return individual trials or average

        Returns
        -------
        Nothing : None

        Notes
        -----
        This class method **has** to be called prior to performing the actual
        computation realized in :meth:`computeFunction`.

        See also
        --------
        compute : core routine performing the actual computation
        """

        # First store `keeptrial` keyword value (important for output shapes below)
        self.keeptrials = keeptrials

        # Determine if data-selection was provided; if so, extract trials and check
        # whether selection requires fancy array indexing
        if data._selection is not None:
            self.trialList = data._selection.trials
            self.useFancyIdx = data._selection._useFancy
        else:
            self.trialList = list(range(len(data.trials)))
            self.useFancyIdx = False
        numTrials = len(self.trialList)

        # If lists/tuples are in positional arguments, ensure `len == numTrials`
        # Scalars are duplicated to fit trials, e.g., ``self.argv = [3, [0, 1, 1]]``
        # then ``argv = [[3, 3, 3], [0, 1, 1]]``
        for ak, arg in enumerate(self.argv):

            # Ensure arguments are within reasonable size for distribution across workers
            # (protect against circular object references by imposing max. calls)
            self._callCount = 0
            argsize = self._sizeof(arg)
            if argsize > self._maxArgSize:
                lgl = "positional arguments less than 100 MB each"
                act = "positional argument with memory footprint of {0:4.2f} MB"
                raise SPYValueError(legal=lgl, varname="argv", actual=act.format(argsize))

            if isinstance(arg, (list, tuple)):
                if len(arg) != numTrials:
                    lgl = "list/tuple of positional arguments for each trial"
                    act = "length of list/tuple does not correspond to number of trials"
                    raise SPYValueError(legal=lgl, varname="argv", actual=act)
                continue
            elif isinstance(arg, np.ndarray):
                if arg.size == numTrials:
                    msg = "found NumPy array with size == #Trials. " +\
                        "Regardless, every worker will receive an identical copy " +\
                        "of this array. To propagate elements across workers, use " +\
                        "a list or tuple instead!"
                    SPYWarning(msg)
            self.argv[ak] = [arg] * numTrials

        # Prepare dryrun arguments and determine geometry of trials in output
        dryRunKwargs = copy(self.cfg)
        dryRunKwargs["noCompute"] = True
        chk_list = []
        dtp_list = []
        trials = []
        for tk, trialno in enumerate(self.trialList):
            trial = data._preview_trial(trialno)
            trlArg = tuple(arg[tk] for arg in self.argv)
            chunkShape, dtype = self.computeFunction(trial,
                                                     *trlArg,
                                                     **dryRunKwargs)
            chk_list.append(list(chunkShape))
            dtp_list.append(dtype)
            trials.append(trial)

        # The aggregate shape is computed as max across all chunks
        chk_arr = np.array(chk_list)
        if np.unique(chk_arr[:, 0]).size > 1 and not self.keeptrials:
            err = "Averaging trials of unequal lengths in output currently not supported!"
            raise NotImplementedError(err)
        if np.any([dtp_list[0] != dtp for dtp in dtp_list]):
            lgl = "unique output dtype"
            act = "{} different output dtypes".format(np.unique(dtp_list).size)
            raise SPYValueError(legal=lgl, varname="dtype", actual=act)
        chunkShape = tuple(chk_arr.max(axis=0))
        self.outputShape = (chk_arr[:, 0].sum(),) + chunkShape[1:]
        self.cfg["chunkShape"] = chunkShape
        self.dtype = np.dtype(dtp_list[0])

        # Ensure channel parallelization can be done at all
        if chan_per_worker is not None and "channel" not in data.dimord:
            msg = "input object does not contain `channel` dimension for parallelization!"
            SPYWarning(msg)
            chan_per_worker = None
        if chan_per_worker is not None and self.keeptrials is False:
            msg = "trial-averaging does not support channel-block parallelization!"
            SPYWarning(msg)
            chan_per_worker = None
        if data._selection is not None:
            if chan_per_worker is not None and data._selection.channel != slice(None, None, 1):
                msg = "channel selection and simultaneous channel-block " +\
                    "parallelization not yet supported!"
                SPYWarning(msg)
                chan_per_worker = None

        # Allocate control variables
        trial = trials[0]
        trlArg0 = tuple(arg[0] for arg in self.argv)
        chunkShape0 = chk_arr[0, :]
        lyt = [slice(0, stop) for stop in chunkShape0]
        sourceLayout = []
        targetLayout = []
        targetShapes = []
        ArgV = []

        # If parallelization across channels is requested the first trial is
        # split up into several chunks that need to be processed/allocated
        if chan_per_worker is not None:

            # Set up channel-chunking
            nChannels = data.channel.size
            rem = int(nChannels % chan_per_worker)
            n_blocks = [chan_per_worker] * int(nChannels//chan_per_worker) + [rem] * int(rem > 0)
            inchanidx = data.dimord.index("channel")

            # Perform dry-run w/first channel-block of first trial to identify
            # changes in output shape w.r.t. full-trial output (`chunkShape`)
            shp = list(trial.shape)
            idx = list(trial.idx)
            shp[inchanidx] = n_blocks[0]
            idx[inchanidx] = slice(0, n_blocks[0])
            trial.shape = tuple(shp)
            trial.idx = tuple(idx)
            res, _ = self.computeFunction(trial, *trlArg0, **dryRunKwargs)
            outchan = [dim for dim in res if dim not in chunkShape0]
            if len(outchan) != 1:
                lgl = "exactly one output dimension to scale w/channel count"
                act = "{0:d} dimensions affected by varying channel count".format(len(outchan))
                raise SPYValueError(legal=lgl, varname="chan_per_worker", actual=act)
            outchanidx = res.index(outchan[0])

            # Get output chunks and grid indices for first trial
            chanstack = 0
            blockstack = 0
            for block in n_blocks:
                shp = list(trial.shape)
                idx = list(trial.idx)
                shp[inchanidx] = block
                idx[inchanidx] = slice(blockstack, blockstack + block)
                trial.shape = tuple(shp)
                trial.idx = tuple(idx)
                res, _ = self.computeFunction(trial, *trlArg0, **dryRunKwargs)
                lyt[outchanidx] = slice(chanstack, chanstack + res[outchanidx])
                targetLayout.append(tuple(lyt))
                targetShapes.append(tuple([slc.stop - slc.start for slc in lyt]))
                sourceLayout.append(trial.idx)
                ArgV.append(trlArg0)
                chanstack += res[outchanidx]
                blockstack += block

        # Simple: consume all channels simultaneously, i.e., just take the entire trial
        else:
            targetLayout.append(tuple(lyt))
            targetShapes.append(chunkShape0)
            sourceLayout.append(trial.idx)
            ArgV.append(trlArg0)

        # Construct dimensional layout of output
        stacking = targetLayout[0][0].stop
        for tk in range(1, len(self.trialList)):
            trial = trials[tk]
            trlArg = tuple(arg[tk] for arg in self.argv)
            chkshp = chk_list[tk]
            lyt = [slice(0, stop) for stop in chkshp]
            lyt[0] = slice(stacking, stacking + chkshp[0])
            stacking += chkshp[0]
            if chan_per_worker is None:
                targetLayout.append(tuple(lyt))
                targetShapes.append(tuple([slc.stop - slc.start for slc in lyt]))
                sourceLayout.append(trial.idx)
                ArgV.append(trlArg)
            else:
                chanstack = 0
                blockstack = 0
                for block in n_blocks:
                    shp = list(trial.shape)
                    idx = list(trial.idx)
                    shp[inchanidx] = block
                    idx[inchanidx] = slice(blockstack, blockstack + block)
                    trial.shape = tuple(shp)
                    trial.idx = tuple(idx)
                    res, _ = self.computeFunction(trial, *trlArg, **dryRunKwargs) # FauxTrial
                    lyt[outchanidx] = slice(chanstack, chanstack + res[outchanidx])
                    targetLayout.append(tuple(lyt))
                    targetShapes.append(tuple([slc.stop - slc.start for slc in lyt]))
                    sourceLayout.append(trial.idx)
                    chanstack += res[outchanidx]
                    blockstack += block
                    ArgV.append(trlArg)

        # If the determined source layout contains unordered lists and/or index
        # repetitions, set `self.useFancyIdx` to `True` and prepare a separate
        # `sourceSelectors` list that is used in addition to `sourceLayout` for
        # data extraction.
        # In this case `sourceLayout` uses ABSOLUTE indices (indices wrt to size
        # of ENTIRE DATASET) that are SORTED W/O REPS to extract a NumPy array
        # of appropriate size from HDF5.
        # Then `sourceLayout` uses RELATIVE indices (indices wrt to size of CURRENT
        # TRIAL) that can be UNSORTED W/REPS to actually perform the requested
        # selection on the NumPy array extracted w/`sourceLayout`.
        for grd in sourceLayout:
            if any([np.diff(sel).min() <= 0 if isinstance(sel, list)
                    and len(sel) > 1 else False for sel in grd]):
                self.useFancyIdx = True
                break
        if self.useFancyIdx:
            sourceSelectors = []
            for gk, grd in enumerate(sourceLayout):
                ingrid = list(grd)
                sigrid = []
                for sk, sel in enumerate(grd):
                    if isinstance(sel, list):
                        selarr = np.array(sel, dtype=np.intp)
                    else: # sel is a slice
                        step = sel.step
                        if sel.step is None:
                            step = 1
                        selarr = np.array(list(range(sel.start, sel.stop, step)), dtype=np.intp)
                    if selarr.size > 0:
                        sigrid.append(np.array(selarr) - selarr.min())
                        ingrid[sk] = slice(selarr.min(), selarr.max() + 1, 1)
                    else:
                        sigrid.append([])
                        ingrid[sk] = []
                sourceSelectors.append(tuple(sigrid))
                sourceLayout[gk] = tuple(ingrid)
        else:
            sourceSelectors = [Ellipsis] * len(sourceLayout)

        # Store determined shapes and grid layout
        self.sourceLayout = sourceLayout
        self.sourceSelectors = sourceSelectors
        self.targetLayout = targetLayout
        self.targetShapes = targetShapes
        self.ArgV = ArgV

        # Compute max. memory footprint of chunks
        if chan_per_worker is None:
            self.chunkMem = np.prod(self.cfg["chunkShape"]) * self.dtype.itemsize
        else:
            self.chunkMem = max([np.prod(shp) for shp in self.targetShapes]) * self.dtype.itemsize

        # Get data access mode (only relevant for parallel reading access)
        self.dataMode = data.mode

    def compute(self, data, out, parallel=False, parallel_store=None,
                method=None, mem_thresh=0.5, log_dict=None, parallel_debug=False):
        """
        Central management and processing method

        Parameters
        ----------
        data : syncopy data object
           Syncopy data object to be processed (has to be the same object
           that was used by :meth:`initialize` in the pre-calculation
           dry-run).
        out : syncopy data object
           Empty object for holding results
        parallel : bool
           If `True`, processing is performed in parallel (i.e.,
           :meth:`computeFunction` is executed concurrently across trials).
           If `parallel` is `False`, :meth:`computeFunction` is executed
           consecutively trial after trial (i.e., the calculation realized
           in :meth:`computeFunction` is performed sequentially).
        parallel_store : None or bool
           Flag controlling saving mechanism. If `None`,
           ``parallel_store = parallel``, i.e., the compute-paradigm
           dictates the employed writing method. Thus, in case of parallel
           processing, results are written in a fully concurrent
           manner (each worker saves its own local result segment on disk as
           soon as it is done with its part of the computation). If `parallel_store`
           is `False` and `parallel` is `True` the processing result is saved
           sequentially using a mutex. If both `parallel` and `parallel_store`
           are `False` standard single-process HDF5 writing is employed for
           saving the result of the (sequential) computation.
        method : None or str
           If `None` the predefined methods :meth:`compute_parallel` or
           :meth:`compute_sequential` are used to control the actual computation
           (specifically, calling :meth:`computeFunction`) depending on whether
           `parallel` is `True` or `False`, respectively. If `method` is a
           string, it has to specify the name of an alternative (provided)
           class method that is invoked using `getattr`.
        mem_thresh : float
           Fraction of available memory required to perform computation. By
           default, the largest single trial result must not occupy more than
           50% (``mem_thresh = 0.5``) of available single-machine or worker
           memory (if `parallel` is `False` or `True`, respectively).
        log_dict : None or dict
           If `None`, the `log` properties of `out` is populated with the employed
           keyword arguments used in :meth:`computeFunction`.
           Otherwise, `out`'s `log` properties are filled  with items taken
           from `log_dict`.
        parallel_debug : bool
           If `True`, concurrent processing is performed using a single-threaded
           scheduler, i.e., all parallel computing task are run in the current
           Python thread permitting usage of tools like `pdb`/`ipdb`, `cProfile`
           and the like in :meth:`computeFunction`.
           Note that enabling parallel debugging effectively runs the given computation
           on the calling local machine thereby requiring sufficient memory and
           CPU capacity.

        Returns
        -------
        Nothing : None
           The result of the computation is available in `out` once
           :meth:`compute` terminated successfully.

        Notes
        -----
        This routine calls several other class methods to perform all necessary
        pre- and post-processing steps in a fully automatic manner without
        requiring any user-input. Specifically, the following class methods
        are invoked consecutively (in the given order):

        1. :meth:`preallocate_output` allocates a (virtual) HDF5 dataset
           of appropriate dimension for storing the result
        2. :meth:`compute_parallel` (or :meth:`compute_sequential`) performs
           the actual computation via concurrently (or sequentially) calling
           :meth:`computeFunction`
        3. :meth:`process_metadata` attaches all relevant meta-information to
           the result `out` after successful termination of the calculation
        4. :meth:`write_log` stores employed input arguments in `out.cfg`
           and `out.log` to reproduce all relevant computational steps that
           generated `out`.

        See also
        --------
        initialize : pre-calculation preparations
        preallocate_output : storage provisioning
        compute_parallel : concurrent computation using :meth:`computeFunction`
        compute_sequential : sequential computation using :meth:`computeFunction`
        process_metadata : management of meta-information
        write_log : log-entry organization
        """

        # By default, use VDS storage for parallel computing
        if parallel_store is None:
            parallel_store = parallel

        # Do not spill trials on disk if they're supposed to be removed anyway
        if parallel_store and not self.keeptrials:
            msg = "trial-averaging only supports sequential writing!"
            SPYWarning(msg)
            parallel_store = False

        # Concurrent processing requires some additional prep-work...
        if parallel:

            # First and foremost, make sure a dask client is accessible
            try:
                client = dd.get_client()
            except ValueError as exc:
                msg = "parallel computing client: {}"
                raise SPYIOError(msg.format(exc.args[0]))

            # Check if the underlying cluster hosts actually usable workers
            if not len(client.cluster.workers):
                raise SPYParallelError("No active workers found in distributed computing cluster",
                                       client=client)

            # Note: `dask_jobqueue` may not be available even if `__dask__` is `True`,
            # hence the `__name__` shenanigans instead of a simple `isinstance`
            if isinstance(client.cluster, dd.LocalCluster):
                memAttr = "memory_limit"
            elif client.cluster.__class__.__name__ == "SLURMCluster":
                memAttr = "worker_memory"
            else:
                msg = "`ComputationalRoutine` only supports `LocalCluster` and " +\
                    "`SLURMCluster` dask cluster objects. Proceed with caution. "
                SPYWarning(msg)
                memAttr = None

            # Check if trials actually fit into memory before we start computation
            if memAttr:
                wrk_size = max(getattr(wrkr, memAttr) for wrkr in client.cluster.workers.values())
                if self.chunkMem >= mem_thresh * wrk_size:
                    self.chunkMem /= 1024**3
                    wrk_size /= 1000**3
                    msg = "Single-trial result sizes ({0:2.2f} GB) larger than available " +\
                        "worker memory ({1:2.2f} GB) currently not supported"
                    raise NotImplementedError(msg.format(self.chunkMem, wrk_size))

            # In some cases distributed dask workers suffer from spontaneous
            # dementia and forget the `sys.path` of their parent process. Fun!
            def init_syncopy(dask_worker):
                spy_path = os.path.abspath(os.path.split(__path__[0])[0])
                if spy_path not in sys.path:
                    sys.path.insert(0, spy_path)
            client.register_worker_callbacks(init_syncopy)

            # Store provided debugging state
            self.parallelDebug = parallel_debug

        # For sequential processing, just ensure enough memory is available
        else:
            mem_size = psutil.virtual_memory().available
            if self.chunkMem >= mem_thresh * mem_size:
                self.chunkMem /= 1024**3
                mem_size /= 1024**3
                msg = "Single-trial result sizes ({0:2.2f} GB) larger than available " +\
                      "memory ({1:2.2f} GB) currently not supported"
                raise NotImplementedError(msg.format(self.chunkMem, mem_size))

        # Create HDF5 dataset of appropriate dimension
        self.preallocate_output(out, parallel_store=parallel_store)

        # The `method` keyword can be used to override the `parallel` flag
        if method is None:
            if parallel:
                computeMethod = self.compute_parallel
            else:
                computeMethod = self.compute_sequential
        else:
            computeMethod = getattr(self, "compute_" + method, None)

        # Ensure `data` is openend read-only to permit (potentially concurrent)
        # reading access to backing device on disk
        data.mode = "r"

        # Take care of `VirtualData` objects
        self.hdr = getattr(data, "hdr", None)

        # Perform actual computation
        computeMethod(data, out)

        # Reset data access mode
        data.mode = self.dataMode

        # Attach computed results to output object
        out.data = h5py.File(out.filename, mode="r+")[self.datasetName]

        # Store meta-data, write log and get outta here
        self.process_metadata(data, out)
        self.write_log(data, out, log_dict)

    def preallocate_output(self, out, parallel_store=False):
        """
        Storage allocation and provisioning

        Parameters
        ----------
        out : syncopy data object
           Empty object for holding results
        parallel_store : bool
           If `True`, a directory for virtual source files is created
           in Syncopy's temporary on-disk storage (defined by `syncopy.__storage__`).
           Otherwise, a dataset of appropriate type and shape is allocated
           in a new regular HDF5 file created inside Syncopy's temporary
           storage folder.

        Returns
        -------
        Nothing : None

        See also
        --------
        compute : management routine controlling memory pre-allocation
        """

        # Set name of target HDF5 dataset in output object
        self.datasetName = "data"

        # In case parallel writing via VDS storage is requested, prepare
        # directory for by-chunk HDF5 files and construct virutal HDF layout
        if parallel_store:
            vdsdir = os.path.splitext(os.path.basename(out.filename))[0]
            self.virtualDatasetDir = os.path.join(__storage__, vdsdir)
            os.mkdir(self.virtualDatasetDir)

            layout = h5py.VirtualLayout(shape=self.outputShape, dtype=self.dtype)
            for k, idx in enumerate(self.targetLayout):
                fname = os.path.join(self.virtualDatasetDir, "{0:d}.h5".format(k))
                layout[idx] = h5py.VirtualSource(fname, "chk", shape=self.targetShapes[k])
            self.VirtualDatasetLayout = layout

        # Create regular HDF5 dataset for sequential writing
        else:

            # The shape of the target depends on trial-averaging
            if not self.keeptrials:
                shp = self.cfg["chunkShape"]
            else:
                shp = self.outputShape
            with h5py.File(out.filename, mode="w") as h5f:
                h5f.create_dataset(name=self.datasetName,
                                   dtype=self.dtype, shape=shp)

    def compute_parallel(self, data, out):
        """
        Concurrent computing kernel

        Parameters
        ----------
        data : syncopy data object
           Syncopy data object to be processed
        out : syncopy data object
           Empty object for holding results

        Returns
        -------
        Nothing : None

        Notes
        -----
        This method mereley acts as a concurrent wrapper for :meth:`computeFunction`
        by passing along all necessary information for parallel execution and
        storage of results using a dask bag of dictionaries. The actual reading
        of source data and writing of results is managed by the decorator
        :func:`syncopy.shared.parsers.unwrap_io`. Note that this routine first
        builds an entire parallel instruction tree and only kicks off execution
        on the cluster at the very end of the calculation command assembly.

        See also
        --------
        compute : management routine invoking parallel/sequential compute kernels
        compute_sequential : serial processing counterpart of this method
        """

        # Prepare to write chunks concurrently
        if self.virtualDatasetDir is not None:
            outfilename = os.path.join(self.virtualDatasetDir, "{0:d}.h5")
            outdsetname = "chk"

        # Write chunks sequentially
        else:
            outfilename = out.filename
            outdsetname = self.datasetName

        # Construct a dask bag with all necessary components for parallelization
        mainBag = db.from_sequence([{"hdr": self.hdr,
                                     "keeptrials": self.keeptrials,
                                     "infile": data.filename,
                                     "indset": data.data.name,
                                     "ingrid": self.sourceLayout[chk],
                                     "sigrid": self.sourceSelectors[chk],
                                     "fancy": self.useFancyIdx,
                                     "vdsdir": self.virtualDatasetDir,
                                     "outfile": outfilename.format(chk),
                                     "outdset": outdsetname,
                                     "outgrid": self.targetLayout[chk],
                                     "outshape": self.targetShapes[chk],
                                     "dtype": self.dtype}
                                     for chk in range(len(self.sourceLayout))],
                                   npartitions=len(self.sourceLayout))

        # Convert by-worker argv-list to dask bags to distribute across cluster
        # Format: ``ArgV = [(3, 0, 'a'), (3, 0, 'a'), (3, 1, 'b'), (3, 1, 'b')]``
        # then ``list(zip(*ArgV)) = [(3, 3, 3, 3), (0, 0, 1, 1), ('a', 'a', 'b', 'b')]``
        bags = []
        for arg in zip(*self.ArgV):
            bags.append(db.from_sequence(arg))

        # Map all components (channel-trial-blocks) onto `computeFunction`
        results = mainBag.map(self.computeFunction, *bags, **self.cfg)

        # Let the fun begin...
        if not self.parallelDebug:

            # Make sure that all futures are executed (i.e., data is actually written)
            # Note 1: `dd.progress` does not correctly track worker progress hence
            #         the custom-tailored `while` formulation: that periodically
            #         checks in on the status of all allocated futures
            #         -> Do not use this `dd.progress(futures, notebook=False)`
            # Note 2: the while loop below does not run indefinitely - erring or
            #         stalling futures get status 'error' or 'waiting'. After
            #         some time the status of all futures is one of 'finished',
            #         'error' or 'waiting', but none is 'pending' any more.
            futures = dd.client.futures_of(results.persist())
            totalTasks = len(futures)
            pbar = tqdm(total=totalTasks, bar_format=self.tqdmFormat)
            cnt = 0
            while any(f.status == "pending" for f in futures):
                time.sleep(self.sleepTime)
                new = max(0, sum([f.status == "finished" for f in futures]) - cnt)
                cnt += new
                pbar.update(new)
            pbar.close()

            # Avoid race condition: give futures time to perform switch from 'pending'
            # to 'finished' so that `finishedTasks` is computed correctly
            time.sleep(self.sleepTime)

            # If number of 'finished' tasks is less than expected, go into
            # problem analysis mode: all futures that erred hav an `.exception`
            # method which can be used to track down the worker it was executed by
            # Once we know the worker, we can point to the right log file. If
            # futures were cancelled (by the user or the SLURM controller),
            # `.exception` is `None` and we can't relialby track down the
            # respective executing worker
            finishedTasks = sum([f.status == "finished" for f in futures])
            if finishedTasks < totalTasks:
                client = dd.get_client()
                schedulerLog = list(client.cluster.get_logs(cluster=False, scheduler=True, workers=False).values())[0]
                erredFutures = [f for f in futures if f.status == "error"]
                msg = "Parallel computation failed: {}/{} tasks failed or stalled.\n"
                msg = msg.format(totalTasks - finishedTasks, totalTasks)
                msg += "Concurrent computing scheduler log below: \n\n"
                msg += schedulerLog + "\n"

                # If we're working w/`SLURMCluster`, perform the Herculean task of
                # tracking down which dask worker was executed by which SLURM job...
                if client.cluster.__class__.__name__ == "SLURMCluster":
                    try:
                        erredJobs = [f.exception().last_worker.identity()["id"] for f in erredFutures]
                    except AttributeError:
                        erredJobs = []
                    erredJobs = list(set(erredJobs))
                    erredJobIDs = [client.cluster.workers[job].job_id for job in erredJobs]
                    slurmFiles = client.cluster.job_header.split("--output=")[1].replace("%j", "{}")
                    slurmOutDir = os.path.split(slurmFiles)[0]
                    errFiles = glob(slurmOutDir + os.sep + "*.err")
                    if len(erredFutures) or len(errFiles):
                        msg += "Please consult the following SLURM log files for details:\n"
                        msg += "".join(slurmFiles.format(id) + "\n" for id in erredJobIDs)
                        msg += "".join(errfile + "\n" for errfile in errFiles)
                    else:
                        msg += "Please check SLURM logs in {}".format(slurmOutDir)

                # In case of a `LocalCluster`, syphon worker logs
                else:
                    msg += "\nParallel worker logs below: \n"
                    workerLogs = client.cluster.get_logs(cluster=False, scheduler=False, workers=True).values()
                    for wLog in workerLogs:
                        if "Failed" in wLog:
                            msg += wLog

                raise SPYParallelError(msg, client=client)

        # If debugging is requested, drop existing client and enforce use of
        # single-threaded scheduler
        else:
            results.compute(scheduler="single-threaded")

        # When writing concurrently, now's the time to finally create the virtual dataset
        if self.virtualDatasetDir is not None:
            with h5py.File(out.filename, mode="w") as h5f:
                h5f.create_virtual_dataset(self.datasetName, self.VirtualDatasetLayout)

        # If trial-averaging was requested, normalize computed sum to get mean
        if not self.keeptrials:
            with h5py.File(out.filename, mode="r+") as h5f:
                h5f[self.datasetName][()] /= len(self.trialList)
                h5f.flush()

        return

    def compute_sequential(self, data, out):
        """
        Sequential computing kernel

        Parameters
        ----------
        data : syncopy data object
           Syncopy data object to be processed
        out : syncopy data object
           Empty object for holding results

        Returns
        -------
        Nothing : None

        Notes
        -----
        This method most closely reflects classic iterative process execution:
        trials in `data` are passed sequentially to :meth:`computeFunction`,
        results are stored consecutively in a regular HDF5 dataset (that was
        pre-allocated by :meth:`preallocate_output`). Since the calculation result
        is immediately stored on disk, propagation of arrays across routines
        is avoided and memory usage is kept to a minimum.

        See also
        --------
        compute : management routine invoking parallel/sequential compute kernels
        compute_parallel : concurrent processing counterpart of this method
        """

        # Initialize on-disk backing device (either HDF5 file or memmap)
        if self.hdr is None:
            try:
                sourceObj = h5py.File(data.filename, mode="r")[data.data.name]
                isHDF = True
            except OSError:
                sourceObj = open_memmap(data.filename, mode="c")
                isHDF = False
            except Exception as exc:
                raise exc

        # Iterate over (selected) trials and write directly to target HDF5 dataset
        fmt = "{desc}: {percentage:3.0f}% |{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]"
        with h5py.File(out.filename, "r+") as h5fout:
            target = h5fout[self.datasetName]

            for nblock in tqdm(range(len(self.trialList)), bar_format=self.tqdmFormat):

                # Extract respective indexing tuples from constructed lists
                ingrid = self.sourceLayout[nblock]
                sigrid = self.sourceSelectors[nblock]
                outgrid = self.targetLayout[nblock]
                argv = self.ArgV[nblock]

                # Catch empty source-array selections; this workaround is not
                # necessary for h5py version 2.10+ (see https://github.com/h5py/h5py/pull/1174)
                if any([not sel for sel in ingrid]):
                    res = np.empty(self.targetShapes[nblock], dtype=self.dtype)
                else:
                    # Get source data as NumPy array
                    if self.hdr is None:
                        if isHDF:
                            if self.useFancyIdx:
                                arr = np.array(sourceObj[tuple(ingrid)])[np.ix_(*sigrid)]
                            else:
                                arr = np.array(sourceObj[tuple(ingrid)])
                        else:
                            if self.useFancyIdx:
                                arr = sourceObj[np.ix_(*ingrid)]
                            else:
                                arr = np.array(sourceObj[ingrid])
                        sourceObj.flush()
                    else:
                        idx = ingrid
                        if self.useFancyIdx:
                            idx = np.ix_(*ingrid)
                        stacks = []
                        for fk, fname in enumerate(data.filename):
                            stacks.append(np.memmap(fname, offset=int(self.hdr[fk]["length"]),
                                                    mode="r", dtype=self.hdr[fk]["dtype"],
                                                    shape=(self.hdr[fk]["M"], self.hdr[fk]["N"]))[idx])
                        arr = np.vstack(stacks)[ingrid]

                    # Perform computation
                    res = self.computeFunction(arr, *argv, **self.cfg)

                # Either write result to `outgrid` location in `target` or add it up
                if self.keeptrials:
                    target[outgrid] = res
                else:
                    target[()] = np.nansum([target, res], axis=0)

                # Flush every iteration to avoid memory leakage
                h5fout.flush()

            # If trial-averaging was requested, normalize computed sum to get mean
            if not self.keeptrials:
                target[()] /= len(self.trialList)

        # If source was HDF5 file, close it to prevent access errors
        if isHDF:
            sourceObj.file.close()

        return

    def write_log(self, data, out, log_dict=None):
        """
        Processing of output log

        Parameters
        ----------
        data : syncopy data object
           Syncopy data object that has been processed
        out : syncopy data object
           Syncopy data object holding calculation results
        log_dict : None or dict
           If `None`, the `log` properties of `out` is populated with the employed
           keyword arguments used in :meth:`computeFunction`.
           Otherwise, `out`'s `log` properties are filled  with items taken
           from `log_dict`.

        Returns
        -------
        Nothing : None

        See also
        --------
        process_metadata : Management of meta-information
        """

        # Copy log from source object and write header
        out._log = str(data._log) + out._log
        logHead = "computed {name:s} with settings\n".format(name=self.computeFunction.__name__)

        # Prepare keywords used by `computeFunction` (sans implementation-specific stuff)
        cfg = dict(self.cfg)
        for key in ["noCompute", "chunkShape"]:
            cfg.pop(key)

        # Write log and store `cfg` constructed above in corresponding prop of `out`
        if log_dict is None:
            log_dict = cfg
        logOpts = ""
        for k, v in log_dict.items():
            logOpts += "\t{key:s} = {value:s}\n".format(key=k,
                                                        value=str(v) if len(str(v)) < 80
                                                        else str(v)[:30] + ", ..., " + str(v)[-30:])
        out.log = logHead + logOpts
        out.cfg = cfg

    @abstractmethod
    def process_metadata(self, data, out):
        """
        Meta-information manager

        Parameters
        ----------
        data : syncopy data object
           Syncopy data object that has been processed
        out : syncopy data object
           Syncopy data object holding calculation results

        Returns
        -------
        Nothing : None

        Notes
        -----
        This routine is an abstract method and is thus intended to be overloaded.
        Consult the developer documentation (:doc:`/developer/compute_kernels`) for
        further details.

        See also
        --------
        write_log : Logging of calculation parameters
        """
        pass

    def _sizeof(self, obj):
        """
        Estimate memory consumption of Python objects

        Parameters
        ----------
        obj : Python object
           Any valid Python object whose memory footprint is of interest.

        Returns
        -------
        objsize : float
           Approximate memory footprint of `obj` in megabytes (MB).

        Notes
        -----
        Memory consumption is is estimated by recursively calling :meth:`sys.getsizeof`.
        Circular object references are followed up to a (preset) maximal recursion
        depth. This method was inspired by a routine in
        `Nifty <https://github.com/mwojnars/nifty/blob/master/util.py>`_.
        """

        # Protect against circular object references by adhering to max. no. of
        # recursive calls `self._callMax`
        self._callCount += 1
        if self._callCount >= self._callMax:
            lgl = "minimally nested positional arguments"
            act = "argument with nesting depth >= {}"
            raise SPYValueError(legal=lgl, varname="argv", actual=act.format(self._callMax))

        # Use `sys.getsizeof` to estimate memory consumption of primitive objects
        objsize = sys.getsizeof(obj) / 1024**2
        if isinstance(obj, dict):
            return objsize + sum(list(map(self._sizeof, obj.keys()))) + sum(list(map(self._sizeof, obj.values())))
        if isinstance(obj, (list, tuple, set)):
            return objsize + sum(list(map(self._sizeof, obj)))
        return objsize


# -*- coding: utf-8 -*-
import numpy as np
import sys
import gym
import math
import os

from nodes import IzhikevichNode, InputEncoder, InputNode
from connections import AbstractConnection, EncoderConnection
from network import Network
from monitors import SpikeMonitor
from encoders import probabilities, normalized_probabilities, rate_based_encoding
from decodingMethod import DecodingMethod
import matplotlib.pyplot as plt

class CartPoleRSTDP:

    def __init__(self, NeuronsPerLayer, ExposurePeriod: int = 50, Decoding: DecodingMethod = DecodingMethod.F2F_RESET, filename: str = "test.txt"):
        # Create the network according to the given parameters

        self.Environment = gym.make('CartPole-v1')
        self.Environment.reset()

        # Network Values
        #self.LearningRate = 1.0 # TODO: Need to recover values from hyperparameter experiment

        self.Decoding = Decoding
        self.ExposurePeriod = ExposurePeriod
        self.FileName = filename

        # Create the network
        self.Network = self.CreateNetwork(NeuronsPerLayer)

        self.LearningRate = 0.1
        #self.LearningRate = 0.1
        #print(self.LearningRate)

    def CreateNetwork(self, NeuronsPerLayer):

        AllLayers = []

        # 1: Create the layers
        for i in range(len(NeuronsPerLayer)):
            if i == 0:
                LayerInput = IzhikevichNode(NeuronsPerLayer[0])
                AllLayers.append(LayerInput)
            elif i == len(NeuronsPerLayer) - 1:
                LayerOutput = IzhikevichNode(NeuronsPerLayer[i])
                AllLayers.append(LayerOutput)
            else:
                AllLayers.append(IzhikevichNode(NeuronsPerLayer[i]))

        # 2: Add the connections between the layers

        AllConnections = []

        for i in range(len(AllLayers) - 1):
            if i >= 0:
                AllConnections.append(AbstractConnection(
                    source=AllLayers[i],
                    target=AllLayers[i+1],
                    wmin=-20.0,
                    wmax=80.0))
                #AllConnections[0].w.fill(18.75)
            else:
                 AllConnections.append(EncoderConnection(
                    source=AllLayers[i],
                    target=AllLayers[i+1],
                    wmin=40.0,
                    wmax=150.0))
                 AllConnections[0].w.fill(50)
            #AllConnections[i].presynaptic_normalization(20)

        # 3: Create the network by giving the input and output layers
        SNN = Network(LayerInput, LayerOutput)
        for i in range(len(AllLayers)):
            if i == 0:
                SNN.AddLayer(LayerInput, "Input Layer")
            elif i == len(AllLayers) - 1:
                SNN.AddLayer(LayerOutput, "Output Layer")
            else:
                SNN.AddLayer(AllLayers[i], "Layer " + str(i))

        for i in range(len(AllConnections)):
            SNN.AddConnection(AllConnections[i])

        #SNN.Connections[1].postsynaptic_normalization(150.0)

        # Add the reward monitor to the nework
        #RewMonitor = RewardMonitor()
        #SNN.AddMonitor(RewMonitor, "Reward Monitor")
        return SNN

    def Preprocess(self, obs):
        newObs = np.asarray([
            obs[0]*0.208 if obs[0] >= 0 else 0,
            abs(obs[0])*0.208 if obs[0] < 0 else 0,
            obs[1]*0.250 if obs[1] >= 0 else 0,
            abs(obs[1])*0.250 if obs[1] < 0 else 0,
            obs[2]*2.4 if obs[2] >= 0 else 0,
            abs(obs[2])*2.4 if obs[2] < 0 else 0,
            obs[3]*0.25 if obs[3] >= 0 else 0,
            abs(obs[3])*0.25 if obs[3] < 0 else 0
            #obs[4] if obs[4] >= 0 else 0,
            #abs(obs[4]) if obs[4] < 0 else 0,
            #obs[5] if obs[5] >= 0 else 0,
            #abs(obs[5]) if obs[5] < 0 else 0,
            #obs[6] if obs[6] >= 0 else 0,
            #abs(obs[6]) if obs[6] < 0 else 0,
            #obs[7] if obs[7] >= 0 else 0,
            #abs(obs[7]) if obs[7] < 0 else 0
        ], dtype=np.float32)
        return newObs

    def Run(self, NumEpisodes: int = 10):
        #e = 0
        #RewardTotal = 0
        self.AllTotalRewards = []
        for i in range(NumEpisodes):
            self.AllSpikesInputLayer = []
            self.AllSpikesOutputLayer = []
            self.AllActions = []
            self.AllEpisodeRewards = []
            self.RunEpisode(False)
            #RewardTotal += Reward
            #print(np.asarray(self.AllSpikesInputLayer[1][50]))
            #print(np.asarray(self.AllSpikesOutputLayer))
            #print(self.Network.Connections[0].w)
            #print(self.Network.Connections[1].w[:])
            #print(self.Network.Connections[1].w[1])
            #print(self.Network.Connections[1].w[:,0])
            #print (self.AllSpikesInputLayer[1][0].size)
            #print("Avg Fitness: " + str(np.average(self.AllTotalRewards)))
            self.UpdateWeights()
            #self.Network.Connections[0].postsynaptic_normalization(200)
            #self.Network.Connections[0].presynaptic_normalization(100)
            self.Network.Connections[0].decay(0.005)
            #print(np.sum(self.Network.Connections[0].w))
            if (i%100 == 0 and i != 0):
                self.PrintToFile(self.FileName, self.Network.Connections[0].w)
                self.PrintToFile(self.FileName, str(sum(self.AllTotalRewards[-100:])/100))
                print(sum(self.AllTotalRewards[-100:])/100)
                print(self.Network.Connections[0].w)
            #self.Network.Connections[1].random_mutation(0.01)

            #e += 1
            #if (e%100 == 0):
                #print(RewardTotal / 100)
                #print (self.AllSpikesInputLayer[1][0].size)
                #print(self.Network.Connections[1].w)
                #RewardTotal = 0

    def RunEpisode(self, DisplayEpisode: bool):
        self.AllActions = []
        #self.AllRewards = []

        self.AllSpikesInputLayer = []
        self.AllSpikesOutputLayer = []
        #self.AllEpisodeRewards = []
        obs = self.Environment.reset()
        done = False
        # Will hold all of the actions taken in the next to run episode
        Actions = []
        RewardTotal = 0
        for _ in range(200):
            newObs = self.Preprocess(obs)
            #print(newObs)
            newObs = newObs*10.0
            #print(newObs)
            action, _ = self.Network.Run(newObs, self.ExposurePeriod, True, self.Decoding)
            # Add the action to the set of actions being recorded for this episode
            Actions.append(action)
            #print(action)
            # Get the output spike train and make a decision based on that

            obs, reward, done, _ = self.Environment.step(action)
            self.AllEpisodeRewards.append(reward)
            #if DisplayEpisode:
                #self.Environment.render()
            #self.Population[AgentIndex].Monitors["Reward Monitor"].Append(reward)
            RewardTotal += reward
            if done:
                break

        # Add all of the current actions for this episode to the AllActions list
        self.AllActions.append(Actions)
        #print(e)
        # TODO: Generalise this for multiple layers
        # Keep track of all of the spiking activity
        self.AllSpikesInputLayer.append(np.asarray(self.Network.Monitors["Input Layer"].Values))
        #print(self.Network.Monitors["Input Layer"].Values)
        self.AllSpikesOutputLayer.append(np.asarray(self.Network.Monitors["Output Layer"].Values))

        #if e == NumEpisodes - 1:
            #self.Network.Monitors["Input Layer"].Plot("Input Layer")
        self.Network.ResetStateVariables(True)

        self.AllTotalRewards.append(RewardTotal)
        #AvgReward = sum(self.AllTotalRewards)
        #print(AvgReward)

        #return (RewardTotal >= 30.0), RewardTotal

    def PrintToFile(self, TextFile, ItemToAppend):
        if os.path.exists(TextFile):
            f = open(TextFile, 'a')
            f.write(str(ItemToAppend) + "\n")
            f.close()
        else:
            f = open(TextFile, 'w')
            f.write(str(ItemToAppend) + "\n")
            f.close()

    def UpdateWeights(self):
        # Can retrieve the actions and the rewards from AllActions and AllTotalRewards

        EpisodeLength = int(len(self.AllActions[0]))
        NetworkTotalSteps = int(len(self.AllSpikesInputLayer[0]))
        #print (self.AllEpisodeRewards)
        #print (self.AllActions)
        #print(self.AllSpikesInputLayer[0])
        #self.Network.Monitors["Layer 1"]
        for t in range(len(self.AllSpikesInputLayer[0])):
            for i in range(self.AllSpikesInputLayer[0][0].size):

                if self.AllSpikesInputLayer[0][t][i] == 1:
                    forward = 0
                    back = 1
                    while True:
                        #print("Starting Update Weights")
                        if (t+forward >= NetworkTotalSteps or forward > 10):
                            break

                        # So if the post-synaptic spike was a result of the pre-synaptic spike then set
                        # the Long Term Potentiation change to 1
                        TargetLTP = np.where(self.AllSpikesOutputLayer[0][t+forward] == 1, 1, 0)
                        #print (TargetLTP)
                        # If the TargetLTP is 1 then trigger the strengthening of the weight
                        if (np.sum(TargetLTP) != 0):
                            self.Network.Connections[0].w[:,i] += np.where(TargetLTP == 1, self.LearningRate, 0)

                        else:
                            if (t-back < 0 or back > 10):
                                break
                            TargetLTD = np.where(self.AllSpikesOutputLayer[0][t-back] == 1, 1, 0)
                            if (np.sum(TargetLTD) != 0):
                                self.Network.Connections[0].w[:,i] -= np.where(TargetLTD == 1, self.LearningRate, 0)

                        forward += 1
                        back += 1

                        # If not check for LTD
                        # If LTD then trigger the weakening of the weight
                        # If neither event occured, look at one further timestep forward and back
        #self.Network.Connections[0].decay(0.05)
'''
    def UpdateWeightsAlternate(self, Success: bool):
        NetworkTotalSteps = int(len(self.AllSpikesInputLayer[0]))
        if Success:
            for t in range(len(self.AllSpikesInputLayer[0])):
                for i in range(self.AllSpikesInputLayer[0][0].size):

                    if self.AllSpikesInputLayer[0][t][i] == 1:
                        forward = 0
                        back = 1
                        while True:
                            if (t+forward >= NetworkTotalSteps or forward > 10):
                                break

                            # So if the post-synaptic spike was a result of the pre-synaptic spike then set
                            # the Long Term Potentiation change to 1
                            TargetLTP = np.where(self.AllSpikesOutputLayer[0][t+forward] == 1, 1, 0)

                            # If the TargetLTP is 1 then trigger the strengthening of the weight
                            if (np.sum(TargetLTP) != 0):
                                self.Network.Connections[1].w[:,i] += np.where(TargetLTP == 1, self.LearningRate * np.random.rand(), 0)
                                break
                            else:
                                if (t-back < 0 or back >= 10):
                                    break
                                TargetLTD = np.where(self.AllSpikesOutputLayer[0][t-back] == 1, 1, 0)
                                if (np.sum(TargetLTD) != 0):
                                    self.Network.Connections[1].w[:,i] -= np.where(TargetLTD == 1, self.LearningRate * np.random.rand(), 0)
                                    break
                            forward += 1
                            back += 1
        else:
            for t in range(len(self.AllSpikesInputLayer[0])):
                for i in range(self.AllSpikesInputLayer[0][0].size):

                    if self.AllSpikesInputLayer[0][t][i] == 1:
                        forward = 0
                        back = 1
                        while True:
                            if (t+forward >= NetworkTotalSteps or forward > 10):
                                break

                            # So if the post-synaptic spike was a result of the pre-synaptic spike then set
                            # the Long Term Potentiation change to 1
                            TargetLTP = np.where(self.AllSpikesOutputLayer[0][t+forward] == 1, 1, 0)

                            # If the TargetLTP is 1 then trigger the strengthening of the weight
                            if (np.sum(TargetLTP) != 0):
                                self.Network.Connections[1].w[:,i] -= np.where(TargetLTP == 1, self.LearningRate * np.random.rand(), 0)
                                break
                            else:
                                if (t-back < 0 or back >= 10):
                                    break
                                TargetLTD = np.where(self.AllSpikesOutputLayer[0][t-back] == 1, 1, 0)
                                if (np.sum(TargetLTD) != 0):
                                    self.Network.Connections[1].w[:,i] += np.where(TargetLTD == 1, self.LearningRate * np.random.rand(), 0)
                                    break
                            forward += 1
                            back += 1'''


# Copyright (C) 2020 Intel Corporation
#
# SPDX-License-Identifier: MIT

import logging as log
import numpy as np

from datumaro.components.extractor import (Transform, AnnotationType,
    DEFAULT_SUBSET_NAME)
from datumaro.components.cli_plugin import CliPlugin

NEAR_ZERO = 1e-7

class _TaskSpecificSplit(Transform, CliPlugin):
    _default_split = [('train', 0.5), ('val', 0.2), ('test', 0.3)]

    @classmethod
    def build_cmdline_parser(cls, **kwargs):
        parser = super().build_cmdline_parser(**kwargs)
        parser.add_argument('-s', '--subset', action='append',
            type=cls._split_arg, dest='splits',
            help="Subsets in the form: '<subset>:<ratio>' "
                "(repeatable, default: %s)" % dict(cls._default_split))
        parser.add_argument('--seed', type=int, help="Random seed")
        return parser

    @staticmethod
    def _split_arg(s):
        parts = s.split(':')
        if len(parts) != 2:
            import argparse
            raise argparse.ArgumentTypeError()
        return (parts[0], float(parts[1]))

    def __init__(self, dataset, splits, seed):
        super().__init__(dataset)

        if splits is None:
            splits = self._default_split

        snames, sratio = self._validate_splits(splits)

        self._snames = snames
        self._sratio = sratio

        self._seed = seed

        self._subsets = {"train", "val", "test"}  # output subset names
        self._parts = []
        self._length = "parent"

        self._initialized = False

    def _set_parts(self, by_splits):
        self._parts = []
        for subset in self._subsets:
            self._parts.append((set(by_splits[subset]), subset))

    @staticmethod
    def _get_uniq_annotations(dataset):
        annotations = []
        for item in dataset:
            labels = [a for a in item.annotations
                if a.type == AnnotationType.label]
            if len(labels) != 1:
                raise Exception("Item '%s' contains %s labels, "
                    "but exactly one is expected" % (item.id, len(labels)))
            annotations.append(labels[0])
        return annotations

    @staticmethod
    def _validate_splits(splits, valid=None):
        snames = []
        ratios = []
        if valid is None:
            valid = ["train", "val", "test"]
        for subset, ratio in splits:
            assert subset in valid, \
                "Subset name must be one of %s, but got %s" % (valid, subset)
            assert 0.0 <= ratio and ratio <= 1.0, \
                "Ratio is expected to be in the range " \
                "[0, 1], but got %s for %s" % (ratio, subset)
            snames.append(subset)
            ratios.append(float(ratio))
        ratios = np.array(ratios)

        total_ratio = np.sum(ratios)
        if not abs(total_ratio - 1.0) <= NEAR_ZERO:
            raise Exception(
                "Sum of ratios is expected to be 1, got %s, which is %s"
                % (splits, total_ratio)
            )
        return snames, ratios

    @staticmethod
    def _get_required(ratio):
        min_value = np.max(ratio)
        for i in ratio:
            if NEAR_ZERO < i and i < min_value:
                min_value = i
        required = int(np.around(1.0) / min_value)
        return required

    @staticmethod
    def _get_sections(dataset_size, ratio):
        n_splits = [int(np.around(dataset_size * r)) for r in ratio[:-1]]
        n_splits.append(dataset_size - np.sum(n_splits))

        # if there are splits with zero samples even if ratio is not 0,
        # borrow one from the split who has one or more.
        for ii, num_split in enumerate(n_splits):
            if num_split == 0 and NEAR_ZERO < ratio[ii]:
                midx = np.argmax(n_splits)
                if n_splits[midx] > 0:
                    n_splits[ii] += 1
                    n_splits[midx] -= 1
        sections = np.add.accumulate(n_splits[:-1])
        return sections

    @staticmethod
    def _group_by_attr(items):
        """
        Args:
            items: list of (idx, ann). ann is the annotation from Label object.
        Returns:
            by_attributes: dict of { combination-of-attrs : list of index }
        """
        # group by attributes
        by_attributes = dict()
        for idx, ann in items:
            attributes = tuple(sorted(ann.attributes.items()))
            if attributes not in by_attributes:
                by_attributes[attributes] = []
            by_attributes[attributes].append(idx)
        return by_attributes

    def _split_by_attr(self, datasets, snames, ratio, out_splits,
            dataset_key=None):
        required = self._get_required(ratio)
        if dataset_key is None:
            dataset_key = "label"
        for key, items in datasets.items():
            np.random.shuffle(items)
            by_attributes = self._group_by_attr(items)
            for attributes, indice in by_attributes.items():
                gname = "%s: %s, attrs: %s" % (dataset_key, key, attributes)
                splits = self._split_indice(indice, gname, ratio, required)
                for subset, split in zip(snames, splits):
                    if 0 < len(split):
                        out_splits[subset].extend(split)

    def _split_indice(self, indice, group_name, ratio, required):
        filtered_size = len(indice)
        if filtered_size < required:
            log.warning("Not enough samples for a group, '%s'" % group_name)
        sections = self._get_sections(filtered_size, ratio)
        splits = np.array_split(indice, sections)
        return splits

    def _find_split(self, index):
        for subset_indices, subset in self._parts:
            if index in subset_indices:
                return subset
        return DEFAULT_SUBSET_NAME  # all the possible remainder --> default

    def _split_dataset(self):
        raise NotImplementedError()

    def __iter__(self):
        # lazy splitting
        if self._initialized is False:
            self._split_dataset()
            self._initialized = True
        for i, item in enumerate(self._extractor):
            yield self.wrap_item(item, subset=self._find_split(i))

class ClassificationSplit(_TaskSpecificSplit):
    """
    Splits dataset into train/val/test set in class-wise manner. |n
    Splits dataset images in the specified ratio, keeping the initial class
    distribution.|n
    |n
    Notes:|n
    - Each image is expected to have only one Label|n
    - If Labels also have attributes, also splits by attribute values.|n
    - If there is not enough images in some class or attributes group,
      the split ratio can't be guaranteed.|n
    |n
    Example:|n
    |s|s%(prog)s --subset train:.5 --subset val:.2 --subset test:.3
    """
    def __init__(self, dataset, splits, seed=None):
        """
        Parameters
        ----------
        dataset : Dataset
        splits : list
            A list of (subset(str), ratio(float))
            Subset is expected to be one of ["train", "val", "test"].
            The sum of ratios is expected to be 1.
        seed : int, optional
        """
        super().__init__(dataset, splits, seed)

    def _split_dataset(self):
        np.random.seed(self._seed)

        # support only single label for a DatasetItem
        # 1. group by label
        by_labels = dict()
        annotations = self._get_uniq_annotations(self._extractor)
        for idx, ann in enumerate(annotations):
            label = getattr(ann, 'label', None)
            if label not in by_labels:
                by_labels[label] = []
            by_labels[label].append((idx, ann))

        by_splits = dict()
        for subset in self._subsets:
            by_splits[subset] = []

        # 2. group by attributes
        self._split_by_attr(by_labels, self._snames, self._sratio, by_splits)
        self._set_parts(by_splits)

class ReidentificationSplit(_TaskSpecificSplit):
    """
    Splits a dataset for re-identification task.|n
    Produces a split with a specified ratio of images, avoiding having same
    labels in different subsets.|n
    |n
    In this task, the test set should consist of images of unseen
    people or objects during the training phase. |n
    This function splits a dataset in the following way:|n
    1. Splits the dataset into 'train + val' and 'test' sets|n
    |s|sbased on person or object ID.|n
    2. Splits 'test' set into 'test-gallery' and 'test-query' sets|n
    |s|sin class-wise manner.|n
    3. Splits the 'train + val' set into 'train' and 'val' sets|n
    |s|sin the same way.|n
    The final subsets would be
    'train', 'val', 'test-gallery' and 'test-query'. |n
    |n
    Notes:|n
    - Each image is expected to have a single Label|n
    - Object ID can be described by Label, or by attribute (--attr parameter)|n
    - The splits of the test set are controlled by '--query' parameter. |n
    |s|sGallery ratio would be 1.0 - query.|n
    |n
    Example: split a dataset in the specified ratio, split the test set|n
    |s|s|s|sinto gallery and query in 1:1 ratio|n
    |s|s%(prog)s --subset train:.5 --subset val:.2 --subset test:.3 --query .5|n
    Example: use 'person_id' attribute for splitting|n
    |s|s%(prog)s --attr person_id
    """

    _default_query_ratio = 0.5

    @classmethod
    def build_cmdline_parser(cls, **kwargs):
        parser = super().build_cmdline_parser(**kwargs)
        parser.add_argument('--query', type=float,
            help="Query ratio in the test set (default: %.3f)"
            % cls._default_query_ratio)
        parser.add_argument('--attr', type=str, dest='attr_for_id',
            help="Attribute name representing the ID (default: use label)")
        return parser

    def __init__(self, dataset, splits, query=None,
            attr_for_id=None, seed=None):
        """
        Parameters
        ----------
        dataset : Dataset
        splits : list
            A list of (subset(str), ratio(float))
            Subset is expected to be one of ["train", "val", "test"].
            The sum of ratios is expected to be 1.
        query : float
            The ratio of 'test-query' set.
            The ratio of 'test-gallery' set would be 1.0 - query.
        attr_for_id: str
            attribute name representing the person/object id.
            if this is not specified, label would be used.
        seed : int, optional
        """
        super().__init__(dataset, splits, seed)

        if query is None:
            query = self._default_query_ratio

        assert 0.0 <= query and query <= 1.0, \
            "Query ratio is expected to be in the range " \
            "[0, 1], but got %f" % query
        test_splits = [('test-query', query), ('test-gallery', 1.0 - query)]

        # reset output subset names
        self._subsets = {"train", "val", "test-gallery", "test-query"}
        self._test_splits = test_splits
        self._attr_for_id = attr_for_id

    def _split_dataset(self):
        np.random.seed(self._seed)

        id_snames, id_ratio = self._snames, self._sratio

        attr_for_id = self._attr_for_id
        dataset = self._extractor

        # group by ID(attr_for_id)
        by_id = dict()
        annotations = self._get_uniq_annotations(dataset)
        if attr_for_id is None:  # use label
            for idx, ann in enumerate(annotations):
                ID = getattr(ann, 'label', None)
                if ID not in by_id:
                    by_id[ID] = []
                by_id[ID].append((idx, ann))
        else:  # use attr_for_id
            for idx, ann in enumerate(annotations):
                attributes = dict(ann.attributes.items())
                assert attr_for_id in attributes, \
                    "'%s' is expected as an attribute name" % attr_for_id
                ID = attributes[attr_for_id]
                if ID not in by_id:
                    by_id[ID] = []
                by_id[ID].append((idx, ann))

        required = self._get_required(id_ratio)
        if len(by_id) < required:
            log.warning("There's not enough IDs, which is %s, "
                "so train/val/test ratio can't be guaranteed."
                % len(by_id)
            )

        # 1. split dataset into trval and test
        #    IDs in test set should not exist in train/val set.
        test = id_ratio[id_snames.index("test")] if "test" in id_snames else 0
        if NEAR_ZERO < test:  # has testset
            split_ratio = np.array([test, 1.0 - test])
            IDs = list(by_id.keys())
            np.random.shuffle(IDs)
            sections = self._get_sections(len(IDs), split_ratio)
            splits = np.array_split(IDs, sections)
            testset = {pid: by_id[pid] for pid in splits[0]}
            trval = {pid: by_id[pid] for pid in splits[1]}

            # follow the ratio of datasetitems as possible.
            # naive heuristic: exchange the best item one by one.
            expected_count = int(len(self._extractor) * split_ratio[0])
            testset_total = int(np.sum([len(v) for v in testset.values()]))
            self._rebalancing(testset, trval, expected_count, testset_total)
        else:
            testset = dict()
            trval = by_id

        by_splits = dict()
        for subset in self._subsets:
            by_splits[subset] = []

        # 2. split 'test' into 'test-gallery' and 'test-query'
        if 0 < len(testset):
            test_snames = []
            test_ratio = []
            for sname, ratio in self._test_splits:
                test_snames.append(sname)
                test_ratio.append(float(ratio))

            self._split_by_attr(testset, test_snames, test_ratio, by_splits,
                dataset_key=attr_for_id)

        # 3. split 'trval' into  'train' and 'val'
        trval_snames = ["train", "val"]
        trval_ratio = []
        for subset in trval_snames:
            if subset in id_snames:
                val = id_ratio[id_snames.index(subset)]
            else:
                val = 0.0
            trval_ratio.append(val)
        trval_ratio = np.array(trval_ratio)
        total_ratio = np.sum(trval_ratio)
        if total_ratio < NEAR_ZERO:
            trval_splits = list(zip(["train", "val"], trval_ratio))
            log.warning("Sum of ratios is expected to be positive, "
                "got %s, which is %s"
                % (trval_splits, total_ratio)
            )
        else:
            trval_ratio /= total_ratio  # normalize
            self._split_by_attr(trval, trval_snames, trval_ratio, by_splits,
                dataset_key=attr_for_id)

        self._set_parts(by_splits)

    @staticmethod
    def _rebalancing(test, trval, expected_count, testset_total):
        diffs = dict()
        for id_test, items_test in test.items():
            count_test = len(items_test)
            for id_trval, items_trval in trval.items():
                count_trval = len(items_trval)
                diff = count_trval - count_test
                if diff == 0:
                    continue  # exchange has no effect
                if diff not in diffs:
                    diffs[diff] = [(id_test, id_trval)]
                else:
                    diffs[diff].append((id_test, id_trval))
        if len(diffs) == 0:  # nothing would be changed by exchange
            return

        exchanges = []
        while True:
            target_diff = expected_count - testset_total
            # find nearest diff.
            keys = np.array(list(diffs.keys()))
            idx = (np.abs(keys - target_diff)).argmin()
            nearest = keys[idx]
            if abs(target_diff) <= abs(target_diff - nearest):
                break
            choice = np.random.choice(range(len(diffs[nearest])))
            id_test, id_trval = diffs[nearest][choice]
            testset_total += nearest
            new_diffs = dict()
            for diff, IDs in diffs.items():
                new_list = []
                for id1, id2 in IDs:
                    if id1 == id_test or id2 == id_trval:
                        continue
                    new_list.append((id1, id2))
                if 0 < len(new_list):
                    new_diffs[diff] = new_list
            diffs = new_diffs
            exchanges.append((id_test, id_trval))

        # exchange
        for id_test, id_trval in exchanges:
            test[id_trval] = trval.pop(id_trval)
            trval[id_test] = test.pop(id_test)

class DetectionSplit(_TaskSpecificSplit):
    """
    Splits a dataset into train/val/test subsets for detection task,
    using object annotations as a basis for splitting.|n
    Tries to produce an image split with the specified ratio, keeping the
    initial distribution of class objects.|n
    |n
    In a detection dataset, each image can have multiple object annotations -
    instance bounding boxes. Since an image shouldn't be included
    in multiple subsets at the same time, and image annotations
    shoudln't be split, in general, dataset annotations are unlikely to be split
    exactly in the specified ratio. |n
    This split tries to split dataset images as close as possible
    to the specified ratio, keeping the initial class distribution.|n
    |n
    Notes:|n
    - Each image is expected to have one or more Bbox annotations.|n
    - Only Bbox annotations are considered.|n
    |n
    Example: split dataset so that each object class annotations were split|n
    |s|s|s|sin the specified ratio between subsets|n
    |s|s%(prog)s --subset train:.5 --subset val:.2 --subset test:.3
    """
    def __init__(self, dataset, splits, seed=None):
        """
        Parameters
        ----------
        dataset : Dataset
        splits : list
            A list of (subset(str), ratio(float))
            Subset is expected to be one of ["train", "val", "test"].
            The sum of ratios is expected to be 1.
        seed : int, optional
        """
        super().__init__(dataset, splits, seed)

    @staticmethod
    def _group_by_bbox_labels(dataset):
        by_labels = dict()
        for idx, item in enumerate(dataset):
            bbox_anns = [a for a in item.annotations
                if a.type == AnnotationType.bbox]
            assert 0 < len(bbox_anns), \
                "Expected more than one bbox annotation in the dataset"
            for ann in bbox_anns:
                label = getattr(ann, 'label', None)
                if label not in by_labels:
                    by_labels[label] = [(idx, ann)]
                else:
                    by_labels[label].append((idx, ann))
        return by_labels

    def _split_dataset(self):
        np.random.seed(self._seed)

        subsets, sratio = self._snames, self._sratio

        # 1. group by bbox label
        by_labels = self._group_by_bbox_labels(self._extractor)

        # 2. group by attributes
        by_combinations = dict()
        for label, items in by_labels.items():
            by_attributes = self._group_by_attr(items)
            for attributes, indice in by_attributes.items():
                gname = "label: %s, attributes: %s" % (label, attributes)
                by_combinations[gname] = indice

        # total number of GT samples per label-attr combinations
        n_combs = {k: len(v) for k, v in by_combinations.items()}

        # 3-1. initially count per-image GT samples
        scores_all = {}
        init_scores = {}
        for idx, _ in enumerate(self._extractor):
            counts = {k: v.count(idx) for k, v in by_combinations.items()}
            scores_all[idx] = counts
            init_scores[idx] = np.sum(
                [v / n_combs[k] for k, v in counts.items()]
            )

        by_splits = dict()
        for sname in self._subsets:
            by_splits[sname] = []

        total = len(self._extractor)
        target_size = dict()
        expected = []  # expected numbers of per split GT samples
        for sname, ratio in zip(subsets, sratio):
            target_size[sname] = total * ratio
            expected.append(
                (sname, {k: v * ratio for k, v in n_combs.items()})
            )

        # functions for keep the # of annotations not exceed the expected num
        def compute_penalty(counts, n_combs):
            p = 0
            for k, v in counts.items():
                p += max(0, (v / n_combs[k]) - 1.0)
            return p

        def update_nc(counts, n_combs):
            for k, v in counts.items():
                n_combs[k] = max(0, n_combs[k] - v)
                if n_combs[k] == 0:
                    n_combs[k] = -1
            return n_combs

        # 3-2. assign each DatasetItem to a split, one by one
        for idx, _ in sorted(
            init_scores.items(), key=lambda item: item[1], reverse=True
        ):
            counts = scores_all[idx]

            # shuffling split order to add randomness
            # when two or more splits have the same penalty value
            np.random.shuffle(expected)

            pp = []
            for sname, nc in expected:
                if target_size[sname] <= len(by_splits[sname]):
                    # the split has enough images,
                    # stop adding more images to this split
                    pp.append(1e08)
                else:
                    # compute penalty based on the number of GT samples
                    # added in the split
                    pp.append(compute_penalty(counts, nc))

            # we push an image to a split with the minimum penalty
            midx = np.argmin(pp)

            sname, nc = expected[midx]
            by_splits[sname].append(idx)
            update_nc(counts, nc)

        self._set_parts(by_splits)


from django.shortcuts import render
from . models import Check
from django.http import HttpResponse
from django.views.decorators.csrf import csrf_exempt

@csrf_exempt
def main(request):
    if request.method == 'POST':
        # image = request.POST['image']
        objects = request.POST['objects']
        time = request.POST['time']
        check = Check(Objects = objects, Time = time)
        check.save()
        return HttpResponse("Created!")
    else:
        context = {'checks': Check.objects.all()}
        return render(request, 'food/index.html', context)


import matplotlib.image as mpimg

class getmap:
	def __init__(self, filepath, sep=' ', format='png'):
		self.image = []
		self.maxVal = 1
		self.minVal = 0

		if format == 'png':
			img = mpimg.imread(filepath, format='png')

			x = img.shape[0]
			y = img.shape[1]

			for i in range(x):
				self.image.append([])
				for j in range(y):
					newValue = int(img[i][j][0] == 0)
					self.image[i].append(newValue)
					self.maxVal = max(self.maxVal, newValue)
					self.minVal = min(self.minVal, newValue)
		elif format == 'txt':
			with open(filepath, 'r') as file:
				for line in file:
					newLine = list(map(int, line.strip().split(sep)))
					self.image.append(newLine)
					self.maxVal = max(self.maxVal, max(newLine))
					self.minVal = min(self.minVal, min(newLine))


# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals

import sys
from argparse import ArgumentParser
from warnings import simplefilter
from warnings import warn

from bravado.client import SwaggerClient
from cookiecutter.log import configure_logger
from six import iteritems

from swagger_code_generator.generator_metaclass import MetaClassRegistry
from swagger_code_generator.util.cli_parser import path_or_uri_type
from swagger_code_generator.util.cli_parser import path_type

def main(argv=()):
    simplefilter('error')   # Ensure that in case if warning the tool execution is interrupted

    cli_parser = ArgumentParser('Swagger Models Class Generator')

    cli_parser.add_argument(
        '-s', '--swagger',
        type=path_or_uri_type,
        help='Root Swagger Spec URI (ie: file://.., http[s]://.., etc.)',
    )

    cli_parser.add_argument(
        '-d', '--destination',
        type=path_type,
        help='Destination directory of the generated files',
    )
    cli_parser.add_argument(
        '-o', '--operations',
        action='store_true',
        help='Generate Operations Handling',
    )
    cli_parser.add_argument(
        '-v', '--verbose',
        action='store_true',
        help='Verbose Logging',
    )

    if not MetaClassRegistry.registry_classes():
        warn('No language classes registered to {}'.format(MetaClassRegistry))

    subparsers = cli_parser.add_subparsers(help='Language to generate')
    for registry_class in MetaClassRegistry.registry_classes():
        subparser = subparsers.add_parser(
            registry_class.language(),
            **{k: v for k, v in iteritems({'help': registry_class.subparser_help()}) if v}
        )
        subparser.set_defaults(func=registry_class.run, validate=registry_class.validate_args, subparser=subparser)
        registry_class.set_subparser(subparser)

    args = cli_parser.parse_args(argv)

    if not args.swagger:
        cli_parser.error('swagger is a required parameter')

    if not args.destination:
        cli_parser.error('destination is a required parameter')

    args.validate(args, args.subparser.error)

    swagger_client = SwaggerClient.from_url(spec_url=args.swagger)

    if args.verbose:
        configure_logger()

    args.func(swagger_client.swagger_spec, args)

if __name__ == '__main__':
    exit(main(sys.argv[1:]))


import ConfigParser
import logging
import logging.config
import datetime
from dateutil.relativedelta import *

class Setting:
	'''Read config info with ConfigParser'''
	def __init__(self):
		logger = logging.getLogger(__name__)
		logger.debug('bgn...')

		self.today = datetime.datetime.now()

		# read the general setting
		setting = ConfigParser.ConfigParser()
		setting.read("../setting.ini")

		# database
		self.steve_id_mkv = setting.get('database', 'steve_id_mkv')
		self.steve_id_edw = setting.get('database', 'steve_id_edw')
		self.steve_ldap = setting.get('database', 'steve_ldap')
		self.edw = setting.get('database', 'edw')
		self.mkv = setting.get('database', 'mkv')

		# hostname
		self.steve_id_unix = setting.get('unix_server', 'steve_id_unix')
		self.ciahost = setting.get('unix_server', 'ciahost')
		self.prchost = setting.get('unix_server', 'prchost')

		# read the VY setting
		setting = ConfigParser.ConfigParser()
		setting.read("setting.ini")
		self.cut_date_yyyy = setting.get('date', 'cut_date_yyyy')
		self.cut_date_mm = setting.get('date', 'cut_date_mm')
		self.cut_date = datetime.date(int(self.cut_date_yyyy), int(self.cut_date_mm), 1)

		# cut_friday = 'first_friday_after_cut_date', 4 is Friday
		self.cut_friday = self.cut_date + relativedelta(weekday=4)
		# previous Friday from today but not today.
		self.end_friday = self.today + relativedelta(days=-1, weekday=FR(-1))
		# the friday to begin with in VY
		self.bgn_friday = self.end_friday + datetime.timedelta(weeks=-51)

		self.show_bp = setting.get('chart_line', 'show_bp')
		logger.debug('end.')


from prody import *
from matplotlib.pylab import *
from Bio.Blat import NCBIWWW
from Bio import SeqIO
from Bio.ExPASy import ScanProsite
from bioservices.kegg import KEGG
import json
k = KEGG()

"""
*How annotations are queried*

blast -> seq
pfam -> seq
prosite -> seq
kegg -> original gene annotation
go -> pfam and prosite ontology later
"""

DNA=open('UP000006737.fasta.txt', 'r').read().split('>')
queries = list(SeqIO.parse("project_seqs.fasta", format="fasta"))
dna = DNA[301:401] #first index is blank from split so index is "normal" for sequences
dna = dna[50:] #can be used to pick which sequences to query
queries = queries[50:] #can be used to pick which sequences to query
fasta = ['>'+dna[i] for i in range(len(dna))]
dna = [''.join(v.split('\n')[1:-1]) for v in dna]
gn = [''.join(v.split('\n')[0]).split('GN=')[1].split()[0] for v in fasta]

def blast(query):
    try:
        result_handle = NCBIWWW.qblast("blastp", "nr", query.seq)
        result_handle = result_handle.read()
        return result_handle
    except:
        return 'noHomology'

def pfam(seq):
    try:
        matches = searchPfam(seq)
        return matches
    except:
        return 'noHomology'

def prosite(seq):
    try:
        handle = ScanProsite.scan(seq=seq, lowscore=10) #add lowscore=# to get the next # 'possible' additional hits
        result = ScanProsite.read(handle)
        if result == []:
            return 'noHomology'
        else:
            return result
    except:
        return 'noHomology'

def kegg(gn):
    try:
        hit = k.get_pathway_by_gene(gn, "acb")
        if not hit:
            return 'noHomology'
        else:
            return hit
    except:
        return 'noHomology'

def go(go_id):
    "Will be inputed later"
    pass

memo = {num:{'blast':blast(queries[num]), 'pfam':pfam(seq), 'prosite':prosite(seq), 'kegg':kegg(gn[num]), 'go':'future go data'} for num, seq in list(enumerate(dna[:]))}

with open('data50_99.json', 'w') as outfile:
    json.dump(memo, outfile)

print('second half done')


import logging

from rq.timeouts import JobTimeoutException

from django.conf import settings
from django.db import connection

from django_push.subscriber.models import Subscription

from ..tasks import raven, enqueue

logger = logging.getLogger('feedupdater')

@raven
def update_feed(feed_url, use_etags=True):
    from .models import UniqueFeed
    try:
        UniqueFeed.objects.update_feed(feed_url, use_etags)
    except JobTimeoutException:
        feed = UniqueFeed.objects.get(url=feed_url)
        feed.backoff()
        feed.save()
        logger.info("Job timed out, backing off %s to %s" % (
            feed.url, feed.backoff_factor,
        ))
    close_connection()

@raven
def read_later(entry_pk):
    from .models import Entry  # circular imports
    Entry.objects.get(pk=entry_pk).read_later()
    close_connection()

@raven
def update_unique_feed(feed_url):
    from .models import UniqueFeed, Feed
    feed, created = UniqueFeed.objects.get_or_create(
        url=feed_url,
        defaults={'subscribers': 1},
    )
    if created:
        enqueue(update_favicon, args=[feed_url], queue='high')
    else:
        feed.subscribers = Feed.objects.filter(url=feed_url).count()
        feed.save(update_fields=['subscribers'])

@raven
def update_favicon(feed_url):
    from .models import Favicon
    Favicon.objects.update_favicon(feed_url)

@raven
def subscribe(topic_url, hub_url):
    Subscription.objects.subscribe(topic_url, hub_url)

def close_connection():
    """Close the connection only if not in eager mode"""
    if hasattr(settings, 'RQ'):
        if not settings.RQ.get('eager', False):
            connection.close()