# *******************************************************************************
# Copyright (C) 2020 INAF
#
# This software is distributed under the terms of the BSD-3-Clause license
#
# Authors:
# Ambra Di Piano <ambra.dipiano@inaf.it>
# *******************************************************************************

import time
import sys
import os
texp = sys.argv[1]
first = sys.argv[2]

# start timing
t = time.time()
clock0 = time.time()
#from astropy.coordinates import SkyCoord
from astropy import units as u
from astropy.coordinates import SkyCoord
from gammapy.analysis import Analysis, AnalysisConfig
from gammapy.data import EventList, GTI, Observation, Observations
from gammapy.irf import load_cta_irfs
from gammapy.modeling import Fit
from gammapy.modeling.models import PowerLawSpectralModel, SkyModel, PointSpatialModel
timport = time.time() - t
print(f'Imports : {timport} s\n')

t = time.time()
rootpath = str(os.path.dirname(os.path.abspath(__file__))).replace('cta-sag-sci/RTAscience/timing', '')
caldb = f'{rootpath}/caldb/data/cta/prod3b-v2/bcf/South_z20_0.5h/irf_file.fits'
irfs = load_cta_irfs(caldb)
filename = f'{rootpath}/DATA/obs/crab/crab_offax_texp{texp}s_n01.fits'
obs_id = 1
print(f'Fits: {filename.replace(rootpath, "")}\n')
tsetup = time.time() - t
print(f'Setup : {tsetup} s\n')

# read phlist
t = time.time()
events = EventList.read(filename, hdu='EVENTS')
# get GTI
gti = GTI.read(filename, hdu='GTI')
# get pointing
pointing = events.pointing_radec
# create observation
observation = Observation.create(
    pointing=pointing, obs_id=f'{obs_id:02d}', tstart=gti.table['START'] * u.s,
    tstop=gti.table['STOP'] * u.s, irfs=irfs, reference_time=gti.time_ref)
observation._events = events
#print(observation.gti)
observations = Observations()
observations.append(observation)
# fix pointing info
observation.fixed_pointing_info
# target
#target = SkyCoord(pointing.ra, pointing.dec - 0.5 * u.deg, unit='deg', frame='icrs')
target = {'ra': 83.6331, 'dec': 22.0145}
tobs = time.time() - t
print(f'Create observation : {tobs} s\n')

# configure a 1d analysis
t = time.time()
config_1d = AnalysisConfig()
config_1d.general.log = {'level': 'warning'}
config_1d.datasets.type = "1d"
config_1d.datasets.stack = False
# define the ON region and make sure that PSF leakage is corrected
config_1d.datasets.on_region = dict(frame="icrs", lon='%s deg' %target['ra'], lat='%s deg' %target['dec'], radius='0.1 deg')
config_1d.datasets.containment_correction = True
# background
config_1d.datasets.background=dict(method="reflected", exclusion=None)
#config_1d.datasets.safe_mask.methods = ["edisp-bias"]
# define the energy binning for the spectra
config_1d.datasets.geom.axes.energy = dict(min='0.05 TeV', max='20 TeV', nbins=30)
config_1d.datasets.geom.axes.energy_true = dict(min='0.03 TeV', max='30 TeV', nbins=40)
config_1d.datasets.geom.selection.offset_max = '2.5 deg'
# fit
#config_1d.fit.fit_range = dict(min='0.03 TeV', max='30 TeV')
#config_1d.flux_points.energy = dict(min='0.03 TeV', max='30 TeV', nbins=3)
#config_1d.flux_points.source = 'Crab'
# write
#config_1d.write("config1d.yaml", overwrite=True)
#config_1d = AnalysisConfig.read("config1d.yaml")
tconf = time.time() - t
print(f'Configuration : {tconf} s\n')
print(target)
# instantiate data reduction passing directly the config object
t = time.time()
analysis_1d = Analysis(config_1d)
analysis_1d.observations = observations
analysis_1d.get_datasets()
tred = time.time() - t
print(f'Data Reduction : {tred} s\n')

# statistics
t = time.time()
stats = analysis_1d.datasets.info_table()
print(stats['sqrt_ts'], '\n')
tstat = time.time() - t
print(f'Statistics : {tstat} s\n')

# prepare models
t = time.time()
stacked_1d = analysis_1d.datasets.stack_reduce(name="stacked")
target = SkyCoord(target['ra'], target['dec'], unit='deg', frame='icrs')
spatial_model = PointSpatialModel(lon_0=target.ra, lat_0=target.dec, frame="icrs")
spectral_model = PowerLawSpectralModel(index=2.48, amplitude=2e-12 * u.Unit("1 / (cm2 s TeV)"))
spectral_model.parameters['index'].frozen = True
spatial_model.parameters['lon_0'].frozen = True
spatial_model.parameters['lat_0'].frozen = True
sky_model = SkyModel(spatial_model=spatial_model, spectral_model=spectral_model, name="Crab")
stacked_1d.models = sky_model
tmodel = time.time() - t
print(f'Modelling : {tmodel} s\n')

# fitting
t = time.time()
fit_1d = Fit([stacked_1d])
result_1d = fit_1d.run()
#print(result_1d.parameters.to_table(), '\n')
tfit = time.time() - t
print(f'Fitting : {tfit} s\n')

# flux
t = time.time()
phflux_err = spectral_model.integral_error(0.05 * u.TeV, 20 * u.TeV)
print(f'\nPH-FLUX {phflux_err.value[0]} +/- {phflux_err.value[1]}')
tflux = time.time() - t
print(f'\nFlux : {tflux} s\n')

ttotal = time.time() - clock0
print(f'Total time: {ttotal} s\n')
print('\n\n-----------------------------------------------------\n\n')

logname = f'{rootpath}/DATA/outputs/crab/gammapy1d_binned_fit.csv'
row = f'{texp} {stats["sqrt_ts"][0]} {phflux_err.value[0]} {phflux_err.value[1]} {ttotal} {timport} {tsetup} {tconf} {tred} {tstat} {tmodel} {tfit} {tflux}\n'
if first == 'True':
    hdr = 'texp sqrt_ts flux flux_err ttotal timport tsetup tobs tconf tred tstat tmodel tfit tflux\n'
    log = open(logname, 'w+')
    log.write(hdr)
    log.write(row)
    log.close()
else:
    log = open(logname, 'a')
    log.write(row)
    log.close()

print (row)


import operator

import numpy as np
import numpy.core.umath_tests as ut

from Quaternions import Quaternions

class Animation:
    """
    Animation is a numpy-like wrapper for animation data

    Animation data consists of several arrays consisting
    of F frames and J joints.

    The animation is specified by

        rotations : (F, J) Quaternions | Joint Rotations
        positions : (F, J, 3) ndarray  | Joint Positions

    The base pose is specified by

        orients   : (J) Quaternions    | Joint Orientations
        offsets   : (J, 3) ndarray     | Joint Offsets

    And the skeletal structure is specified by

        parents   : (J) ndarray        | Joint Parents
    """

    def __init__(self, rotations, positions, orients, offsets, parents):

        self.rotations = rotations
        self.positions = positions
        self.orients   = orients
        self.offsets   = offsets
        self.parents   = parents

    def __op__(self, op, other):
        return Animation(
            op(self.rotations, other.rotations),
            op(self.positions, other.positions),
            op(self.orients, other.orients),
            op(self.offsets, other.offsets),
            op(self.parents, other.parents))

    def __iop__(self, op, other):
        self.rotations = op(self.roations, other.rotations)
        self.positions = op(self.roations, other.positions)
        self.orients   = op(self.orients, other.orients)
        self.offsets   = op(self.offsets, other.offsets)
        self.parents   = op(self.parents, other.parents)
        return self

    def __sop__(self, op):
        return Animation(
            op(self.rotations),
            op(self.positions),
            op(self.orients),
            op(self.offsets),
            op(self.parents))

    def __add__(self, other): return self.__op__(operator.add, other)
    def __sub__(self, other): return self.__op__(operator.sub, other)
    def __mul__(self, other): return self.__op__(operator.mul, other)
    def __div__(self, other): return self.__op__(operator.div, other)

    def __abs__(self): return self.__sop__(operator.abs)
    def __neg__(self): return self.__sop__(operator.neg)

    def __iadd__(self, other): return self.__iop__(operator.iadd, other)
    def __isub__(self, other): return self.__iop__(operator.isub, other)
    def __imul__(self, other): return self.__iop__(operator.imul, other)
    def __idiv__(self, other): return self.__iop__(operator.idiv, other)

    def __len__(self): return len(self.rotations)

    def __getitem__(self, k):
        if isinstance(k, tuple):
            return Animation(
                self.rotations[k],
                self.positions[k],
                self.orients[k[1:]],
                self.offsets[k[1:]],
                self.parents[k[1:]])
        else:
            return Animation(
                self.rotations[k],
                self.positions[k],
                self.orients,
                self.offsets,
                self.parents)

    def __setitem__(self, k, v):
        if isinstance(k, tuple):
            self.rotations.__setitem__(k, v.rotations)
            self.positions.__setitem__(k, v.positions)
            self.orients.__setitem__(k[1:], v.orients)
            self.offsets.__setitem__(k[1:], v.offsets)
            self.parents.__setitem__(k[1:], v.parents)
        else:
            self.rotations.__setitem__(k, v.rotations)
            self.positions.__setitem__(k, v.positions)
            self.orients.__setitem__(k, v.orients)
            self.offsets.__setitem__(k, v.offsets)
            self.parents.__setitem__(k, v.parents)

    @property
    def shape(self): return (self.rotations.shape[0], self.rotations.shape[1])

    def copy(self): return Animation(
        self.rotations.copy(), self.positions.copy(),
        self.orients.copy(), self.offsets.copy(),
        self.parents.copy())

    def repeat(self, *args, **kw):
        return Animation(
            self.rotations.repeat(*args, **kw),
            self.positions.repeat(*args, **kw),
            self.orients, self.offsets, self.parents)

    def ravel(self):
        return np.hstack([
            self.rotations.log().ravel(),
            self.positions.ravel(),
            self.orients.log().ravel(),
            self.offsets.ravel()])

    @classmethod
    def unravel(clas, anim, shape, parents):
        nf, nj = shape
        rotations = anim[nf*nj*0:nf*nj*3]
        positions = anim[nf*nj*3:nf*nj*6]
        orients   = anim[nf*nj*6+nj*0:nf*nj*6+nj*3]
        offsets   = anim[nf*nj*6+nj*3:nf*nj*6+nj*6]
        return cls(
            Quaternions.exp(rotations), positions,
            Quaternions.exp(orients), offsets,
            parents.copy())

def transforms_local(anim):
    """
    Computes Animation Local Transforms

    As well as a number of other uses this can
    be used to compute global joint transforms,
    which in turn can be used to compete global
    joint positions

    Parameters
    ----------

    anim : Animation
        Input animation

    Returns
    -------

    transforms : (F, J, 4, 4) ndarray

        For each frame F, joint local
        transforms for each joint J
    """

    transforms = anim.rotations.transforms()
    transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1)
    transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2)
    transforms[:,:,0:3,3] = anim.positions
    transforms[:,:,3:4,3] = 1.0
    return transforms

def transforms_multiply(t0s, t1s):
    """
    Transforms Multiply

    Multiplies two arrays of animation transforms

    Parameters
    ----------

    t0s, t1s : (F, J, 4, 4) ndarray
        Two arrays of transforms
        for each frame F and each
        joint J

    Returns
    -------

    transforms : (F, J, 4, 4) ndarray
        Array of transforms for each
        frame F and joint J multiplied
        together
    """

    return ut.matrix_multiply(t0s, t1s)

def transforms_inv(ts):
    fts = ts.reshape(-1, 4, 4)
    fts = np.array(list(map(lambda x: np.linalg.inv(x), fts)))
    return fts.reshape(ts.shape)

def transforms_blank(anim):
    """
    Blank Transforms

    Parameters
    ----------

    anim : Animation
        Input animation

    Returns
    -------

    transforms : (F, J, 4, 4) ndarray
        Array of identity transforms for
        each frame F and joint J
    """

    ts = np.zeros(anim.shape + (4, 4))
    ts[:,:,0,0] = 1.0; ts[:,:,1,1] = 1.0;
    ts[:,:,2,2] = 1.0; ts[:,:,3,3] = 1.0;
    return ts

def transforms_global(anim):
    """
    Global Animation Transforms

    This relies on joint ordering
    being incremental. That means a joint
    J1 must not be a ancestor of J0 if
    J0 appears before J1 in the joint
    ordering.

    Parameters
    ----------

    anim : Animation
        Input animation

    Returns
    ------

    transforms : (F, J, 4, 4) ndarray
        Array of global transforms for
        each frame F and joint J
    """

    joints  = np.arange(anim.shape[1])
    parents = np.arange(anim.shape[1])
    locals  = transforms_local(anim)
    globals = transforms_blank(anim)

    globals[:,0] = locals[:,0]

    for i in range(1, anim.shape[1]):
        globals[:,i] = transforms_multiply(globals[:,anim.parents[i]], locals[:,i])

    return globals

def positions_global(anim):
    """
    Global Joint Positions

    Given an animation compute the global joint
    positions at at every frame

    Parameters
    ----------

    anim : Animation
        Input animation

    Returns
    -------

    positions : (F, J, 3) ndarray
        Positions for every frame F
        and joint position J
    """

    positions = transforms_global(anim)[:,:,:,3]
    return positions[:,:,:3] / positions[:,:,3,np.newaxis]


"""The main lagrangian-filtering module.

This module contains the crucial datastructure for
lagrangian-filtering, `LagrangeFilter`. See project documentation
for examples on how to construct a filtering workflow using this
library.

"""

import dask.array as da
import numpy as np
from datetime import timedelta, datetime
from glob import iglob
import parcels
from scipy import signal
import netCDF4
import xarray as xr
import netCDF4 as nc
import sys
from .file import LagrangeParticleFile

import multiprocessing as mp
from functools import partial

import pathos.pools as pp

class LagrangeFilter(object):
    """The main class for a Lagrangian filtering workflow.

    The workflow is set up using the input files and the filtering
    parameters. Filtering can be performed all at once, or on
    individual time levels.

    Data must contain horizontal velocity components `U` and `V` to
    perform the Lagrangian frame transformation. Any variables that should
    be filtered must be specified in the `sample_variables` list (this
    includes velocity).

    Note:
        We use the OceanParcels convention for variable names. This means that
        ``U``, ``V``, ``lon``, ``lat``, ``time`` and ``depth`` are canonical
        names for properties required for particle advection. The mapping from
        the actual variable name in your data files to these canonical names
        is contained in the `variables` and `dimensions` dictionaries. When
        specifying `filenames` or `sample_variables`, the canonical names
        must be used, however any other variables may use whatever name you
        would like.

    Once the `LagrangeFilter` has been constructed, you may call it as
    a function to perform the filtering workflow. See :func:`~filter`
    for documentation.

    Example:
        A straightforward filtering workflow::

            f = LagrangeFilter(
                name, filenames, variables, dimensions, sample_variables,
            )
            f()

        Would result in a new file with the given `name` and an appropriate
        extension containing the filtered data for each of the `sample_variables`.

    Args:
        name (str): The name of the workflow
        filenames (Dict[str, str]): A mapping from data variable names
            to the files containing the data.

            Filenames can contain globs if the data is spread across
            multiple files.
        variables_or_data (Union[Dict[str, str], xarray.Dataset]): Either
            a mapping from canonical variable names to the variable
            names in your data files, or an xarray Dataset containing
            the input data.
        dimensions (Dict[str, str]): A mapping from canonical dimension
            names to the dimension names in your data files.
        sample_variables ([str]): A list of variable names that should be sampled
            into the Lagrangian frame of reference and filtered.
        mesh (:obj:`str`, optional): The OceanParcels mesh type, either "flat"
            or "spherical". "flat" meshes are expected to have dimensions
            in metres, and "spherical" meshes in degrees.
        c_grid (:obj:`bool`, optional): Whether to interpolate velocity
            components on an Arakawa C grid (defaults to no).
        indices (:obj:`Dict[str, [int]]`, optional): An optional dictionary
            specifying the indices to which a certain dimension should
            be restricted.
        uneven_window (:obj:`bool`, optional): Whether to allow different
            lengths for the forward and backward advection phases.
        window_size (:obj:`float`, optional): The nominal length of the both
            the forward and backward advection windows, in seconds. A
            longer window may better capture the low-frequency signal to be
            removed.
        highpass_frequency (:obj:`float`, optional): The 3dB cutoff frequency
            for filtering, below which spectral components will be
            attenuated. This should be an angular frequency, in [rad/s].
        advection_dt (:obj:`datetime.timedelta`, optional): The timestep
            to use for advection. May need to be adjusted depending on the
            resolution/frequency of your data.

    """

    def __init__(
        self,
        name,
        filenames_or_dataset,
        variables,
        dimensions,
        sample_variables,
        mesh="flat",
        c_grid=False,
        indices={},
        uneven_window=False,
        window_size=None,
        highpass_frequency=5e-5,
        advection_dt=timedelta(minutes=5),
    ):
        # The name of this filter
        self.name = name
        # Width of window over which our filter computes a meaningful result
        # in seconds. Default to 3.5 days on either side
        if window_size is None:
            self.window_size = timedelta(days=3.5).total_seconds()
        else:
            self.window_size = window_size
        # Whether we're permitted to use uneven windows on either side
        self.uneven_window = uneven_window

        # copy input file dictionaries so we can construct the output file
        # filenames dictionary is modified to expand globs when
        # the fieldset is constructed
        self._filenames = filenames_or_dataset
        self._variables = variables
        self._dimensions = dimensions
        self._indices = indices
        # sample variables without the "var_" prefix
        self._sample_variables = sample_variables

        # choose the fieldset constructor depending on the format
        # of the input data
        if isinstance(filenames_or_dataset, xr.Dataset):
            fieldset_constructor = parcels.FieldSet.from_xarray_dataset
        else:
            fieldset_constructor = parcels.FieldSet.from_netcdf

        # for C-grid data, we have to change the interpolation method
        fieldset_kwargs = {}
        if c_grid:
            interp_method = {}
            for v in variables:
                if v in ["U", "V", "W"]:
                    interp_method[v] = "cgrid_velocity"
                else:
                    interp_method[v] = "cgrid_tracer"

            fieldset_kwargs["interp_method"] = interp_method

        # construct the OceanParcels FieldSet to use for particle advection
        self.fieldset = fieldset_constructor(
            filenames_or_dataset,
            variables,
            dimensions,
            indices=indices,
            mesh=mesh,
            **fieldset_kwargs,
        )
        # save the lon/lat on which to seed particles
        # this is saved here because if the grid is later made periodic, the
        # underlying grids will be modified, and we'll seed particles in the halos
        if self.fieldset.gridset.grids[0].gtype in [
            parcels.GridCode.CurvilinearZGrid,
            parcels.GridCode.CurvilinearSGrid,
        ]:
            self._curvilinear = True
            self._grid_lon = self.fieldset.gridset.grids[0].lon
            self._grid_lat = self.fieldset.gridset.grids[0].lat
        else:
            self._curvilinear = False
            self._grid_lon, self._grid_lat = np.meshgrid(
                self.fieldset.gridset.grids[0].lon, self.fieldset.gridset.grids[0].lat
            )

        # starts off non-periodic
        self._is_zonally_periodic = False
        self._is_meridionally_periodic = False

        # guess the output timestep
        times = self.fieldset.gridset.grids[0].time
        self.output_dt = times[1] - times[0]
        print('timestep =',self.output_dt,'seconds')

        # create the filter - use a 4th order Butterworth for the moment
        # make sure to convert angular frequency back to linear for passing to the
        # filter constructor
        fs = 1.0 / self.output_dt
        self.inertial_filter = signal.butter(
            4, highpass_frequency / (2 * np.pi), "lowpass", fs=fs
        )

        # timestep for advection
        self.advection_dt = advection_dt

        # the sample variable attribute has 'var_' prepended to map to
        # variables on particles
        self.sample_variables = ["var_" + v for v in sample_variables]
        # create the particle class and kernel for sampling
        # map sampled variables to fields
        self.particleclass = ParticleFactory(sample_variables)
        self._create_sample_kernel(sample_variables)
        self.kernel = parcels.AdvectionRK4 + self.sample_kernel

        # compile kernels
        self._compile(self.sample_kernel)
        self._compile(self.kernel)

    def _create_sample_kernel(self, sample_variables):
        """Create the parcels kernel for sampling fields during advection."""

        # make sure the fieldset has C code names assigned, etc.
        self.fieldset.check_complete()

        # string for the kernel itself
        f_str = "def sample_kernel(particle, fieldset, time):\n"
        for v in sample_variables:
            f_str += f"\tparticle.var_{v} = fieldset.{v}[time, particle.depth, particle.lat, particle.lon]\n"
        else:
            f_str += "\tpass"

        # create the kernel
        self.sample_kernel = parcels.Kernel(
            self.fieldset,
            self.particleclass.getPType(),
            funcname="sample_kernel",
            funcvars=["particle", "fieldset", "time"],
            funccode=f_str,
        )

    def _compile(self, kernel):
        """Compile a kernel and tell it to load the resulting shared library."""

        kernel.compile(compiler=parcels.compiler.GNUCompiler())
        kernel.load_lib()

    def make_zonally_periodic(self, width=None):
        """Mark the domain as zonally periodic.

        This will add a halo to the eastern and western edges of the
        domain, so that they may cross over during advection without
        being marked out of bounds. If a particle ends up within the
        halo after advection, it is reset to the valid portion of the
        domain.

        If the domain has already been marked as zonally periodic,
        nothing happens.

        Due to the method of resetting particles that end up in the
        halo, this is incompatible with curvilinear grids.

        Args:
            width (:obj:`int`, optional): The width of the halo,
                defaults to 5 (per parcels). This needs to be less
                than half the number of points in the grid in the x
                direction. This may need to be adjusted for small
                domains, or if particles are still escaping the halo.

        Note:
            This causes the kernel to be recompiled to add another stage
            which resets particles that end up in the halo to the main
            domain.

            If the kernel has already been recompiled for meridional periodicity,
            it is again reset to include periodicity in both
            directions.

        """

        # the method of resetting particles won't work on a curvilinear grid
        if self._curvilinear:
            raise Exception("curvilinear grids can not be periodic")

        # make sure we can't do this twice
        if self._is_zonally_periodic:
            return

        # add constants that are accessible within the kernel denoting the
        # edges of the halo region
        self.fieldset.add_constant("halo_west", self.fieldset.gridset.grids[0].lon[0])
        self.fieldset.add_constant("halo_east", self.fieldset.gridset.grids[0].lon[-1])

        if width is None:
            self.fieldset.add_periodic_halo(zonal=True)
        else:
            self.fieldset.add_periodic_halo(zonal=True, halosize=width)

        # unload the advection-only kernel, and add the periodic-reset kernel
        self.kernel.remove_lib()

        if self._is_meridionally_periodic:
            k = _doubly_periodic_BC
        else:
            k = _zonally_periodic_BC

        periodic_kernel = parcels.Kernel(
            self.fieldset, self.particleclass.getPType(), k
        )

        self.kernel = parcels.AdvectionRK4 + periodic_kernel + self.sample_kernel
        self._compile(self.kernel)

        self._is_zonally_periodic = True

    def make_meridionally_periodic(self, width=None):
        """Mark the domain as meridionally periodic.

        This will add a halo to the northern and southern edges of the
        domain, so that they may cross over during advection without
        being marked out of bounds. If a particle ends up within the
        halo after advection, it is reset to the valid portion of the
        domain.

        If the domain has already been marked as meridionally periodic,
        nothing happens.

        Due to the method of resetting particles that end up in the
        halo, this is incompatible with curvilinear grids.

        Args:
            width (:obj:`int`, optional): The width of the halo,
                defaults to 5 (per parcels). This needs to be less
                than half the number of points in the grid in the y
                direction. This may need to be adjusted for small
                domains, or if particles are still escaping the halo.

        Note:
            This causes the kernel to be recompiled to add another stage
            which resets particles that end up in the halo to the main
            domain.

            If the kernel has already been recompiled for zonal periodicity,
            it is again reset to include periodicity in both
            directions.

        """

        # the method of resetting particles won't work on a curvilinear grid
        if self._curvilinear:
            raise Exception("curvilinear grids can not be periodic")

        # make sure we can't do this twice
        if self._is_meridionally_periodic:
            return

        # add constants that are accessible within the kernel denoting the
        # edges of the halo region
        self.fieldset.add_constant("halo_north", self.fieldset.gridset.grids[0].lat[-1])
        self.fieldset.add_constant("halo_south", self.fieldset.gridset.grids[0].lat[0])

        if width is None:
            self.fieldset.add_periodic_halo(meridional=True)
        else:
            self.fieldset.add_periodic_halo(meridional=True, halosize=width)

        # unload the previous kernel, and add the meridionally-periodic kernel
        self.kernel.remove_lib()

        if self._is_zonally_periodic:
            k = _doubly_periodic_BC
        else:
            k = _meridionally_periodic_BC

        periodic_kernel = parcels.Kernel(
            self.fieldset, self.particleclass.getPType(), k
        )

        self.kernel = parcels.AdvectionRK4 + periodic_kernel + self.sample_kernel
        self._compile(self.kernel)

        self._is_meridionally_periodic = True

    def particleset(self, time):
        """Create a ParticleSet initialised at the given time.

        Args:
            time (float): The origin time for forward and backward advection
                on this ParticleSet.

        Returns:
            parcels.ParticleSet: A new ParticleSet containing a single particle
                at every gridpoint, initialised at the specified time.

        """

        # reset the global particle ID counter so we can rely on particle IDs making sense
        parcels.particle.lastID = 0

        return parcels.ParticleSet(
            self.fieldset,
            pclass=self.particleclass,
            lon=self._grid_lon,
            lat=self._grid_lat,
            time=time,
        )

    def advection_step(self, time, output_time=False):
        """Perform forward-backward advection at a single point in time.

        This routine is responsible for creating a new ParticleSet at
        the given time, and performing the forward and backward
        advection steps in the Lagrangian transformation.

        Args:
            time (float): The point in time at which to calculate filtered data.
            output_time (:obj:`bool`, optional): Whether to include "time" as
                a numpy array in the output dictionary, for doing manual analysis.

        Note:
            If ``output_time`` is True, the output object will not be compatible
            with the default filtering workflow, :func:`~filter_step`!

        Returns:
            Dict[str, (int, dask.array)]: A dictionary of the advection
                data, mapping variable names to a pair. The first element is
                the index of the sampled timestep in the data, and the
                second element is a lazy dask array concatenating the forward
                and backward advection data.

        """

        # seed all particles at gridpoints
        ps = self.particleset(time)
        # execute the sample-only kernel to efficiently grab the initial condition
        ps.kernel = self.sample_kernel
        ps.execute(self.sample_kernel, runtime=0, dt=self.advection_dt)

        # set up the temporary output file for the initial condition and
        # forward advection
        outfile = LagrangeParticleFile(ps, self.output_dt, self.sample_variables)

        # now the forward advection kernel can run
        outfile.set_group("forward")
        ps.kernel = self.kernel
        ps.execute(
            self.kernel,
            runtime=self.window_size,
            dt=self.advection_dt,
            output_file=outfile,
        )

        # reseed particles back on the grid, then advect backwards
        # we don't need any initial condition sampling since we've already done it
        outfile.set_group("backward")
        ps = self.particleset(time)
        ps.kernel = self.kernel
        ps.execute(
            self.kernel,
            runtime=self.window_size,
            dt=-self.advection_dt,
            output_file=outfile,
        )

        # stitch together and filter all sample variables from the temporary
        # output data
        da_out = {}
        for v in self.sample_variables:
            # load data lazily as dask arrays, for forward and backward segments
            var_array_forward = da.from_array(
                outfile.data("forward")[v], chunks=(None, "auto")
            )[:-1, :]
            var_array_backward = da.from_array(
                outfile.data("backward")[v], chunks=(None, "auto")
            )[:-1, :]

            # get an index into the middle of the array
            time_index_data = var_array_backward.shape[0] - 1

            # construct proper sequence by concatenating data and flipping the backward segment
            # for var_array_forward, skip the initial output for both the sample-only and
            # sample-advection kernels, which have meaningless data
            var_array = da.concatenate(
                (da.flip(var_array_backward[1:, :], axis=0), var_array_forward)
            )

            da_out[v] = (time_index_data, var_array)

        if output_time:
            da_out["time"] = np.concatenate(
                (
                    outfile.data("backward").attrs["time"][1:-1][::-1],
                    outfile.data("forward").attrs["time"][:-1],
                )
            )

        return da_out

    def filter_step(self, advection_data):
        """Perform filtering of a single step of advection data.

        The Lagrangian-transformed data from :func:`~advection_step` is
        high-pass filtered in time, leaving only the signal at the
        origin point (i.e. the filtered forward and backward advection
        data is discarded).

        Args:
            advection_data (Dict[str, (int, dask.array)]): A dictionary of
                particle advection data from a single timestep, returned
                from :func:`~advection_step`.

        Returns:
            Dict[str, dask.array]: A dictionary mapping sampled
                variable names to a 1D dask array containing the
                filtered data at the specified time. This data is not
                lazy, as it has already been computed out of the
                temporary advection data.

        """

        da_out = {}
        for v, a in advection_data.items():
            time_index_data, var_array = a

            def filter_select(x):
                return signal.filtfilt(*self.inertial_filter, x)[..., time_index_data]

            # apply scipy filter as a ufunc
            # mapping an array to scalar over the first axis, automatically vectorize execution
            # and allow rechunking (since we have a chunk boundary across the first axis)
            filtered = da.apply_gufunc(
                filter_select,
                "(i)->()",
                var_array,
                axis=0,
                output_dtypes=var_array.dtype,
                allow_rechunk=True,
            )

            da_out[v] = filtered.compute()

        return da_out

    def filter(self, *args, **kwargs):
        """Run the filtering process on this experiment.

        Note:
            Instead of `f.filter(...)`, you can call `f(...)` directly.

        This is main method of the filtering workflow. The timesteps
        to filter may either be specified manually, or determined from
        the window size and the timesteps within the input files. In
        this latter case, only timesteps that have the full window
        size on either side are selected.

        Note:
            If `absolute` is True, the times must be the same datatype
            as those the input data. For dates with a calendar, this
            is likely :obj:`np.datetime64` or :obj:`cftime.datetime`.
            For abstract times, this may simply be a number.

        Args:
            times (:obj:`[float]`, optional): A list of timesteps at
                which to run the filtering. If this is omitted, all
                timesteps that are fully covered by the filtering
                window are selected.
            clobber (:obj:`bool`, optional): Whether to overwrite any
                existing output file with the same name as this
                experiment. Default behaviour will not clobber an
                existing output file.

            absolute (:obj:`bool`, optional): If `times` is provided,
                this argument determines whether to interpret them
                as relative to the first timestep in the input dataset
                (False, default), or as absolute, following the actual
                time dimension in the dataset (True).

        """

        self(*args, **kwargs)

    def create_out(self, clobber=False, date=None):
        """Create a netCDF dataset to hold filtered output.

        Here we create a new ``netCDF4.Dataset`` for filtered
        output. For each sampled variable in the input files, a
        corresponding variable in created in the output file, with
        the same dimensions.

        Returns:
            netCDF4.Dataset: A single dataset that will hold all
                filtered output.

        """

        # the output dataset we're creating
        filename = self.name
        if date is not None:
            date_str = str(date)[:13]
            filename += '_' + date_str
        filename += ".nc"
        print(filename)
        ds = netCDF4.Dataset(filename, "w", clobber=clobber)

        # helper function to create the dimensions in the ouput file
        def create_dimension(dims, dim, var):
            # translate from parcels -> file convention
            # and check whether we've already created this dimension
            # (e.g. for a previous variable)
            file_dim = dims[dim]
            if file_dim in ds.variables:
                return ds.variables[file_dim].dimensions[
                    1
                    if len(ds.variables[file_dim].dimensions) > 1 and dim == "lon"
                    else 0
                ]

            # get the file containing the dimension data
            v_orig = self._variables.get(var, var)
            if isinstance(self._filenames, xr.Dataset):
                ds_orig = self._filenames[file_dim]
            else:
                if isinstance(self._filenames[var], dict):
                    filename = self._filenames[var][dim]
                else:
                    filename = self._filenames[var]

                if isinstance(filename, list):
                    filename = filename[0]

                ds_orig = xr.open_dataset(next(iglob(filename)))[file_dim]

            # create dimensions if needed
            for d in ds_orig.dims:
                if d not in ds.dimensions:
                    ds.createDimension(d, ds_orig[d].size)

            # create the dimension variable
            ds.createVariable(file_dim, ds_orig.dtype, dimensions=ds_orig.dims)
            ds.variables[file_dim][:] = ds_orig

            # curvilinear grid case
            return ds_orig.dims[1 if len(ds_orig.dims) > 1 and dim == "lon" else 0]

        # create a time dimension if dimensions are uniform across all variables
        if "time" in self._dimensions:
            dim_time = self._dimensions["time"]
            ds.createDimension(dim_time)
            # Add by FlG
            if date is not None:
                ds.createVariable(
                "time",
                "float32",
                dimensions=(dim_time,),
                )
                calendar = 'standard'
                units = 'seconds since 1900-01-01 00:00'
                ts = (date - np.datetime64('1970-01-01T00:00:00Z')) / np.timedelta64(1, 's')
                dt = datetime.utcfromtimestamp(ts)
                ds["time"][:] = nc.date2num(dt, units=units, calendar=calendar)

        else:
            dim_time = None

        for v in self._sample_variables:
            # translate if required (parcels -> file convention)
            v_orig = self._variables.get(v, v)

            # open all the relevant files for this variable
            if isinstance(self._filenames, xr.Dataset):
                ds_orig = self._filenames[v_orig]
            else:
                if isinstance(self._filenames[v], dict):
                    # variable -> dictionary (for separate coordinate files)
                    filename = self._filenames[v]["data"]
                else:
                    # otherwise, we just have a plain variable -> file mapping
                    filename = self._filenames[v]

                # globs can give us a list, but we only need the first item
                # to get the metadata
                if isinstance(filename, list):
                    filename = filename[0]

                ds_orig = xr.open_dataset(next(iglob(filename)))[v_orig]

            # are dimensions defined specifically for this variable?
            if v in self._dimensions:
                dims = self._dimensions[v]
            else:
                dims = self._dimensions

            # are indices defined specifically for this variable?
            if v in self._indices:
                indices = self._indices[v]
            else:
                indices = self._indices

            # translate to variable names from the source files
            indices = {dims[v]: ind for v, ind in indices.items()}

            # select only the relevant indices
            # in particular, squeeze to drop z dimension if we index it out
            local_indices = {k: v for k, v in indices.items() if k in ds_orig.dims}
            ds_orig = ds_orig.isel(**local_indices).squeeze()

            # create time dimension if required (i.e. not already in the
            # output file we've created)
            out_dims = {}
            if dim_time is None:
                out_dims["time"] = dims["time"]
                if dims["time"] not in ds.dimensions:
                    ds.createDimension(dims["time"])
            else:
                out_dims["time"] = dim_time

            # for each non-time dimension, create it if it doesn't already exist
            # in the output file
            for d in ["lat", "lon"]:
                out_dims[d] = create_dimension(dims, d, v)

            # create the variable in the dataset itself
            ds.createVariable(
                "var_" + v,
                "float32",
                dimensions=(out_dims["time"], out_dims["lat"], out_dims["lon"]),
            )

        return ds

    def _window_times(self, times, absolute):

        """Restrict an array of times to those which have an adequate window,
        optionally converting from absolute to relative first.

        """

        tgrid = self.fieldset.gridset.grids[0].time

        if times is None:
            times = tgrid.copy()

        if absolute:
            times = self.fieldset.gridset.grids[0].time_origin.reltime(times)

        times = np.array(times)
        window_left = times - tgrid[0] >= self.window_size
        window_right = times <= tgrid[-1] - self.window_size
        return times[window_left & window_right]

    def _do_filtering(self, time):

        date = self.fieldset.gridset.grids[0].time_origin.fulltime(time)

        print(date)

        ds = self.create_out(clobber=True, date=date)

        # returns a dictionary of sample_variable -> dask array
        filtered = self.filter_step(self.advection_step(time))
        for v, a in filtered.items():
            ds[v][0, ...] = a

        ds.close()

    def __call__(self, times=None, absolute=False, clobber=False):
        """Run the filtering process on this experiment."""

        if self.uneven_window:
            raise NotImplementedError("uneven windows aren't supported")

        # either restrict the specified times to period covered by window,
        # or use the full range of times covered by window
        times = self._window_times(times, absolute)

        for time in times:
            self._do_filtering(time)

        print('End of the program')

def ParticleFactory(variables, name="SamplingParticle", BaseClass=parcels.JITParticle):
    """Create a Particle class that samples the specified variables.

    The variables that should be sampled will be prepended by ``var_`` as
    class attributes, in case there are any namespace clashes with existing
    variables on the base class.

    Args:
        variables ([str]): A list of variable names which should be sampled.
        name (str): The name of the generated particle class.
        BaseClass (Type[parcels.particle._Particle]): The base particles class upon
            which to append the required variables.

    Returns:
        Type[parcels.particle._Particle]: The new particle class

    """

    var_dict = {"var_" + v: parcels.Variable("var_" + v) for v in variables}

    newclass = type(name, (BaseClass,), var_dict)
    return newclass

def _zonally_periodic_BC(particle, fieldset, time):
    if particle.lon < fieldset.halo_west:
        particle.lon += fieldset.halo_east - fieldset.halo_west
    elif particle.lon > fieldset.halo_east:
        particle.lon -= fieldset.halo_east - fieldset.halo_west

def _meridionally_periodic_BC(particle, fieldset, time):
    if particle.lat < fieldset.halo_south:
        particle.lat += fieldset.halo_north - fieldset.halo_south
    elif particle.lat > fieldset.halo_north:
        particle.lat -= fieldset.halo_north - fieldset.halo_south

def _doubly_periodic_BC(particle, fieldset, time):
    # because the kernel is run through code generation, we can't simply
    # call the above kernels, so we unfortunately have to reproduce them
    # in full here
    if particle.lon < fieldset.halo_west:
        particle.lon += fieldset.halo_east - fieldset.halo_west
    elif particle.lon > fieldset.halo_east:
        particle.lon -= fieldset.halo_east - fieldset.halo_west

    if particle.lat < fieldset.halo_south:
        particle.lat += fieldset.halo_north - fieldset.halo_south
    elif particle.lat > fieldset.halo_north:
        particle.lat -= fieldset.halo_north - fieldset.halo_south


"""
Reads collapsing matrix and summary file from ATAV and generates an
empirical permutation-based Q-Q plot.

Written by Charlie Wolock <cw3026@cumc.columbia.edu>
"""

from functools import partial
from operator import le, lt
from scipy import stats
import argparse
import ctypes
import matplotlib
matplotlib.use('pdf')
import matplotlib.pyplot as plt
import multiprocessing as mp
import numpy as np
import random
import seaborn as sns
import sys

sns.set_style('darkgrid')

def valid_numerical_argument(
        arg, arg_name, arg_type=int, min_value=0, max_value=sys.maxint,
        left_op=lt, right_op=le):
    """
    Confirm that the specified value is valid in the range
    (minimum_value, maximum_value] (by default)
    :param arg: the value to be tested
    :param arg_name: the name of the parameter
    :param arg_type: the type of the parameter, e.g. int or float
    :param min_value: the minimum value for the parameter, exclusive
    :param max_value: the maximum value for the parameter, inclusive
    :param left_op: the operator for testing left_op(min_value, value)
    :param right_op: the operator testing right_op(value, max_value)
    :return: arg_type(arg) if arg is valid
    """
    try:
        value = arg_type(arg)
        if left_op(min_value, value) and right_op(value, max_value):
            return value
        else:
            raise argparse.ArgumentTypeError(
                "{arg_name} ({arg}) is not in the range "
                "{left_endpoint}{min_value}, {max_value}{right_endpoint}".format(
                    arg_name=arg_name, arg=arg, min_value=min_value,
                    max_value=max_value, left_endpoint="(" if left_op == lt else "[",
                    right_endpoint="]" if right_op == le else ")"))
    except TypeError:
        raise argparse.ArgumentTypeError(
            "{arg_name} ({arg}) is not a valid {arg_type}".format(
                arg_name=arg_name, arg=arg, arg_type=arg_type.__name__))

def precalc(pair, num_case, num_ctrl):
    """
    Calculate FET p-value for pair of qual case + qual ctrl
    :param pair: (tuple) number of qualified cases and controls
    :param num_case: (int) number of case samples
    :param num_ctrl: (int) number of control samples
    """
    odds, pval = stats.fisher_exact([[pair[0], num_case - pair[0]],
                                     [pair[1], num_ctrl - pair[1]]])
    return (pair[0], pair[1], pval)

def shared_mem(result):
    """
    Populate lookup table with pre-calculated p-values
    :param result: (tuple) num. qual. cases, num. qual ctrls, associated FET p-value
    """
    lookup[result[0], result[1]] = result[2]

def permute(counter, statuses):
    """
    Permute affection statuses and look up p-vals for this configuration
    :param counter: (int) iteration counter for multiprocessing to keep track
    :param statuses: (list) randomized list of case/control statuses
    """
    # randomize the statuses
    perm_statuses = np.array(random.sample(statuses, len(statuses)))
    case_indices = np.where(perm_statuses == '2')[0]
    ctrl_indices = np.where(perm_statuses == '1')[0]
    # make contingency table and calculate pvals for each gene
    pvals = []
    for row in range(col_matrix.shape[0]):
        q_case = np.where(col_matrix[row, case_indices] == 1)[0].shape[0]
        q_ctrl = np.where(col_matrix[row, ctrl_indices] == 1)[0].shape[0]
        # fisher exact test
        # if number of qualified samples is in the lookup table, use that
        if q_case + 1 <= lookup.shape[0] and q_ctrl + 1 <= lookup.shape[1]:
            pvalue = lookup[q_case, q_ctrl]
        # otherwise just do the FET
        else:
            odds, pvalue = stats.fisher_exact([[q_case, len(case_indices) - q_case],
                                               [q_ctrl, len(ctrl_indices) - q_ctrl]])
        pvals.append(pvalue)
    pvals = sorted(pvals)
    return pvals

def read_summary_file(summary_file):
    """
    Read in partial summary and add gene info to dictionary
    :param summary_file: (string) name of partial summary file
    """
    genes = {}
    case_qual = []
    ctrl_qual = []
    with open(summary_file, 'r') as infile:
        header = infile.readline().strip().split(',')
        gene_index = header.index('Gene Name')
        qcase_index = header.index('Qualified Case')
        uqcase_index = header.index('Unqualified Case')
        qctrl_index = header.index('Qualified Ctrl')
        uqctrl_index = header.index('Unqualified Ctrl')
        # read first line to determine number of cases and ctrls
        first = infile.readline().strip().split(',')
        genes[first[gene_index]] = [x for x in first[2:]]
        ncase = int(first[qcase_index]) + int(first[uqcase_index])
        nctrl = int(first[qctrl_index]) + int(first[uqctrl_index])
        case_qual.append(int(first[qcase_index]))
        ctrl_qual.append(int(first[qctrl_index]))
        for line in infile:
            line = line.strip().split(',')
            genes[line[gene_index]] = [x for x in line[2:]]
            case_qual.append(int(line[qcase_index]))
            ctrl_qual.append(int(line[qctrl_index]))
    return (genes, ncase, nctrl, case_qual, ctrl_qual)

def read_matrix_file(matrix_file, ngenes):
    """
    Read in matrix file
    :param matrix_file: (string) name of matrix file
    :param ngenes: (int) number of genes (or general collapsing units)
    """
    with open(matrix_file, 'r') as infile:
        samps = infile.readline().strip().split('\t')
        nsamps = len(samps) - 1
        col_matrix = np.zeros((ngenes, nsamps))
        for i, line in enumerate(infile):
            line = line.strip().split('\t')
            col_matrix[i, :] = line[1:]
    return col_matrix

def plot_qq(qq_file, exp, obs, lower, upper, sorted_genes):
    """
    Create QQ-plot of observed and expected FET p-values, with 2.5%ile and 97.5%ile
        bounds.
    Lambda = slope of regression on chi-sq transformed obs and exp p-values, after
         removal of p-values of 1 and genome-wide sig. p-values
    :param qq_file: (string) name of QQ-plot file to be created
    :param exp: (np array) expected p-values
    :param obs: (np array) observed p-values
    :param lower: (np array) 2.5%ile exp p-values
    :param upper: (np array) 97.5%ile exp p-values
    :param sorted_genes: (list) genes sorted by descending rank
    """
    # change p-values > 1 to 1
    exp[exp > 1] = 1
    obs[obs > 1] = 1

    # calculate lambda by regression method
    # first remove p-values of 1 and < genome-wide significance
    gws = 0.05 / len(obs)
    reg_exp = exp[(obs > gws) & (exp > gws) & (exp < 1) & (obs < 1)]
    reg_obs = obs[(obs > gws) & (exp > gws) & (exp < 1) & (obs < 1)]
    reg_exp = stats.chi2.ppf(1 - reg_exp, 1)
    reg_obs = stats.chi2.ppf(1 - reg_obs, 1)
    # for lstsq, explanatory vars must be in column form
    reg_exp = reg_exp[:, np.newaxis]
    # this least squares regression forces intercept to be 0
    slope = np.linalg.lstsq(reg_exp, reg_obs)[0][0]
    lambda_factor = slope
    lambda_factor = np.around(lambda_factor, 5)

    # transform p-values to log10
    exp = -np.log10(exp)
    obs = -np.log10(obs)
    # note that "upper" and "lower" switch here because of the transformation
    lower = -np.log10(lower)
    upper = -np.log10(upper)

    # set limits of qq plot axes
    axisMax_obs = np.ceil(max(obs))
    axisMax_exp = np.ceil(max(exp))

    # initialize qqplot with axes and labels
    fig = plt.figure(figsize=(12, 12))
    plt.xlim([0, axisMax_exp])
    plt.xlabel('Expected -log10(p)', fontsize=20)
    plt.ylim([0, axisMax_obs])
    plt.ylabel('Observed -log10(p)', fontsize=20)
    plt.title('QQ Plot: Observed vs. expected p-values. Lambda = {l}'.format(
              l=lambda_factor), fontsize=20)

    # change size of axis tick labels
    plt.tick_params(axis='both', which='major', labelsize=12)

    # plot the points, exp on x axis, obs on y axis
    dataAx = fig.add_subplot(1, 1, 1)
    dataAx.plot(exp, obs, 'r.', label='_nolegend_', markersize=12)

    # plot a diagonal line for comparison
    lineAx = fig.add_subplot(1, 1, 1)
    lineAx.plot([0, max(axisMax_obs, axisMax_exp)], [0, max(axisMax_obs, axisMax_exp)],
                'b-', label='_nolegend_')
    uppAx = fig.add_subplot(1, 1, 1)
    uppAx.plot(exp, upper, 'g-')
    lowAx = fig.add_subplot(1, 1, 1)
    lowAx.plot(exp, lower, 'y-')
    plt.legend(['2.5th percentile of expected p-values',
                '97.5th percentile of expected p-values'],
               loc=2)
    plt.tight_layout()
    plt.savefig(qq_file)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('partial_summary_file',
                        help='Specify name of summary file generated by summarize.py')
    parser.add_argument('matrix_file',
                        help='Specify name of matrix file generated by summarize.py')
    parser.add_argument('output', help='Specify prefix for output files')
    parser.add_argument('--nperms',
                        type=partial(valid_numerical_argument, arg_name='nperms',
                                     min_value=0, max_value=1000, arg_type=int),
                        default=1000,
                        help='Specify number of permutations')
    parser.add_argument('--nprocs',
                        type=partial(valid_numerical_argument, arg_name='nprocs',
                                     min_value=0, max_value=10, arg_type=int),
                        default=4,
                        help='Specify number of concurrent processes')
    args = parser.parse_args()

    # read summary file to get dictionary of gene-level info
    # num of cases, controls, and maximum number of qualified samples in any
    # gene
    genes, ncase, nctrl, case_qual, ctrl_qual = read_summary_file(
        args.partial_summary_file)

    # determine 99th percentile of qualified
    case_qual = np.array(case_qual)
    ctrl_qual = np.array(ctrl_qual)
    qual_sum = np.sort(case_qual + ctrl_qual)
    high_qual = np.percentile(qual_sum, 95, interpolation='higher')

    # calculate number of genes
    ngenes = len(genes)

    # read matrix file to build collapsing matrix
    col_matrix = read_matrix_file(args.matrix_file, ngenes)

    # calculate number of samples
    nsamps = col_matrix.shape[1]

    # create shared memory version of collapsing matrix
    shared_array_base1 = mp.Array(ctypes.c_int, ngenes * nsamps)
    shared_col_matrix = np.ctypeslib.as_array(shared_array_base1.get_obj())
    shared_col_matrix = shared_col_matrix.reshape(ngenes, nsamps)

    # calculate necessary dimensions of FET lookup table
    max_case = min([high_qual, ncase])
    max_ctrl = min([high_qual, nctrl])

    # create shared memory FET lookup table
    shared_array_base2 = mp.Array(
        ctypes.c_double, (max_case + 1) * (max_ctrl + 1))
    lookup = np.ctypeslib.as_array(shared_array_base2.get_obj())
    lookup = lookup.reshape(max_case + 1, max_ctrl + 1)

    # make list of qual case/qual ctrl pairs that need to be calculated for
    # lookup
    pairs = []

    # make list of (qual case, qual ctrl) tuples for FET lookup table
    for i in xrange(max_case + 1):
        for j in xrange(max_ctrl + 1):
            pairs.append((i, j))

    # generate FET p-values for lookup table
    pool = mp.Pool(processes=args.nprocs)
    fisher_results = pool.map(
        partial(precalc, num_case=ncase, num_ctrl=nctrl), pairs)
    pool.close()
    pool.join()

    # populate shared memory FET lookup table
    pool = mp.Pool(processes=args.nprocs)
    pool.map(shared_mem, fisher_results)
    pool.close()
    pool.join()

    # perform permutations to generate expected p-values
    status_l = ['2'] * ncase
    status_l.extend(['1'] * nctrl)
    pool = mp.Pool(processes=args.nprocs)
    permute_results = pool.map(
        partial(permute, statuses=status_l), range(args.nperms))
    pool.close()
    pool.join()

    # populate array of permuted p-values
    perm_pvals = np.ones((ngenes, args.nperms))
    for i, pvals in enumerate(permute_results):
        perm_pvals[:, i] = pvals

    # calculate 2.5%ile and 97.5%ile of permuted pvals
    bottom_perc = np.percentile(perm_pvals, 2.5, axis=1)
    top_perc = np.percentile(perm_pvals, 97.5, axis=1)

    # calculate observed p-values
    for gene, v in genes.iteritems():
        qcase = int(v[3])
        qctrl = int(v[6])
        if qcase + 1 <= lookup.shape[0] and qctrl + 1 <= lookup.shape[1]:
            v.append(lookup[qcase, qctrl])
        else:
            v.append(stats.fisher_exact(
                [[qcase, ncase - qcase], [qctrl, nctrl - qctrl]])[1])
    ordered = sorted(genes.items(), key=lambda x: x[1][-1])

    # create sorted gene list
    sorted_genes = [gene for gene, info in ordered]

    # make arrays of ordered observed and expected p-values
    obs_pvals = np.array([x[1][-1] for x in ordered])
    exp_pvals = np.mean(perm_pvals, axis=1)

    # generate qq plot
    plot_qq(args.output, exp_pvals, obs_pvals, bottom_perc, top_perc, sorted_genes)


import numpy as np

from numba import boolean, float64, int64, njit, optional, uint8
from numba.core.types.containers import Tuple

__author__ = "Peter Maxwell"
__copyright__ = "Copyright 2007-2021, The Cogent Project"
__credits__ = ["Peter Maxwell", "Gavin Huttley", "Stephen Ma"]
__license__ = "BSD-3"
__version__ = "2021.5.7a"
__maintainer__ = "Gavin Huttley"
__email__ = "Gavin.Huttley@anu.edu.au"
__status__ = "Production"

@njit(
    Tuple(types=(Tuple(types=(int64, int64)), int64, float64,))(
        int64[::1],
        int64[::1],
        int64[::1],
        int64,
        int64,
        int64,
        int64,
        optional(int64[::1]),
        optional(int64[::1]),
        optional(int64[::1]),
        optional(int64[::1]),
        int64[:, ::1],
        float64[:, ::1],
        float64[:, ::1],
        float64[:, ::1],
        float64[:, :, ::1],
        optional(float64[:, :, ::1]),
        float64,
        optional(int64[:, :, ::1]),
        optional(uint8[:, :, ::1]),
        optional(int64[::1]),
        boolean,
        boolean,
        boolean,
        boolean,
    ),
    cache=True,
)
def calc_rows(
    plan,
    x_index,
    y_index,
    i_low,
    i_high,
    j_low,
    j_high,
    i_sources,
    i_sources_offsets,
    j_sources,
    j_sources_offsets,
    state_directions,
    T,
    xgap_scores,
    ygap_scores,
    match_scores,
    mantissas,
    mantissa,
    exponents,
    track,
    track_enc,
    viterbi,
    local=False,
    use_scaling=False,
    use_logs=False,
):
    assert not (use_logs and not viterbi)
    assert not (use_logs and use_scaling)
    assert not (local and not viterbi)

    MIN_SCALE = -10000
    MAX_SCALE = +10000
    SCALE_STEP = 2.0 ** 50
    MIN_FLOAT_VALUE = 1.0 / SCALE_STEP
    source_row_index_cache = np.zeros(256)

    N = max(T.shape[0], T.shape[1])

    dest_states = max(0, state_directions.shape[0])

    row_count = x_index.shape[0]
    row_length = y_index.shape[0]

    max_x = match_scores.shape[1]
    max_y = match_scores.shape[2]

    max_x = max(xgap_scores.shape[1], max_x)
    max_y = max(ygap_scores.shape[1], max_y)

    for i in range(row_count):
        assert 0 <= x_index[i] <= max_x

    for j in range(row_length):
        assert 0 <= y_index[j] <= max_y

    assert j_low >= 0 and j_high > j_low and j_high <= row_length

    row_length = max(mantissas.shape[1], row_length)
    N = max(mantissas.shape[2], N)

    if use_scaling:
        row_length = max(exponents.shape[1], row_length)
        N = max(exponents.shape[2], N)

    if use_logs:
        impossible = -np.inf
    else:
        impossible = 0.0

    if viterbi and track is not None and track_enc is not None:
        N = max(track.shape[2], N)
        (tcode_x, tcode_y, tcode_s) = track_enc
    else:
        track = None
        tcode_x = tcode_y = tcode_s = 0

    overall_max_exponent = MIN_SCALE
    overall_max_mantissa = impossible
    last_i = last_j = last_state = -1

    max_exponent = MIN_SCALE

    for i in range(i_low, i_high):
        x = x_index[i]

        i_sources_start = i_sources_offsets[i]
        i_sources_end = i_sources_offsets[i + 1]

        current_row_index = plan[i]
        source_row_index_cache[0] = current_row_index

        a_count = i_sources_end - i_sources_start
        for a in range(a_count):
            prev_i = i_sources[a + i_sources_start]
            source_row_index_cache[a + 1] = plan[prev_i]

        if i == 0:
            if use_logs:
                mantissas[current_row_index, 0, 0] = 0.0
            else:
                mantissas[current_row_index, 0, 0] = 1.0
            if use_scaling:
                exponents[current_row_index, 0, 0] = 0
        else:
            mantissas[current_row_index, 0, 0] = impossible
            if use_scaling:
                exponents[current_row_index, 0, 0] = MIN_SCALE

        j_sources_end = j_sources_offsets[j_low]
        for j in range(j_low, j_high):
            j_sources_start = j_sources_end
            j_sources_end = j_sources_offsets[j + 1]

            for dest_state in range(dest_states):
                state = state_directions[dest_state, 0]
                bin = state_directions[dest_state, 1]
                dx = state_directions[dest_state, 2]
                dy = state_directions[dest_state, 3]

                max_mantissa = impossible
                max_exponent = MIN_SCALE
                partial_sum = 0.0
                pointer_state = N

                if dx:
                    a_low = 1
                    a_high = a_count + 1
                else:
                    a_low = 0
                    a_high = 1

                if dy:
                    b_low = 1
                    b_high = j_sources_end - j_sources_start + 1
                else:
                    b_low = 0
                    b_high = 1

                pointer_a = 0
                pointer_b = 0

                if use_scaling:
                    sub_partial_sum = 0.0

                    for a in range(a_low, a_high):
                        source_row_index = int(source_row_index_cache[a])
                        for b in range(b_low, b_high):
                            if dy:
                                prev_j = j_sources[b - 1 + j_sources_start]
                            else:
                                prev_j = j
                            min_prev_state = prev_j > 0

                            for prev_state in range(min_prev_state, N):
                                exponent = exponents[
                                    source_row_index, prev_j, prev_state
                                ]
                                if exponent == MIN_SCALE:
                                    continue

                                transition = T[prev_state, state]

                                mantissa = mantissas[
                                    source_row_index, prev_j, prev_state
                                ]
                                mantissa *= transition

                                if mantissa < MIN_FLOAT_VALUE:
                                    if mantissa == 0.0:
                                        continue
                                    assert mantissa >= 0.0 and transition >= 0.0

                                    while mantissa < MIN_FLOAT_VALUE:
                                        mantissa *= SCALE_STEP
                                        exponent += -1
                                        assert exponent > MIN_SCALE

                                elif mantissa > 1.0:
                                    mantissa *= MIN_FLOAT_VALUE
                                    exponent += 1
                                    assert exponent <= MAX_SCALE

                                if exponent > max_exponent:
                                    if exponent == max_exponent + 1:
                                        sub_partial_sum = partial_sum
                                    else:
                                        sub_partial_sum = 0.0
                                    partial_sum = 0.0
                                    max_mantissa = 0.0
                                    max_exponent = exponent

                                if exponent == max_exponent:
                                    partial_sum += mantissa
                                    if viterbi and mantissa > max_mantissa:
                                        max_mantissa = mantissa
                                        pointer_state = prev_state
                                        pointer_a = a
                                        pointer_b = b

                                elif exponent == max_exponent - 1:
                                    sub_partial_sum += mantissa

                            partial_sum += sub_partial_sum * MIN_FLOAT_VALUE
                else:
                    for a in range(a_low, a_high):
                        source_row_index = int(source_row_index_cache[a])
                        for b in range(b_low, b_high):
                            if dy:
                                prev_j = j_sources[b - 1 + j_sources_start]
                            else:
                                prev_j = j
                            min_prev_state = prev_j > 0

                            for prev_state in range(min_prev_state, N):
                                mantissa = mantissas[
                                    source_row_index, prev_j, prev_state
                                ]
                                transition = T[prev_state, state]
                                if use_logs:
                                    mantissa += transition
                                else:
                                    mantissa *= transition
                                    partial_sum += mantissa
                                if viterbi and mantissa > max_mantissa:
                                    max_mantissa = mantissa
                                    pointer_state = prev_state
                                    pointer_a = a
                                    pointer_b = b

                if viterbi:
                    mantissa = max_mantissa
                    if track is not None:
                        track[i, j, state] = (
                            (pointer_a << tcode_x)
                            | (pointer_b << tcode_y)
                            | (pointer_state << tcode_s)
                        )
                else:
                    mantissa = partial_sum

                if dy:
                    y = y_index[j]
                    if dx:
                        d_score = match_scores[bin, x, y]
                    else:
                        d_score = ygap_scores[bin, y]
                elif dx:
                    d_score = xgap_scores[bin, x]
                elif use_logs:
                    d_score = 0.0
                else:
                    d_score = 1.0

                if use_logs:
                    mantissa += d_score
                else:
                    mantissa *= d_score

                mantissas[current_row_index, j, state] = mantissa

                if use_scaling:
                    exponents[current_row_index, j, state] = max_exponent

                if local and dx and dy:
                    if (use_scaling and max_exponent > overall_max_exponent) or (
                        (not use_scaling or max_exponent == overall_max_exponent)
                        and (mantissa > overall_max_mantissa)
                    ):
                        overall_max_exponent = max_exponent
                        overall_max_mantissa = mantissa
                        last_i = i
                        last_j = j
                        last_state = state

    if not local:
        last_i = i_high - 1
        last_j = j_high - 1
        last_state = state
    else:
        mantissa = overall_max_mantissa
        max_exponent = overall_max_exponent

    if use_scaling:
        score = np.log(mantissa) + np.log(SCALE_STEP) * max_exponent
    elif use_logs:
        score = mantissa
    else:
        score = np.log(mantissa)
    return ((last_i, last_j), last_state, score)


import os
import warnings
import numpy as np
from pyrates.utility.genetic_algorithm import CGSGeneticAlgorithm
from pandas import DataFrame, read_hdf
from copy import deepcopy

class CustomGOA(CGSGeneticAlgorithm):

    def eval_fitness(self, target: list, **kwargs):

        # define simulation conditions
        worker_file = self.cgs_config['worker_file'] if 'worker_file' in self.cgs_config else None
        param_grid = self.pop.drop(['fitness', 'sigma', 'results'], axis=1)
        result_vars = ['r_e', 'r_i', 'r_a']
        freq_targets = [0.0, 0.0, 0.0, 0.0, np.nan]
        #param_grid, invalid_params = eval_params(param_grid)
        conditions = [{},  # healthy control
                      {'k_pe': 0.2, 'k_ae': 0.2},  # AMPA blockade in GPe
                      {'k_pe': 0.2, 'k_pp': 0.2, 'k_pa': 0.2, 'k_ae': 0.2, 'k_aa': 0.2, 'k_ap': 0.2,
                       'k_ps': 0.2, 'k_as': 0.2},  # AMPA blockade and GABAA blockade in GPe
                      {'k_pp': 0.2, 'k_pa': 0.2, 'k_aa': 0.2, 'k_ap': 0.2, 'k_ps': 0.2,
                       'k_as': 0.2},  # GABAA blockade in GPe
                      #{'k_pe': 0.0, 'k_ae': 0.0},  # STN blockade
                      #{'k_pe': 0.0, 'k_ae': 0.0, 'k_pp': 0.2, 'k_pa': 0.2, 'k_aa': 0.2, 'k_ap': 0.2,
                      # 'k_ps': 0.2, 'k_as': 0.2},  # STN blockade + GABAA blockade in GPe
                      {'k_ep': 0.2}  # GABAA blocker in STN
                      ]
        param_scalings = [
            ('delta_e', 'tau_e', 2.0),
            ('delta_p', 'tau_p', 2.0),
            ('delta_a', 'tau_a', 2.0),
            ('k_ee', 'delta_e', 0.5),
            ('k_ep', 'delta_e', 0.5),
            ('k_pe', 'delta_p', 0.5),
            ('k_pp', 'delta_p', 0.5),
            ('k_pa', 'delta_p', 0.5),
            ('k_ps', 'delta_p', 0.5),
            ('k_ae', 'delta_a', 0.5),
            ('k_ap', 'delta_a', 0.5),
            ('k_aa', 'delta_a', 0.5),
            ('k_as', 'delta_a', 0.5),
            ('eta_e', 'delta_e', 1.0),
            ('eta_p', 'delta_p', 1.0),
            ('eta_a', 'delta_a', 1.0),
            ]
        chunk_size = [
            #60,  # carpenters
            200,  # osttimor
            150,  # spanien
            #100,  # animals
            100,  # kongo
            100,  # tschad
            150,  # uganda
            100,  # tiber
            150,  # giraffe
            100,  # lech
            50,  # rilke
            50,  # dinkel
            #10,  # rosmarin
            #10,  # mosambik
        ]

        # perform simulations
        if len(param_grid) > 0:
            self.gs_config['init_kwargs'].update(kwargs)
            res_file = self.cgs.run(
                circuit_template=self.gs_config['circuit_template'],
                param_grid=deepcopy(param_grid),
                param_map=self.gs_config['param_map'],
                simulation_time=self.gs_config['simulation_time'],
                dt=self.gs_config['step_size'],
                inputs=self.gs_config['inputs'],
                outputs=self.gs_config['outputs'],
                sampling_step_size=self.gs_config['sampling_step_size'],
                permute=False,
                chunk_size=chunk_size,
                worker_file=worker_file,
                worker_env=self.cgs_config['worker_env'],
                gs_kwargs={'init_kwargs': self.gs_config['init_kwargs'], 'conditions': conditions,
                           'param_scalings': param_scalings},
                worker_kwargs={'freq_targets': freq_targets, 'y': target, 'time_lim': 5400.0, 'cpu_lim': False,
                               'nproc_lim': False, 'memory_lim': False, 'T': self.gs_config['simulation_time']},
                result_concat_axis=0)
            results_tmp = read_hdf(res_file, key=f'Results/results')

            # calculate fitness
            for gene_id in param_grid.index:
                self.pop.at[gene_id, 'fitness'] = 1.0 / results_tmp.at[gene_id, 'fitness']
                self.pop.at[gene_id, 'results'] = [results_tmp.at[gene_id, v] for v in result_vars]

        # set fitness of invalid parametrizations
        #for gene_id in invalid_params.index:
        #    self.pop.at[gene_id, 'fitness'] = 0.0
        #    self.pop.at[gene_id, 'results'] = [0. for _ in result_vars]

def fitness(y, t):
    y = np.asarray(y).flatten()
    t = np.asarray(t).flatten()
    diff = np.asarray([0.0 if np.isnan(t_tmp) else y_tmp - t_tmp for y_tmp, t_tmp in zip(y, t)]).flatten()**2
    t[np.isnan(t)] = 1.0
    t[t == 0] = 1.0
    weights = 1/np.abs(t)
    return weights @ diff

if __name__ == "__main__":
    warnings.filterwarnings("ignore")

    pop_size = 2048
    pop_genes = {
        'k_ee': {'min': 0, 'max': 6, 'size': pop_size, 'sigma': 0.2, 'loc': 4.0, 'scale': 1.0},
        'k_ae': {'min': 1, 'max': 100, 'size': pop_size, 'sigma': 1.0, 'loc': 50.0, 'scale': 5.0},
        'k_pe': {'min': 1, 'max': 100, 'size': pop_size, 'sigma': 1.0, 'loc': 50.0, 'scale': 5.0},
        'k_pp': {'min': 0, 'max': 6, 'size': pop_size, 'sigma': 0.2, 'loc': 3.0, 'scale': 1.0},
        'k_ep': {'min': 1, 'max': 100, 'size': pop_size, 'sigma': 1.0, 'loc': 30.0, 'scale': 3.0},
        'k_ap': {'min': 1, 'max': 100, 'size': pop_size, 'sigma': 1.0, 'loc': 50.0, 'scale': 5.0},
        'k_aa': {'min': 0, 'max': 6, 'size': pop_size, 'sigma': 0.2, 'loc': 3.0, 'scale': 1.0},
        'k_pa': {'min': 1, 'max': 100, 'size': pop_size, 'sigma': 1.0, 'loc': 50.0, 'scale': 5.0},
        'k_ps': {'min': 3, 'max': 300, 'size': pop_size, 'sigma': 2.0, 'loc': 100.0, 'scale': 10.0},
        'k_as': {'min': 3, 'max': 300, 'size': pop_size, 'sigma': 2.0, 'loc': 100.0, 'scale': 10.0},
        'eta_e': {'min': -4, 'max': 4, 'size': pop_size, 'sigma': 0.1, 'loc': 0.0, 'scale': 1.0},
        'eta_p': {'min': -4, 'max': 4, 'size': pop_size, 'sigma': 0.1, 'loc': 0.0, 'scale': 1.0},
        'eta_a': {'min': -6, 'max': 2, 'size': pop_size, 'sigma': 0.1, 'loc': -1.0, 'scale': 1.0},
        'delta_e': {'min': 0.02, 'max': 0.06, 'size': pop_size, 'sigma': 0.005, 'loc': 0.04, 'scale': 0.01},
        'delta_p': {'min': 0.15, 'max': 0.35, 'size': pop_size, 'sigma': 0.01, 'loc': 0.2, 'scale': 0.05},
        'delta_a': {'min': 0.05, 'max': 0.15, 'size': pop_size, 'sigma': 0.01, 'loc': 0.1, 'scale': 0.05},
        'tau_e': {'min': 13, 'max': 13, 'size': pop_size, 'sigma': 0.0, 'loc': 13.0, 'scale': 0.0},
        'tau_p': {'min': 25, 'max': 25, 'size': pop_size, 'sigma': 0.0, 'loc': 25.0, 'scale': 0.0},
        'tau_a': {'min': 20, 'max': 20, 'size': pop_size, 'sigma': 0.0, 'loc': 20.0, 'scale': 0.0},
    }

    param_map = {
        'k_ee': {'vars': ['weight'], 'edges': [('stn', 'stn')]},
        'k_ae': {'vars': ['weight'], 'edges': [('stn', 'gpe_a')]},
        'k_pe': {'vars': ['weight'], 'edges': [('stn', 'gpe_p')]},
        'k_pp': {'vars': ['weight'], 'edges': [('gpe_p', 'gpe_p')]},
        'k_ep': {'vars': ['weight'], 'edges': [('gpe_p', 'stn')]},
        'k_ap': {'vars': ['weight'], 'edges': [('gpe_p', 'gpe_a')]},
        'k_aa': {'vars': ['weight'], 'edges': [('gpe_a', 'gpe_a')]},
        'k_pa': {'vars': ['weight'], 'edges': [('gpe_a', 'gpe_p')]},
        'k_ps': {'vars': ['weight'], 'edges': [('str', 'gpe_p')]},
        'k_as': {'vars': ['weight'], 'edges': [('str', 'gpe_a')]},
        'eta_e': {'vars': ['stn_op/eta_e'], 'nodes': ['stn']},
        'eta_p': {'vars': ['gpe_proto_op/eta_i'], 'nodes': ['gpe_p']},
        'eta_a': {'vars': ['gpe_arky_op/eta_a'], 'nodes': ['gpe_a']},
        'delta_e': {'vars': ['stn_op/delta_e'], 'nodes': ['stn']},
        'delta_p': {'vars': ['gpe_proto_op/delta_i'], 'nodes': ['gpe_p']},
        'delta_a': {'vars': ['gpe_arky_op/delta_a'], 'nodes': ['gpe_a']},
        'tau_e': {'vars': ['stn_op/tau_e'], 'nodes': ['stn']},
        'tau_p': {'vars': ['gpe_proto_op/tau_i'], 'nodes': ['gpe_p']},
        'tau_a': {'vars': ['gpe_arky_op/tau_a'], 'nodes': ['gpe_a']},
    }

    T = 200.0
    dt = 1e-2
    dts = 1e-1

    # perform genetic optimization
    compute_dir = f"{os.getcwd()}/stn_gpe_healthy_opt"

    ga = CustomGOA(fitness_measure=fitness,
                   gs_config={
                       'circuit_template': f"{os.getcwd()}/config/stn_gpe/stn_gpe",
                       'permute_grid': True,
                       'param_map': param_map,
                       'simulation_time': T,
                       'step_size': dt,
                       'sampling_step_size': dts,
                       'inputs': {},
                       'outputs': {'r_e': "stn/stn_op/R_e",
                                   'r_i': 'gpe_p/gpe_proto_op/R_i',
                                   'r_a': 'gpe_a/gpe_arky_op/R_a'},
                       'init_kwargs': {'backend': 'numpy', 'solver': 'scipy', 'step_size': dt},
                   },
                   cgs_config={'nodes': [
                                           #'carpenters',
                                           'osttimor',
                                           'spanien',
                                           #'animals',
                                           'kongo',
                                           'tschad',
                                           'uganda',
                                           'tiber',
                                           'giraffe',
                                           'lech',
                                           'rilke',
                                           'dinkel',
                                           #'rosmarin',
                                           #'mosambik',
                                         ],
                               'compute_dir': compute_dir,
                               'worker_file': f'{os.getcwd()}/stn_gpe_healthy_worker.py',
                               'worker_env': "/data/u_rgast_software/anaconda3/envs/pyrates/bin/python3",
                               })

    drop_save_dir = f'{compute_dir}/PopulationDrops/'
    os.makedirs(drop_save_dir, exist_ok=True)

    winner = ga.run(
        initial_gene_pool=pop_genes,
        gene_sampling_func=np.random.normal,
        new_member_sampling_func=np.random.uniform,
        target=[[19, 62, 31],  # healthy control
                [np.nan, 35, np.nan],  # ampa blockade in GPe
                [np.nan, 76, np.nan],  # ampa and gabaa blockade in GPe
                [np.nan, 135, np.nan],  # GABAA blockade in GPe
                #[np.nan, 30, np.nan],  # STN blockade
                #[np.nan, 60, np.nan],  # STN blockade + gabaa blockade in GPe
                [35, 124, np.nan]  # GABAA blockade in STN
                ],
        max_iter=500,
        enforce_max_iter=True,
        min_fit=1.0,
        n_winners=30,
        n_parent_pairs=120,
        n_new=248,
        sigma_adapt=0.05,
        candidate_save=f'{compute_dir}/GeneticCGSCandidatestn.h5',
        drop_save=drop_save_dir,
        new_pop_on_drop=True,
        pop_save=f'{drop_save_dir}/pop_summary',
        permute=False,
        max_stagnation_steps=5,
        stagnation_decimals=3
    )
    #winner.to_hdf(f'{drop_save_dir}/winner.h5', key='data')


"""
Task:           Meta-Path Based Random Walk
Coder:          Haoyu Huang
Source Data:    Foursquare
the slow version
"""
import time
import datetime
import argparse
import numpy as np
import pickle as pickle
from collections import Counter
import collections
from math import radians, cos, sin, asin, sqrt
from tqdm import tqdm

class MPBasedRandomWalk:
    def __init__(self, meta_path, walk_length, walk_num, data):
        self.meta_path = meta_path
        self.walk_length = walk_length
        self.walk_num = walk_num
        self.data = data

        self.tmp_path = "../data"
        self.triple_pc_path = self.tmp_path + "/triple_pc.txt"
        self.triple_ptp_path = self.tmp_path + "/triple_ptp.txt"
        self.triple_utp_path = self.tmp_path + "/triple_utp.txt"

        self.word_loc_adlist = {}
        self.loc_word_adlist = {}  # for convenience
        self.user_loc_adlist = {}
        self.loc_user_adlist = {}  # for convenience
        self.loc_loc_adlist = {}  # dict的dict
        self.loc2neighbor_LL_dict = {}
        self.loc2neighbor_LUL_dict = {}
        self.loc2neighbor_LVL_dict = {}
        self.loc2neighbor_dict = {}  # 上面三者的merge
        self.loc2paths_LL_dict = {}
        self.loc2paths_LUL_dict = {}
        self.loc2paths_LVL_dict = {}
        self.v = {}

    # load data from three triple file,或者可以理解为生成各自的转移矩阵
    def make_adlist(self):
        print('Reading ptp file:\n')
        with open(self.triple_ptp_path, 'r') as f:
            for line in f:
                line = line.strip('\n').split('\t')
                pid1, _, pid2 = line  # 无向图，pid1->pid2 and pid2->pid1
                if pid1 not in self.loc_loc_adlist:
                    self.loc_loc_adlist.update({pid1: {pid2: 1}})
                    if pid2 not in self.loc_loc_adlist:
                        self.loc_loc_adlist.update({pid2: {pid1: 1}})
                    else:
                        if pid1 not in self.loc_loc_adlist[pid2]:
                            self.loc_loc_adlist[pid2].update({pid1: 1})
                        else:
                            self.loc_loc_adlist[pid2][pid1] += 1
                else:
                    if pid2 not in self.loc_loc_adlist[pid1]:
                        self.loc_loc_adlist[pid1].update({pid2: 1})
                    else:
                        self.loc_loc_adlist[pid1][pid2] += 1

        print("Reading utp file:\n")
        with open(self.triple_utp_path, 'r') as f:
            for line in f:
                line = line.strip('\n').split('\t')
                uid, _, pid = line
                if uid not in self.user_loc_adlist:
                    self.user_loc_adlist.update({uid: {pid: 1}})
                    if pid not in self.loc_user_adlist:
                        self.loc_user_adlist.update({pid: {uid: 1}})
                    else:
                        if uid not in self.loc_user_adlist[pid]:
                            self.loc_user_adlist[pid].update({uid: 1})
                        else:
                            self.loc_user_adlist[pid][uid] += 1
                else:
                    if pid not in self.user_loc_adlist[uid]:
                        self.user_loc_adlist[uid].update({pid: 1})
                    else:
                        self.user_loc_adlist[uid][pid] += 1

        print("Reading pc file:\n")
        with open(self.triple_pc_path, 'r') as f:
            for line in f:
                line = line.strip('\n').split('\t')
                pid, _, wid = line
                if pid not in self.loc_word_adlist:
                    self.loc_word_adlist.update({pid: {wid: 1}})
                    if wid not in self.word_loc_adlist:
                        self.word_loc_adlist.update({wid: {pid: 1}})
                    else:
                        if pid not in self.word_loc_adlist[wid]:
                            self.word_loc_adlist[wid].update({pid: 1})
                        else:
                            self.word_loc_adlist[wid][pid] += 1
                else:
                    if wid not in self.loc_word_adlist[pid]:
                        self.loc_word_adlist[pid].update({pid: 1})
                    else:
                        self.loc_word_adlist[pid][wid] += 1  # 逻辑上讲不可能出现这样的情况

    # 生成一个单一类型元素的dict
    def make_vec(self, keys):
        result = dict()
        for key in keys:
            result.update({key: 1})
        return result

    # 对单一类型的元素dict进行平均分布
    def nom_vec(self, vec):
        tmp_vec = vec.copy()  # 不改变原本数据
        total = sum(tmp_vec.values())
        for key in tmp_vec.keys():
            tmp_vec[key] = float(tmp_vec[key]) / total
        return tmp_vec

    # 起始的向量，只有一个起始点loc，也就是只有这一个元素为1
    def begin_vec(self, vec, begin):
        tmp_vec = vec.copy()
        for key in tmp_vec.keys():
            if key == begin:
                tmp_vec[key] = float(1)
            else:
                tmp_vec[key] = float(0)
        return tmp_vec

    # 对adlist进行平均分布，也就是概率
    def nom_adlist(self, adlist):
        tmp_adlist = adlist.copy()                      # value->key,为什么以有向图的方式存，因为要方便不同种类的entity进行转移

        for key1 in tqdm(tmp_adlist.keys()):            # 遍历所有的node（5000）
            count = 0
            for key2 in tmp_adlist.keys():              # 再次遍历所有的node, 计算key1的出度count
                if key1 in tmp_adlist[key2]:
                    count += 1
            for key2 in tmp_adlist.keys():              # key1->key2，1/key1的出度
                if key1 in tmp_adlist[key2]:
                    tmp_adlist[key2][key1] = float(tmp_adlist[key2][key1]) / count

        return tmp_adlist                               # 返回的结果是转移矩阵

    # 矩阵与向量相乘 (n,n)·(n,1)=(n,1),<too slow>
    def adlist_vec_multiply(self, adlist, vec):
        result = dict()
        for key1 in adlist.keys():
            tmp1 = adlist[key1]                         # 是所有指向key1的node
            tmp_sum = 0
            for key2 in vec.keys():                     # 遍历向量的每一个元素
                if key2 in tmp1:                        # 如果说向量的这个位置所对应的node，在矩阵的这一行中存在，就讲这两个值相乘
                    tmp_sum += tmp1[key2] * vec[key2]   # 矩阵这一行的乘积加起来，没有的就是0，得到结果向量的某个位置的值
            result.update({key1: tmp_sum})
        return result

    # 决定元路径
    def method(self):
        if self.meta_path == 'LL':
            for pid in tqdm(self.loc_loc_adlist):
                total_paths = self.run_with_LL(pid)
                self.loc2paths_LL_dict[pid] = total_paths
        # elif self.meta_path == 'LVL':
            for pid in tqdm(self.loc_loc_adlist):
                total_paths = self.run_with_LVL(pid)
                self.loc2paths_LVL_dict[pid] = total_paths
        # elif self.meta_path == 'LUL':
            for pid in tqdm(self.loc_loc_adlist):
                total_paths = self.run_with_LUL(pid)
                self.loc2paths_LUL_dict[pid] = total_paths

    # 三类随机游走的实现
    def run_with_LL(self, pid):
        print('LL based random walk:')
        self.loc_loc_adlist = self.nom_adlist(self.loc_loc_adlist)
        # self.v = self.nom_vec(self.make_vec(self.loc_loc_adlist.keys()))
        total_paths = []
        for i in range(self.walk_num):
            self.v = self.begin_vec(self.make_vec(self.loc_loc_adlist.keys()), pid)
            path = [pid]
            for j in range(self.walk_length):                                                 # walk_length长度的path
                self.v = self.adlist_vec_multiply(self.loc_loc_adlist, self.v)
                tmp_vec = np.array(list(self.nom_vec(self.v).values()))                       # 以当前的权重向量为权重进行random walk下一跳的选择
                next_pid = np.random.choice(list(self.v.keys()), p=tmp_vec.ravel())
                path.append(next_pid)
                self.v = self.begin_vec(self.make_vec(self.loc_loc_adlist.keys()), next_pid)  # 重置起点
            total_paths.append(path)
        return total_paths

    def run_with_LUL(self, pid):
        print('LUL based random walk:')
        self.user_loc_adlist = self.nom_adlist(self.user_loc_adlist)
        self.loc_user_adlist = self.nom_adlist(self.loc_user_adlist)
        total_paths = []
        for i in range(self.walk_num):
            self.v = self.begin_vec(self.make_vec(self.loc_loc_adlist.keys()), pid)                # 依然将pid作为起点,loc的起点权重向量
            path = [pid]
            for j in range(self.walk_length):
                self.v = self.adlist_vec_multiply(self.user_loc_adlist, self.v)                # L->U
                tmp_vec = np.array(list(self.nom_vec(self.v).values()))                        # 得到下一个uid的分布权重
                next_uid = np.random.choice(list(self.v.keys()), p=tmp_vec.ravel())            # 找到置信uid
                path.append(next_uid)
                self.v = self.begin_vec(self.make_vec(self.user_loc_adlist.keys()), next_uid)  # user的起点权重向量
                self.v = self.adlist_vec_multiply(self.loc_user_adlist, self.v)                # U->L
                tmp_vec = np.array(list(self.nom_vec(self.v).values()))                        # 得到下一个pid的分布权重
                next_pid = np.random.choice(list(self.v.keys()), p=tmp_vec.ravel())            # 找到置信pid
                path.append(next_pid)
                self.v = self.begin_vec(self.make_vec(self.loc_user_adlist.keys()), next_uid)  # loc的起点权重向量
            total_paths.append(path)
        return total_paths

    def run_with_LVL(self, pid):
        print('LVL based random walk:')
        self.loc_word_adlist = self.nom_adlist(self.loc_word_adlist)
        self.word_loc_adlist = self.nom_adlist(self.word_loc_adlist)
        total_paths = []
        for i in range(self.walk_length):
            self.v = self.begin_vec(self.make_vec(self.loc_loc_adlist.keys()), pid)
            path = [pid]
            for j in range(self.walk_length):
                self.v = self.adlist_vec_multiply(self.word_loc_adlist, self.v)                # L->V
                tmp_vec = np.array(list(self.nom_vec(self.v).values()))
                next_wid = np.random.choice(list(self.v.keys()), p=tmp_vec.ravel())
                path.append(next_wid)
                self.v = self.begin_vec(self.make_vec(self.word_loc_adlist.keys()), next_wid)
                self.v = self.adlist_vec_multiply(self.loc_word_adlist, self.v)                # V->L
                tmp_vec = np.array(list(self.nom_vec(self.v).values()))
                next_pid = np.random.choice(list(self.v.keys()), p=tmp_vec.ravel())
                path.append(next_pid)
                self.v = self.begin_vec(self.make_vec(self.loc_word_adlist.keys()), next_wid)
            total_paths.append(path)
        return total_paths

    def save_variables(self):
        paths = {'LL': self.loc2paths_LL_dict, 'LUL': self.loc2paths_LUL_dict, 'LVL': self.loc2paths_LVL_dict}
        pickle.dump(paths, open(self.tmp_path + 'paths' + '.pkl', 'wb'))

if __name__ == '__main__':
    generator = MPBasedRandomWalk(meta_path='LL', walk_length=10, walk_num=50, data=[])
    generator.make_adlist()
    generator.method()