from pyspark.sql import SparkSession
spark = SparkSession.builder.master("local[*]").getOrCreate()

import json
import numpy as np
import sys

sc = spark.sparkContext
text = sc.textFile(sys.argv[1])

df_fin =text.map(lambda x: json.loads(x))

def myFuncA(x):
    return ((x['reviewerID'], x['asin']),(x['overall'], x['unixReviewTime'], x['verified']))

OneRatingRDD = df_fin.map(lambda x: myFuncA(x)).groupByKey().map(lambda x: (x[0], list(x[1])[-1]))#.map(lambda x: func2(x))

#print(OneRatingRDD.take(2))

item25RDD = OneRatingRDD.map(lambda x: (x[0][1], (x[0][0],x[1][0]))).groupByKey().filter(lambda x: len(list(x[1]))>=25)

#print(item25RDD.take(2))

def interchange(x):
    result = []
    for i in x[1]:
        result.append((i[0], (x[0], i[1])))
    return result

User5itemsRDD = item25RDD.flatMap(lambda x: interchange(x)).groupByKey().filter(lambda x: len(list(x[1]))>=5)
#User5itemsRDD = OneRatingRDD.map(lambda x: (x[0][0], (x[0][1],x[1][0]))).groupByKey().filter(lambda x: len(list(x[1]))>=5)

#print(User5itemsRDD.take(1))

itemRDD = User5itemsRDD.flatMap(lambda x: interchange(x)).groupByKey()

#print(len((itemRDD.collect())))

def latFunc(x):
    #result = []
    count =0
    # userList = []
    # for i in fwBC.value:
    #     userList.append(i[0])
    for i in x[1]:
        if i[0] in fwBC.value:
            count+=1

    if count>1:
        return x

def meanCenter(x):
    result = []
    total = 0
    for i in x[1]:
        #i2 = list(i)
        total+= i[1]
    mean = float(total)/len(x[1])

    for i in x[1]:
        i2 = list(i)
        i2[1] = i2[1]-mean
        result.append((i2[0], i2[1], i2[1]+mean))

    return (x[0],result)

def cosine(x):
    l = meanList.value
    item1 = l[0][0]
    item2 = x[0]
    dot =0
    sumA = 0
    sumB = 0
    sim =0
    if item1 != item2:
        for a in l[0][1]:
            for b in x[1]:
                if b[0] == a[0]:
                    dot += a[1]*b[1]
        for a in l[0][1]:
            sumA += a[1]*a[1]
        for b in x[1]:
            sumB += b[1]*b[1]
            #sumB = 1

        sim = float(dot)/(np.sqrt(sumA) * np.sqrt(sumB))

    return ((x[0], sim), x[1])

def filterUsers(x):
    myList = []
    for itemId in itemList.value:
        myList.append(itemId[0][0])
    #result = []
    count = 0
    if(x[0] not in fwBC.value):
        for i in x[1]:
            if i[0] in myList:
                count+=1
        if(count>=2):
            return x
    #return result

def utility(x):
    num =0
    den =0
    for a, b in itemList.value:
        for record in b:
            if(record[0]==x[0]):
                num += a[1]*record[2]
                den += a[1]
    result = num/den

    return (x[0], result)

input  = sys.argv[2]
inputList = map(str, input.strip('[]').split(','))
for each in inputList:
    myItem = each.strip()
    myItem = myItem.replace("'", '')
    #print(myItem)
    broadv= itemRDD.filter(lambda x : x[0] == myItem).flatMap(lambda x: list(x[1])).map(lambda x: x[0]).collect()
    fwBC = sc.broadcast(broadv)
    neighborRDD = itemRDD.map(lambda x: (x[0], list(x[1]))).filter(lambda x: latFunc(x)).map(lambda x: meanCenter(x)) # a (ii)
    broadList = neighborRDD.filter(lambda x: x[0] == myItem).collect()
    meanList = sc.broadcast(broadList)
    cosineRDD = neighborRDD.map(lambda x: cosine(x)).filter(lambda x: x[0][1]>0)
    mynewRDD = sc.parallelize(cosineRDD.take(50))
    fItems = mynewRDD.collect()
    itemList = sc.broadcast(fItems)
    allUsers = User5itemsRDD.filter(lambda x: filterUsers(x))
    answerRDD = allUsers.map(lambda x: utility(x))#.sortByKey()
    #print(len(answerRDD.collect()))
    print('')
    print('Product id: '+myItem)
    for x in answerRDD.collect():
        print(x)
    print('\n \n')


from __future__ import print_function
import os

import h5py
import numpy as np
import copy
from collections import OrderedDict
from simnibs.simulation.tms_coil.tms_coil import TmsCoil

from simnibs.utils.mesh_element_properties import ElementTags

from . import pygpc
from ..mesh_tools import mesh_io
from .sim_struct import SimuList
from . import fem
from . import coil_numpy as coil
from ..utils.simnibs_logger import logger

FIELD_NAME = {'v': 'v', 'E': 'E', 'e': 'magnE', 'J': 'J', 'j': 'magnJ'}

def write_data_hdf5(data, data_name, hdf5_fn, path='data/', compression='gzip'):
    ''' Saves a field in an hdf5 file

    Parameters:
    ---------------
    data: np.array
        data to be saved
    data_name: str
        name of data
    path: str (optional)
        path inside hdf5 file (default: data/)
    '''
    with h5py.File(hdf5_fn, 'a') as f:
        f.create_dataset(path + data_name, data=data, compression=compression)

def read_data_hdf5(data_name, hdf5_fn, path='data/'):
    ''' Saves a field in an hdf5 file

    Parameters:
    ---------------
    data_name: str
        name of data
    hdf5_fn: str
        name of hdf5 file
    path: str (optional)
        path inside hdf5 file (default: data/)
    '''
    with h5py.File(hdf5_fn, 'r') as f:
        return np.squeeze(f[path + data_name])

class gPC_regression(pygpc.RegularizedRegression):
    ''' Class defining the parameters of a gPC regression

    Inherits from pygpc's "reg" class

    Attributes:
    ------------------
    pdftype: list
        List of strings with PDF types. should be 'beta' or 'normal'
        uniform distributions are beta distributed with both parameters = 1
    pdfshape: list of lists
        Parameters for pdf distributions [[p1, p2, mean1, ...], [q1, q2, var1, ...]]
    limits: list of lists
        Limits for pdfs [[min1, min2, None, ..], [max1, max2, None, ...]]
        Normal ditributions do not have limits, therefore both are set to None
    poly_idx: np.ndarray
        Indices (i, j, k ,...) for each polynomial
    order: int
        Raximum individual expansion order list [1 x DIM]
        generates individual polynomials also if maximum expansion order
        in order_max is exceeded
    order_max: int
        maximum expansion order (sum of all exponents) [1]
        the maximum expansion order considers the sum of the
        orders of combined polynomials only
    interaction_order: int
        number of random variables, which can interact with each other [1]
        all polynomials are ignored, which have an interaction order greater than the specified
    grid: pygpc.grid
        grid object gnerated in pygpc.grid, including grid.coords and grid.coords_norm
    data_file: str
        Path to file with regression data
    Parameters
    ----------------
    random_vars: list
        List of random variables. Integers is conductivities, else strings
    pdftype: list
        List of strings with PDF types. should be 'beta' or 'normal'
        uniform distributions are beta distributed with both parameters = 1
    pdfshape: list of lists
        Parameters for pdf distributions [[p1, p2, mean1, ...], [q1, q2, var1, ...]]
    limits: list of lists
        Limits for pdfs [[min1, min2, None, ..], [max1, max2, None, ...]]
        Normal ditributions do not have limits, therefore both are set to None
    poly_idx: np.ndarray
        Indices (i, j, k ,...) for each polynomial
    coords_norm: np.ndarray
        List of sampled points in normalized space
    sim_type: {'TMS', 'TCS'}
        Type of siulation
    data_file: str, optional
        Path to file with raw data
    '''
    def __init__(self, random_vars, pdftype, pdfshape, limits, poly_idx, coords_norm,
                 sim_type, data_file=None, regularization_factors=[0.]):
        if isinstance(pdftype, str):
            raise ValueError('pdftype must be a list of strings')

        if sim_type in ['TMS', 'TCS']:
            self._sim_type = sim_type
        else:
            raise ValueError('invalid sim_type')

        self.random_vars = random_vars
        # initialize grid class
        grid = pygpc.randomgrid(pdftype=pdftype,
                                gridshape=pdfshape,
                                limits=limits,
                                N=0)

        # overwrite grid information
        grid.coords = pygpc.denorm(coords_norm, pdftype, pdfshape, limits)
        grid.coords_norm = coords_norm

        # initialize reg class
        super().__init__(pdftype, pdfshape, limits, [0]*len(pdftype), 0, len(pdftype),
                         grid=grid, regularization_factors=regularization_factors)
        # enrich polynomial basis
        self.enrich_polynomial_basis(poly_idx, form_A=False)

        # set up data (HDF5) file
        self.data_file = data_file

    @property
    def sim_type(self):
        '''Simulation type, TMS or TCS'''
        return self._sim_type

    @property
    def mesh_file(self):
        ''' Name of mesh file '''
        return os.path.splitext(self.data_file)[0] + '.msh'

    def postprocessing(self, postprocessing_type, order_sobol_max=1):
        ''' Postprocessing for TMS gPC operation
        Makes the operation for the postprodesssing and saves in in the data_file

        Parameters
        -----------------------
        postprocessing_type: str
            A combination of 'v'(potential), 'E'(electric field vector),
            'e'(electric field magnitude), 'J'(current density vector) and
            'j'(current density magnitude) eg: 'eEJ'

        order_sobol_max: int (Optional)
            Maximum number of components for Sobol coefficients. Default: 1
        Returns
        -----------------------
        Writes mean, std, expansion coefficiets, sobol coefficiets, sensitivity and
        global sensitivity to file
        '''

        if self.data_file is None:
            raise ValueError('Please set a data_file before running the postprocessing')
        if not os.path.exists(self.data_file):
            raise ValueError('Could not find the data_file: {0}'.format(self.data_file))
        # make the vector into a list
        postprocessing_type = [p for p in postprocessing_type]
        for p in postprocessing_type:
            if p not in ['v', 'E', 'e', 'J', 'j']:
                raise ValueError(
                    'Unrecognized postprocessing type: {0}'.format(p))

        # load the cropped mesh
        msh = mesh_io.Msh.read_hdf5(self.data_file, 'mesh_roi/')

        lst = SimuList()
        lst._get_conductivity_from_hdf5(self.data_file)
        lst.mesh = msh

        potentials = read_data_hdf5('v_samples', self.data_file,
                                    'mesh_roi/data_matrices/')

        nr_simu = potentials.shape[0]

        '''
        if 'v' in postprocessing_type:
            self._postprocessing_core(potentials, 'mesh_roi/nodedata', False)
'mesh_roi/' + dtype + '/',
                FIELD_NAME[p],
                order_sobol_max
            postprocessing_type.remove('v')
        '''

        # This will NOT work when changing positions
        if self.sim_type == 'TMS':
            dAdt = read_data_hdf5('dAdt', self.data_file, 'mesh_roi/elmdata/')
            dAdt = mesh_io.ElementData(dAdt, mesh=msh)
        else:
            dAdt = None

        # See which fields have already been calculated
        fields_dict = dict.fromkeys(postprocessing_type)
        to_calc = []
        for postprocess in postprocessing_type:
            try:
                fields_dict[postprocess] = read_data_hdf5(
                    FIELD_NAME[postprocess] +
                    '_samples', self.data_file, 'mesh_roi/data_matrices/')
            except:
                if postprocess in ['e', 'j']:
                    fields_dict[postprocess] = np.nan * np.ones((nr_simu, msh.elm.nr), dtype=float)
                elif postprocess in ['E', 'J']:
                    fields_dict[postprocess] = np.nan * np.ones((nr_simu, msh.elm.nr, 3), dtype=float)
                elif postprocess in ['v']:
                    fields_dict[postprocess] = np.nan * np.ones((nr_simu, msh.nodes.nr), dtype=float)
                else:
                    raise ValueError('Unrecognized postprocessing option: ' + postprocess)
                to_calc.append(postprocess)

        # Calculate the ones which have not been calculated
        # This step is slow! I've opted for greater flexibility and reliability rather
        # than just speed. it uses calc_fields
        if len(to_calc) > 0:
            logger.info('Calculating fields: {0}'.format(''.join(to_calc)))
            for i, phi, gr in zip(range(nr_simu), potentials, self.grid.coords):
                # set the conductivities right
                for rv, g in zip(self.random_vars, gr):
                    if isinstance(rv, int):
                        lst.cond[rv - 1].value = g
                    elmdata = lst.cond2elmdata(logger_level=10)
                # sets the potential
                pot = mesh_io.NodeData(phi, mesh=msh)
                # calclate the remaining fields
                # We can use 'E' to calculate the remaining fields, if it has been
                # defined
                if 'E' in fields_dict.keys() and not np.all(np.isnan(fields_dict['E'][i])):
                    m = fem.calc_fields(pot, to_calc, cond=elmdata, dadt=dAdt,
                                        E=fields_dict['E'][i])
                #Else we also need to calculate E
                else:
                    m = fem.calc_fields(pot, to_calc, cond=elmdata, dadt=dAdt)

                for p in to_calc:
                    fields_dict[p][i] = m.field[FIELD_NAME[p]].value

        for p, f in fields_dict.items():
            logger.info('Expanding field: {0}'.format(FIELD_NAME[p]))
            if p == 'v':
                dtype = 'nodedata'
            else:
                dtype = 'elmdata'
            self._postprocessing_core(
                f, 'mesh_roi/' + dtype + '/',
                FIELD_NAME[p],
                order_sobol_max)

    def _postprocessing_core(self, data, path, name, order_sobol_max):
        ''' Convinience function to calculate postprocessing output '''
        with h5py.File(self.data_file, 'a') as f:
            try:
                f.create_dataset(
                    'mesh_roi/data_matrices/' + name + '_samples',
                    data=data, compression='gzip')
            except (RuntimeError, OSError, ValueError):
                pass

        data_dims = data.shape[1:]
        if data.ndim == 3:
            data = data.reshape(data.shape[0], -1)
        coeffs, error = self.expand(data)
        logger.info('Estimated error: {0:1e} for field {1}'.format(error, name))

        # Mean and std
        mean = self.mean(coeffs)
        std = self.std(coeffs)
        # Sobol
        sobol, sobol_idx = self.sobol(coeffs)
        # Filter sobol
        sobol_order = np.array([len(idx) for idx in sobol_idx])
        sobol = sobol[sobol_order <= order_sobol_max]
        sobol_idx = [s_idx for s_idx in sobol_idx if len(s_idx) <= order_sobol_max]
        # Sensitivity
        sens = self.globalsens(coeffs)

        # Put everything in Data structures
        mean = mesh_io.Data(mean.reshape(*data_dims), name=name + '_mean')
        mean.write_hdf5(self.data_file, path=path)
        std = mesh_io.Data(std.reshape(*data_dims), name=name + '_std')
        std.write_hdf5(self.data_file, path=path)
        for s, s_idx in zip(sobol, sobol_idx):
            s_name = name + '_sobol_' + \
                    '_'.join([str(self.random_vars[s_i]) for s_i in s_idx])
            sobol_dat = mesh_io.Data(s.reshape(*data_dims), name=s_name)
            sobol_dat.write_hdf5(self.data_file, path=path)

        for i, s in enumerate(sens):
            s_name = name + '_sensitivity_' + str(self.random_vars[i])
            sens_dat = mesh_io.Data(s.reshape(*data_dims), name=s_name)
            sens_dat.write_hdf5(self.data_file, path=path)

    def save_hdf5(self, data_file=None):
        """ Saves the gPC information in an hdf5 file
        pdftype, pdfshape, limits, poly_idx and coords_norm are saved
        in "gpc_object" in the data file

        Parameters:
        -----------------
        data_file (optional): str
            name of data file. default: self.data_file

        """
        if data_file is None:
            data_file = self.data_file
        if data_file is None:
            raise ValueError('Please specify a data file')
        with h5py.File(data_file, 'a') as f:
            f.attrs['type'] = self.sim_type
            if 'gpc_object' in f.keys():
                del f['gpc_object']
            f.create_dataset('gpc_object/random_vars',
                             data=np.array(self.random_vars, dtype=np.string_))
            f.create_dataset('gpc_object/pdftype',
                             data=np.array(self.pdftype, dtype='S10'))
            f.create_dataset('gpc_object/pdfshape',
                             data=np.array(self.pdfshape))
            lim = [[], []]
            lim[0] = [l if l is not None else -1e10 for l in self.limits[0]]
            lim[1] = [l if l is not None else 1e10 for l in self.limits[1]]
            f.create_dataset('gpc_object/limits',
                             data=np.array(lim))
            f.create_dataset('gpc_object/poly_idx',
                             data=np.array(self.poly_idx))
            f.create_dataset('gpc_object/grid/coords_norm',
                             data=np.array(self.grid.coords_norm))
            f.create_dataset('gpc_object/regularization_factors',
                             data=np.array(self.regularization_factors))

    @classmethod
    def read_hdf5(cls, fn_hdf5):
        """ Reads gPC information from hdf5 file
        Information must have the same format as in gPC_regression.save_hdf5

        Parameteres:
        -----------------
        fn_hdf5: str
            Name of hdf5 file

        Returns:
        -----------------
        gPC_regression: gPC_regression
            regression object
        """
        with h5py.File(fn_hdf5, 'r') as f:
            sim_type = f.attrs['type']
            random_vars = f['gpc_object/random_vars'][()].tolist()
            processed = []
            for rv in random_vars:
                try:
                    processed.append(int(rv.decode()))
                except(AttributeError, ValueError):
                    processed.append(rv.decode)
            random_vars = processed
            pdftype = f['gpc_object/pdftype'][()].tolist()
            pdftype = [s.decode() for s in pdftype]
            pdfshape = f['gpc_object/pdfshape'][()].tolist()
            limits = f['gpc_object/limits'][()].tolist()
            limits[0] = [None if np.isclose(l, -1e10) else l for l in limits[0]]
            limits[1] = [None if np.isclose(l, 1e10) else l for l in limits[1]]
            poly_idx = f['gpc_object/poly_idx'][()]
            coords_norm = f['gpc_object/grid/coords_norm'][()]
            regularization_factors = f['gpc_object/regularization_factors'][()]

        return cls(random_vars, pdftype, pdfshape, limits,
                   poly_idx, coords_norm, sim_type, data_file=fn_hdf5)
                   #regularization_factors=regularization_factors)

    def visualize(self):
        ''' Creates a mesh file for visualization

        Returns:
        --------
        writes a mesh file in the same folder as the hdf5 file
        msh: simnibs.msh.mesh_io
            mesh with gpc fields
        '''
        msh = mesh_io.Msh.read_hdf5(self.data_file, path='mesh_roi/')
        mesh_io.write_msh(msh, self.mesh_file)
        return msh

    def expand_quantity(self, func=None, field='E'):
        ''' Expand an arbitrary quantity

        Parameters
        --------......
        func (optional): function
            Function which takes up a single argument and returns a single number or a
            vector to be expanded by gpc.
            The arguments corresponds the the field samples in the format
            [N_simulations x N_roi x 3] for vector fields
            [N_simulations x N_roi] for scalar fields

        field (Optinal): 'v', 'e', 'E', 'J' or 'j'
            field to be passed as an argument to the function, must have been previously
            calculated suing the postprocessing method. Default: E

        Returns
        --------
        coeffs: ndarray
            gPC polynomial coefficients. You can then call the attributes mean, std,
            sobol, and globalsens to calculate the mean, standard deviation, sobol coefficients
            and sensitivity of you quantity of interest
        '''
        try:
            field = read_data_hdf5(
                FIELD_NAME[field] + '_samples',
                self.data_file, 'mesh_roi/data_matrices/'
            )
        except KeyError:
            field = read_data_hdf5(
                field + '_samples',
                self.data_file, 'mesh_roi/data_matrices/'
            )
        except RuntimeError:
            raise IOError(
                f'Could not read field "{field}" in file "{self.data_file}"'
            )
        if func:
            f = func(field)
        else:
            f = field.reshape(field.shape[0], -1)
        coeffs, cv = self.expand(f)
        logger.info('CV value: {0:1e}'.format(cv))
        return coeffs

    def roi_mesh(self):
        ''' Returns the mesh where the expansion is defined '''
        return mesh_io.Msh.read_hdf5(self.data_file, 'mesh_roi/')

def prep_gpc(simlist):
    cond = simlist.cond
    random_vars = []
    pdf_types = []
    parameters = []

    for i, c in enumerate(cond):
        if c.distribution_type is not None:
            random_vars.append(i + 1)
            pdf_types.append(c.distribution_type)
            parameters.append(c.distribution_parameters)

    if len(random_vars) == 0:
        raise ValueError('No random variables found for simulation')

    # assign variables in a way gpc understands
    limits = [[], []]
    pdfshape = [[], []]
    for pdf, pars in zip(pdf_types, parameters):
        if pdf == 'uniform':
            if len(pars) != 2:
                raise ValueError('uniform random variables must have 2 parameters')
            pdfshape[0].append(1)
            pdfshape[1].append(1)
            limits[0].append(pars[0])
            limits[1].append(pars[1])
        elif pdf == 'beta':
            if len(pars) != 4:
                raise ValueError('Beta random variables must have 4 parameters')
            pdfshape[0].append(pars[0])
            pdfshape[1].append(pars[1])
            limits[0].append(pars[2])
            limits[1].append(pars[3])
        elif pdf == 'normal':
            if len(pars) != 2:
                raise ValueError('Normal random variables must have 2 parameters')
            pdfshape[0].append(pars[0])
            pdfshape[1].append(pars[1])
            limits[0].append(None)
            limits[1].append(None)
        else:
            raise ValueError('Invalid distribution_type: {0}'.format(pdf))

    # change "uniform" to "beta"
    pdf_type = ['beta' if 'uniform' == p else p for p in pdf_types]
    return random_vars, pdf_type, pdfshape, limits

def run_tms_gpc(poslist, fn_simu, cpus=1, tissues=[ElementTags.GM], eps=1e-2,
                max_iter=1000, min_iter=2, data_poly_ratio=2):
    ''' Runs one TMS gPC for each position in the current TMSLIST

    Parameters
    ------------
    poslist: simnibs.sim_struct.TDCSLIST
        TDCSLIST structure defining the simulation
    fn_simu: str
        Output name
    cpus: int (optional)
        Number of CPUs to use
    tissues: list (Optional)
        List of tissue tags where to evaluate the electric field. Default: [2]
    eps: float (optional)
        Tolerance for gPC expansions. Default:1e-2
    max_iter: int (optinal)
        Maximum number of adaptive gPC expansion. Defaut:1000
    min_iter: int (optinal)
        Minimum number of adaptive gPC expansion interations. Defaut:2
    data_poly_ratio: int
        Ratio of number of new simulation per new polynomial. Default:2

    Returns
    --------
    fns: list
        List of mesh file names
    '''
    poslist._prepare()
    fn_simu = os.path.abspath(os.path.expanduser(fn_simu))
    if cpus > 1:
        logger.warning("Can't run GPC with multiprocessing (for now)")
        logger.warning("Setting cpus=1")
        cpus = 1

    logger.info('Running a gPC expansion with tolerance: {0:1e}'.format(eps))
    # run simulations
    random_vars, pdf_type, pdfshape, limits = prep_gpc(poslist)
    path, basename = os.path.split(fn_simu)
    if not os.path.isdir(path) and path != '':
        os.mkdir(path)

    fns = []
    for i, p in enumerate(poslist.pos):
        fn_hdf5 = fn_simu+'_{0:0=4d}_gpc.hdf5'.format(i + 1)
        if os.path.isfile(fn_hdf5):
            raise IOError('Output file ' + fn_hdf5 + ' already exists')
        matsimnibs = p.calc_matsimnibs(poslist.mesh)
        sampler = TMSgPCSampler(
            poslist.mesh, poslist, fn_hdf5,
            poslist.fnamecoil, matsimnibs, p.didt,
            roi=tissues)
        sampler.create_hdf5()
        reg, phi = pygpc.adaptive.run_reg_adaptive_grid(
            pdf_type, pdfshape, limits,
            sampler.run_simulation,
            data_poly_ratio=data_poly_ratio,
            max_iter=max_iter,
            eps=eps,
            n_cpus=cpus,
            print_function=logger.info,
            min_iter=min_iter)
        gpc_reg = gPC_regression(random_vars,
                                 pdf_type, pdfshape, limits, reg.poly_idx,
                                 reg.grid.coords_norm, 'TMS',
                                 data_file=fn_hdf5)
        gpc_reg.save_hdf5()
        gpc_reg.postprocessing(poslist.postprocess)
        gpc_reg.visualize()
        fns.append(gpc_reg.mesh_file)

    return fns

def run_tcs_gpc(poslist, fn_simu, cpus=1, tissues=[2], eps=1e-2,
                max_iter=1000, min_iter=2, data_poly_ratio=2,
                regularization_factors=[0.]):
    ''' Runs a tDCS gPC expansion

    Parameters
    ------------
    poslist: simnibs.sim_struct.TDCSLIST
        TDCSLIST structure defining the simulation
    fn_simu: str
        Output name
    cpus: int (optional)
        Number of CPUs to use
    tissues: list (Optional)
        List of tissue tags where to evaluate the electric field. Default: [2]
    eps: float (optional)
        Tolerance for gPC expansions. Default:1e-2
    max_iter: int (optinal)
        Maximum number of adaptive gPC expansion iterations. Defaut:1000
    min_iter: int (optinal)
        Minimum number of adaptive gPC expansion interations. Defaut:2
    data_poly_ratio(optional): int
        Ratio of number of new simulation per new polynomial. Default:2

    Returns
    --------
    fns: list
        List of mesh file names
    '''
    poslist._prepare()
    fn_simu = os.path.abspath(os.path.expanduser(fn_simu))
    if cpus > 1:
        logger.warning("Can't run GPC with multiprocessing (for now)")
        logger.warning("Setting cpus=1")
        cpus = 1
    logger.info('Running a gPC expansion with tolerance: {0:1e}'.format(eps))
    fn_hdf5 = fn_simu+'_gpc.hdf5'
    if os.path.isfile(fn_hdf5):
        raise IOError('Output file ' + fn_hdf5 + ' already exists')
    random_vars, pdf_type, pdfshape, limits = prep_gpc(poslist)

    path, basename = os.path.split(fn_simu)
    if not os.path.isdir(path) and path != '':
        os.mkdir(path)
    # place electrodes
    fn_no_extension, extension = os.path.splitext(fn_simu)
    m, electrode_surfaces = poslist._place_electrodes()
    mesh_io.write_msh(m, fn_simu + '_electrodes.msh')
    sampler = TDCSgPCSampler(
        m, poslist, fn_simu+'_gpc.hdf5',
        electrode_surfaces, poslist.currents,
        roi=tissues)
    sampler.create_hdf5()
    reg, phi = pygpc.adaptive.run_reg_adaptive_grid(
        pdf_type, pdfshape, limits,
        sampler.run_simulation,
        data_poly_ratio=data_poly_ratio,
        max_iter=max_iter,
        eps=eps,
        regularization_factors=regularization_factors,
        n_cpus=cpus,
        print_function=logger.info,
        min_iter=min_iter)
    gpc_reg = gPC_regression(random_vars,
                             pdf_type, pdfshape, limits, reg.poly_idx,
                             reg.grid.coords_norm, 'TCS',
                             data_file=fn_hdf5)
    gpc_reg.save_hdf5()
    gpc_reg.postprocessing(poslist.postprocess)
    gpc_reg.visualize()

    return [gpc_reg.mesh_file]

class gPCSampler(object):
    ''' Object used by pygpc to sample

    Attributes
    -----------
    mesh: simnibs.msh.mesh_io.Msh
        Mesh structure
    roi: list of integers
        List of tags defining the ROI
    mesh_roi: simnibs.msh.mesh_io.Msh
        Mesh of the ROI only
    fn_hdf5: string
        Name of hdf5 file with simulation output
    poslist: simnibs.simulation.sim_struct.SimuList
        Structure where the conductivity is defined
    identifiers: list
        List of random variable identifiers
    qoi_function: OrderedDict
        dictionaty with functions for each QOI.
        The first QOI will be passed to the gPC algorithm.

    Parameters
    ----------
    mesh: simnibs.msh.mesh_io.Msh
        Mesh structure
    poslist: simnibs.simulation.sim_struct.SimuList
        Structure where the conductivity is defined
    fn_hdf5: string
        Name of hdf5 file with simulation output
    roi: list of integers (Optional)
        List of tags defining the ROI. Default: [2]
    '''
    def __init__(self, m, poslist, fn_hdf5, roi=[2]):
        self.mesh = m
        self.roi = roi
        if self.roi is not None:
            self.mesh_roi = m.crop_mesh(roi)
        else:
            self.mesh_roi = m
            self.roi = np.unique(m.elm.tag1).tolist()
        self.fn_hdf5 = fn_hdf5
        self.poslist = poslist
        self._gpc_vars = prep_gpc(poslist)
        self.identifiers = self._gpc_vars[0]
        self.qoi_function = OrderedDict([('E', self._calc_E)])

    def create_hdf5(self):
        '''Creates an HDF5 file to store the data '''
        # if the hdf5 file does not exist, create it
        file_exists = os.path.exists(self.fn_hdf5)
        if file_exists:
            raise IOError('Cannot create hdf5 file: {0} '
                          'it already exists!'.format(self.fn_hdf5))

        self.mesh.write_hdf5(self.fn_hdf5, 'mesh/')
        self.mesh_roi.write_hdf5(self.fn_hdf5, 'mesh_roi/')
        self.poslist._write_conductivity_to_hdf5(self.fn_hdf5)
        with h5py.File(self.fn_hdf5, 'a') as f:
            f.create_dataset('roi', data=np.atleast_1d(np.array(self.roi, dtype=int)))

    @classmethod
    def load_hdf5(cls, fn_hdf5):
        '''Loads structure from hdf5 file '''
        mesh = mesh_io.Msh.read_hdf5(fn_hdf5, 'mesh/')
        poslist = SimuList()
        poslist._get_conductivity_from_hdf5(fn_hdf5)
        with h5py.File(fn_hdf5, 'r') as f:
            roi = f['roi'][()].tolist()
        return cls(mesh, poslist, fn_hdf5, roi)

    def record_data_matrix(self, data, name, group):
        ''' Appends or create data to the HDF5 file

        Parameters:
        -------------
        data: np.ndarray
            Data to be appended. Will be appended along the first dimension
        name: str
            Name of data seet
        group: str
            Group where to place data set
        '''
        data = np.array(data).squeeze()
        data = np.atleast_1d(data)
        with h5py.File(self.fn_hdf5, 'a') as f:
            try:
                g = f.create_group(group)
            except:
                g = f[group]
            if name not in g.keys():
                g.create_dataset(name,
                                 shape=(0, ) + data.shape,
                                 maxshape=(None, ) + data.shape,
                                 dtype=data.dtype,
                                 chunks=(1, ) + data.shape)

            dset = g[name]
            dset.resize((dset.shape[0] + 1, ) + data.shape)
            dset[-1, ...] = data

    def run_simulation(self, random_vars):
        raise NotImplementedError('This method is to be implemented in a subclass!')

    def _calc_E(self, v, random_vars, dAdt=None):
        grad = v.gradient()
        grad.assign_triangle_values()
        E = -grad.value * 1e3
        if dAdt is not None:
            E -= dAdt.value
        return E

    def _update_poslist(self, random_vars):
        poslist = copy.deepcopy(self.poslist)
        for i, iden in enumerate(self.identifiers):
            if type(iden) == int:
                poslist.cond[iden-1].value = random_vars[i]
        return poslist

    def run_N_random_simulations(self, N):
        grid = pygpc.randomgrid(pdftype=self._gpc_vars[1],
                                gridshape=self._gpc_vars[2],
                                limits=self._gpc_vars[3],
                                N=N)
        for i, x in enumerate(grid.coords):
            logger.info('Running simulation {0} out of {1}'.format(i + 1, N))
            self.run_simulation(x)

class TDCSgPCSampler(gPCSampler):
    ''' Object used by pygpc to sample a tDCS problem

    Attributes
    -----------
    mesh: simnibs.msh.mesh_io.Msh
        Mesh structure
    roi: list of integers
        List of tags defining the ROI
    mesh_roi: simnibs.msh.mesh_io.Msh
        Mesh of the ROI only
    fn_hdf5: string
        Name of hdf5 file with simulation output
    poslist: simnibs.simulation.sim_struct.SimuList
        Structure where the conductivity is defined
    gpc_vars: list
        List with variables for inputing to pygpc
    identifiers: list
        List of random variable identifiers
    qoi_function: OrderedDict
        dictionaty with functions for each QOI.
        The first QOI will be passed to the gPC algorithm, flattened.
        The QOI function should take only one argument (the potential v)
    el_tags: list of integers
        List of integers with surface tags of electrodes
    el_currents: list of floats
        List with current values for each electrode (in A)

    Parameters
    ----------
    mesh: simnibs.msh.mesh_io.Msh
        Mesh structure
    poslist: simnibs.simulation.sim_struct.SimuList
        Structure where the conductivity is defined
    fn_hdf5: string
        Name of hdf5 file with simulation output
    el_tags: list of integers
        List of integers with surface tags of electrodes
    el_currents: list of floats
        List with current values for each electrode (in A)
    roi: list of integers (Optional)
        List of tags defining the ROI. Default: [2]

    '''

    def __init__(self, mesh, poslist, fn_hdf5, el_tags, el_currents, roi=[2]):
        super(TDCSgPCSampler, self).__init__(mesh, poslist, fn_hdf5, roi=roi)
        self.el_tags = el_tags
        self.el_currents = el_currents

    def create_hdf5(self):
        super(TDCSgPCSampler, self).create_hdf5()
        with h5py.File(self.fn_hdf5, 'a') as f:
            f.create_dataset('el_tags', data=np.array(self.el_tags, dtype=int))
            f.create_dataset('el_currents', data=np.array(self.el_currents,
                                                          dtype=float))

    @classmethod
    def load_hdf5(cls, fn_hdf5):
        s = gPCSampler.load_hdf5(fn_hdf5)
        with h5py.File(fn_hdf5, 'r') as f:
            el_tags = f['el_tags'][()].tolist()
            el_currents = f['el_currents'][()].tolist()
        return cls(
            s.mesh, s.poslist, s.fn_hdf5, el_tags, el_currents,
            roi=s.roi)

    def run_simulation(self, random_vars):
        poslist = self._update_poslist(random_vars)
        cond = poslist.cond2elmdata(self.mesh)
        v = fem.tdcs(
            self.mesh, cond, self.el_currents,
            self.el_tags, units='mm')

        self.mesh.nodedata = [v]
        cropped = self.mesh.crop_mesh(self.roi)
        v_c = cropped.nodedata[0]
        self.mesh.nodedata = []

        qois = []
        for qoi_name, qoi_f in self.qoi_function.items():
            qois.append(qoi_f(v_c, random_vars))

        self.record_data_matrix(random_vars, 'random_var_samples', '/')
        self.record_data_matrix(v.value, 'v_samples', 'mesh/data_matrices')
        self.record_data_matrix(v_c.value, 'v_samples', 'mesh_roi/data_matrices')
        for qoi_name, qoi_v in zip(self.qoi_function.keys(), qois):
            self.record_data_matrix(
                qoi_v, qoi_name + '_samples', 'mesh_roi/data_matrices')

        del cropped
        del cond
        del v
        del v_c

        return np.atleast_1d(qois[0]).reshape(-1)

class TMSgPCSampler(gPCSampler):
    ''' Object used by pygpc to sample a TMS problem

    Attributes
    -----------
    mesh: simnibs.msh.mesh_io.Msh
        Mesh structure
    roi: list of integers
        List of tags defining the ROI
    mesh_roi: simnibs.msh.mesh_io.Msh
        Mesh of the ROI only
    fn_hdf5: string
        Name of hdf5 file with simulation output
    poslist: simnibs.simulation.sim_struct.SimuList
        Structure where the conductivity is defined
    gpc_vars: list
        List with variables for inputing to pygpc
    identifiers: list
        List of random variable identifiers
    qoi_function: OrderedDict
        dictionaty with functions for each QOI.
        The first QOI will be passed to the gPC algorithm.
        The QOI function should take only 2 arguments: (the potential v and dAdt)
    matsimnibs: np.ndarray
        Matrix defining coil position
    dIdt: float
        Current intensity in coil

    Parameters
    ----------
    mesh: simnibs.msh.mesh_io.Msh
        Mesh structure
    poslist: simnibs.simulation.sim_struct.SimuList
        Structure where the conductivity is defined
    fn_hdf5: string
        Name of hdf5 file with simulation output
    pos: simnibs.simulation.sim_struct.POS
       Coil position definition
    roi: list of integers (Optional)
        List of tags defining the ROI. Default: [2]

    '''

    def __init__(self, mesh, poslist, fn_hdf5, fnamecoil, matsimnibs, didt, roi=[2]):
        super(TMSgPCSampler, self).__init__(
            mesh, poslist, fn_hdf5, roi=roi)
        self.matsimnibs = np.ascontiguousarray(matsimnibs)
        self.didt = didt
        self.fnamecoil = fnamecoil
        self.constant_dAdt = True

    def create_hdf5(self):
        super(TMSgPCSampler, self).create_hdf5()
        with h5py.File(self.fn_hdf5, 'a') as f:
            f.create_dataset('matsimnibs', data=self.matsimnibs)
            f.create_dataset('didt', data=np.array(self.didt, dtype=float))
            f.create_dataset('fnamecoil', data=np.array(self.fnamecoil, dtype=np.string_))

    @classmethod
    def load_hdf5(cls, fn_hdf5):
        s = gPCSampler.load_hdf5(fn_hdf5)
        with h5py.File(fn_hdf5, 'r') as f:
            matsimnibs = f['matsimnibs'][()]
            fnamecoil = f['fnamecoil'][()].decode()
            didt = f['didt'][()]

        return cls(
            s.mesh, s.poslist, s.fn_hdf5, fnamecoil,
            matsimnibs, didt, roi=s.roi)

    def run_simulation(self, random_vars):
        poslist = self._update_poslist(random_vars)
        cond = poslist.cond2elmdata(self.mesh)
        if self.constant_dAdt:
            try:
                dAdt = self.dAdt
                dAdt_roi = self.dAdt_roi
            except AttributeError:
                tms_coil = TmsCoil.from_file(self.fnamecoil)
                didt = np.atleast_1d(self.didt)
                if len(didt) == 1:
                    for stimulator in tms_coil.get_elements_grouped_by_stimulators().keys():
                        stimulator.di_dt = didt
                else:
                    for stimulator, stimulator_didt in zip(tms_coil.get_elements_grouped_by_stimulators().keys(), didt):
                        stimulator.di_dt = stimulator_didt
                dAdt = tms_coil.get_da_dt(self.mesh, self.matsimnibs)
                if isinstance(dAdt, mesh_io.NodeData):
                    dAdt = dAdt.node_data2elm_data()
                dAdt.field_name = 'dAdt'
                dAdt.write_hdf5(self.fn_hdf5, 'mesh/elmdata/')
                self.dAdt = dAdt
                self.mesh.elmdata = [dAdt]
                cropped = self.mesh.crop_mesh(self.roi)
                dAdt_roi = cropped.elmdata[0]
                dAdt_roi.write_hdf5(self.fn_hdf5, 'mesh_roi/elmdata/')
                self.mesh.elmdata = []
                self.dAdt_roi = dAdt_roi
        else:
            raise NotImplementedError

        v = fem.tms_dadt(self.mesh, cond, dAdt)
        self.mesh.nodedata = [v]
        cropped = self.mesh.crop_mesh(self.roi)
        v_c = cropped.nodedata[0]
        self.mesh.nodedata = []

        qois = []
        for qoi_name, qoi_f in self.qoi_function.items():
            qois.append(qoi_f(v_c, random_vars, dAdt_roi))

        self.record_data_matrix(random_vars, 'random_var_samples', '/')
        self.record_data_matrix(v.value, 'v_samples', 'mesh/data_matrices')
        self.record_data_matrix(v_c.value, 'v_samples',
                                'mesh_roi/data_matrices')
        for qoi_name, qoi_v in zip(self.qoi_function.keys(), qois):
            self.record_data_matrix(
                qoi_v, qoi_name + '_samples', 'mesh_roi/data_matrices')

        del cropped
        del cond
        del v

        return np.atleast_1d(qois[0]).reshape(-1)


import glob
import asyncio
from dir_settings import DATA_PATH
from asgiref.sync import sync_to_async
import pandas as pd
import spotipy
from spotipy.oauth2 import SpotifyClientCredentials

class Uris:
    '''Handling Uris'''

    spotify_key = None
    spotify_secret = None

    def __init__(self):
        pass

    def get_uris(self, spotify_key, spotify_secret):
        if not spotify_key:
            raise Exception('Api key not provided')

        if not spotify_secret:
            raise Exception('Api secret not provided')

        self.spotify_key = str(spotify_key)
        self.spotify_secret = str(spotify_secret)

        self._run_until_patcher_finish()
        self._concat_uris()

    def _run_until_patcher_finish(self):
        part = 0
        patches = self._create_patches()
        while part < len(patches):
            part = self._uris_patcher(part=part)

    def _uris_patcher(self, part=0):
        patches = self._create_patches()
        try:
            while part < len(patches):
                loop = asyncio.get_event_loop()
                future = asyncio.ensure_future(self._main(patches[part]))
                responses = loop.run_until_complete(future)
                self._save_uris(responses, part)
                part += 1
            return part
        finally:
            return part

    def _concat_uris(self):
        path = DATA_PATH + '/uri_chunks'
        all_files = glob.glob(path + "/*.csv")
        li = []

        for filename in all_files:
            df = pd.read_csv(filename, index_col=None, header=0)
            li.append(df)

        frame = pd.concat(li, axis=0, ignore_index=True)
        frame.to_csv(DATA_PATH + '/uris.csv', index=None, encoding='utf-8')

    def _sp_init(self):
        client_credentials_manager = SpotifyClientCredentials(client_id=self.spotify_key,
                                                              client_secret=self.spotify_secret)
        return spotipy.Spotify(client_credentials_manager=client_credentials_manager)

    def _sp_search(self, q):
        sp = self._sp_init()
        print(q)
        return sp.search(q=q, type='track', limit=1)

    async def _main(self, patch):
        tasks = [asyncio.ensure_future(sync_to_async(self._sp_search)(q)) for q in patch]
        return await asyncio.gather(*tasks)

    def _get_data_from_file(self):
        df = pd.read_csv(DATA_PATH + '/scrobbles.csv')
        artist_names = df['artist'].values.tolist()
        track_names = df['track'].values.tolist()
        return artist_names, track_names

    def _get_queries(self):
        artist_names, track_names = self._get_data_from_file()
        return ['artist: ' + artist_names[i] + ' track: ' + track_names[i] for i in range(len(artist_names))]

    def _create_patches(self):
        queries = self._get_queries()
        return [queries[x:x + 1000] for x in range(0, len(queries), 1000)]

    def _save_uris(self, responses, c):
        uris = []

        for response in responses:
            print(response)
            resp_data = response['tracks']['items']
            if not resp_data:
                continue
            uri = response['tracks']['items'][0]['uri']
            uris.append(uri)

        df = pd.DataFrame()
        df['uri'] = uris
        df.to_csv(DATA_PATH + '/uri_chunks/patch' + str(c) + '.csv', index=None, encoding='utf-8')


import os
import numpy as np
import pyqtgraph as pg
from PyQt5.QtCore import QThread, QObject, pyqtSignal, QRunnable, QThreadPool
from PyQt5.QtGui import QImage
from PyQt5.QtWidgets import QAction, QSpinBox, QLabel, QDoubleSpinBox, QMenu, QTableWidget, QPushButton, QWidget, \
    QVBoxLayout, QHBoxLayout
from mne.viz import plot_topomap

from scipy.signal import savgol_filter
from scipy.fft import fft, fftfreq, rfft
from PyQt5 import QtCore, QtGui, QtWidgets
from datetime import datetime
import mne
from scipy.fftpack import rfftfreq

from pynfb.generators import ch_names32
from pynfb.generators import ch_names27
from pynfb.inlets.montage import Montage

from pynfb.postprocessing.utils import runica
from pynfb.widgets import edf_saver
from vispy import scene

class WorkerSignals(QObject):
    finished = pyqtSignal()
    error = pyqtSignal(tuple)
    result = pyqtSignal(int, object, object)
    progress = pyqtSignal(int)

class DPHandler(QRunnable):
    def __init__(self, y, c, scaler, shift):
        super().__init__()
        self.y = y
        self.c = c
        self.scaler = scaler
        self.shift = shift
        self.signals = WorkerSignals()

    def run(self):
        signal_buffer = (self.y - self.c) * self.scaler + 50 * self.shift
        spectral_buffer = rfft(self.y)
        self.signals.result.emit(self.shift, signal_buffer, spectral_buffer)

class CrossButtonsWidget(QtWidgets.QWidget):
    names = ['left', 'right', 'up', 'down']
    symbols = ['<', '>', '^', 'v']
    positions = [(1, 0), (1, 2), (0, 1), (2, 1)]

    def __init__(self, parent=None):
        super(CrossButtonsWidget, self).__init__(parent)
        buttons = dict([(name, QtWidgets.QPushButton(symbol)) for name, symbol in zip(self.names, self.symbols)])
        layout = QtWidgets.QGridLayout(self)
        self.clicked_dict = {}
        for name, pos in zip(self.names, self.positions):
            buttons[name].setAutoRepeat(True)
            buttons[name].setAutoRepeatDelay(100)
            layout.addWidget(buttons[name], *pos)
            self.clicked_dict[name] = buttons[name].clicked

class LSLPlotDataItem(pg.PlotDataItem):
    def getData(self):
        x, y = super(LSLPlotDataItem, self).getData()
        if self.opts['fftMode']:
            return x, y / (max(y) - min(y) + 1e-20) * 0.75
        return x, y

class PopupWidget(QPushButton):
    def __init__(self, label, elements, checklist, parent=None):
        super(PopupWidget, self).__init__(parent)
        self.parent = parent
        self.setText(label)
        self.menu = QMenu(self)
        self.checklist = checklist
        for i in range(0, len(elements)):
            Action = QAction(elements[i], self.menu)
            Action.triggered.connect(lambda chk, item=i: self.open_spectral(item))
            self.menu.addAction(Action)

    def mouseReleaseEvent(self, event):
        self.menu.exec_(event.globalPos())

    def open_spectral(self, index):
        self.parent.open_spectral(index)

    def getChecklist(self):
        mlen = len(self.menu.actions())
        checklist = np.zeros(mlen)
        for i in range(0, mlen):
            checklist[i] = self.menu.actions()[i].isChecked()
        return checklist

class SpectralWindow(QWidget):
    def __init__(self, name, n, freq):
        super().__init__()
        layout = QHBoxLayout()
        self.setWindowTitle(name)

        color = np.ones((n, 4), dtype=np.float32)
        color[:, 0] = np.linspace(0, 1, n)
        color[:, 1] = color[::-1, 0]
        canvas = scene.SceneCanvas(keys='interactive', show=True)
        grid = canvas.central_widget.add_grid(spacing=0)
        self.viewbox = grid.add_view(row=0, col=1, camera='panzoom')
        x_axis = scene.AxisWidget(orientation='bottom')
        x_axis.stretch = (1, 0.1)
        grid.add_widget(x_axis, row=1, col=1)
        x_axis.link_view(self.viewbox)
        y_axis = scene.AxisWidget(orientation='left')
        y_axis.stretch = (0.1, 1)
        grid.add_widget(y_axis, row=0, col=0)
        y_axis.link_view(self.viewbox)

        self.pos = np.zeros((n, 2), dtype=np.float32)
        self.pos[:, 0] = rfftfreq(n, 1/freq)
        self.line = scene.Line(self.pos, color, parent=self.viewbox.scene)

        self.viewbox.camera.set_range()
        layout.addWidget(canvas.native)
        self.setLayout(layout)

    def set_data(self, y):
        self.pos[:, 1] = y
        self.line.set_data(pos=self.pos)

class RhythmsWindow(QWidget):
    def __init__(self, n, freq):
        super().__init__()
        layout = QHBoxLayout()
        self.setWindowTitle("Rhythms indicators")
        self.n = n
        self.freq = freq

        self.bar1 = pg.BarGraphItem(x=[1], height=1, width=0.2, brush='g')
        self.bar2 = pg.BarGraphItem(x=[2], height=2, width=0.2, brush='w')
        self.bar3 = pg.BarGraphItem(x=[3], height=3, width=0.2, brush='y')
        self.plot = pg.PlotWidget()
        self.plot.addItem(self.bar1)
        self.plot.addItem(self.bar2)
        self.plot.addItem(self.bar3)
        layout.addWidget(self.plot)
        self.setLayout(layout)

        self.update_counter = 1001

    def get_indicator(self, data, num, f, k2):
        x = rfftfreq(self.n, 1/self.freq)
        buffer = rfft(data[:, num])
        sy0 = 0
        sy1 = 0
        sy2 = 0
        sx = 0
        B = []

        for i in range(0, len(x), 1):
            if f[0] < x[i] < f[3]:
                B.append(np.abs(buffer)[i])

        k1 = len(B)-1 if len(B)%2 == 0 else len(B)
        fb = savgol_filter(B, k1, k2)
        #pg.plot(fb)
        kx = len(fb)/(f[3]-f[0])
        for i in range(0, int(kx*(f[1]-f[0])), 1):
            sy0 += fb[i]
        for i in range(int(kx*(f[1]-f[0])), int(kx*(f[2]-f[0])), 1):
            sy1 += fb[i]
        for i in range(int(kx*(f[2]-f[0])), len(fb), 1):
            sy2 += fb[i]

        # print(kx)
        # print(sy1)
        # print(sy0)
        # print(sy2)

        if sy1:
            sx = 1 - (sy0/sy1+sy2/sy1)/2
        if sx < -1:
            sx = -1
        return sx

    def set_data(self, data):
        self.update_counter += 1
        if self.update_counter <= 10:
            return
        else:
            self.update_counter = 0

        ai = self.get_indicator(data[0], 19, [7, 9, 11, 13], 6) # OZ, alpha
        ti = self.get_indicator(data[1], 3, [3, 5, 7, 9], 6) # F7, theta
        di = self.get_indicator(data[2], 10, [0.4, 1.4, 2.4, 3.4], 7) # CZ, delta
        self.bar1.setOpts(height=ai)
        self.bar2.setOpts(height=ti)
        self.bar3.setOpts(height=di)

class RawViewer(QWidget):
    def __init__(self, fs, channels_labels, parent=None, overlap=False, mode=0, toolbar=None, freqbar=None, widgets=None):
        super(RawViewer, self).__init__(parent)
        n_channels = len(channels_labels)
        #print(channels_labels)
        self.n_channels = n_channels
        self.calib_values = np.zeros(self.n_channels)
        self.calib_flag = False
        self.visible_flags = np.zeros(self.n_channels)
        self.channelsanalyze_flags = np.zeros(self.n_channels)
        self.spectraldata = []
        self.tw = None
        self.sw = [QWidget]*n_channels
        for i in range(0, n_channels):
            #self.sw.append(None)
            self.visible_flags[i] = True
            self.channelsanalyze_flags[i] = False

        print("Signal painter started")

        self.channels_labels = channels_labels
        self.toolbar = toolbar
        #self.freqbar = freqbar
        self.widgets = widgets
        self.fs = fs
        self.current_pos = 0
        self.std = None
        self.raw_buffer = None
        self.filtered_buffer1 = None
        self.filtered_buffer2 = None
        self.filtered_buffer3 = None
        self.scaler = 0.0041
        self.step = 0.0001
        self.curves = []
        self.lines = []
        self.x_mesh = None
        #self.setYRange(0, n_channels)
        self.mode = mode
        self.n_samples_to_display = 0
        #self.notch_filter = NotchFilter(0, fs, n_channels)
        #self.butter_filter = ButterFilter((0.5, 45), fs, n_channels)

        print("FS: "+str(fs))
        self.pool = QThreadPool.globalInstance()

        if self.mode == 0:
            buffer_time_sec = 10
            self.n_samples = int(buffer_time_sec*fs)
            self.n_samples_to_display = self.n_samples
            self.raw_buffer = np.zeros((self.n_samples, self.n_channels))
            self.x_mesh = np.linspace(0, self.n_samples_to_display / self.fs, self.n_samples_to_display)
            #self.setXRange(0, self.x_mesh[self.n_samples_to_display - 1])
            self.create_surface()
            self.rw = RhythmsWindow(int(self.n_samples/2), self.fs)

        self.saveAction = QAction(self)
        self.saveAction.setText("Save data...")
        self.recalibAction = QAction(self)
        self.recalibAction.setText("Recalibrate")
        self.scaleupAction = QAction(self)
        self.scaleupAction.setText(" ^ ")
        self.scaledownAction = QAction(self)
        self.scaledownAction.setText(" v ")
        self.saveAction.triggered.connect(lambda: self.save_data())
        self.recalibAction.triggered.connect(lambda: self.calibrate())
        self.scaleupAction.triggered.connect(lambda: self.update_scaler(increase=True))
        self.scaledownAction.triggered.connect(lambda: self.update_scaler(increase=False))
        self.toolbar.addAction(self.saveAction)
        self.toolbar.addAction(self.recalibAction)
        self.toolbar.addAction(self.scaleupAction)
        self.toolbar.addAction(self.scaledownAction)
        self.stepSpinBox = QDoubleSpinBox()
        self.ampLabel = QLabel()
        #font = self.stepSpinBox.font()
        #font.setPointSizeF(20.25)
        #self.stepSpinBox.setFont(font)

        #self.stepSpinBox.setFloatDomain(True)
        self.stepSpinBox.setDecimals(5)
        self.stepSpinBox.setMaximum(1000)
        self.stepSpinBox.setValue(self.step)
        self.stepSpinBox.setFixedWidth(100)
        self.ampLabel.setText("Amplitude scale: x"+str(self.scaler))
        self.ampLabel.setFixedWidth(150)
        self.toolbar.addWidget(self.ampLabel)
        self.toolbar.addWidget(self.stepSpinBox)

        # self.visMenu = QMenu("menu")
        # self.visMenu.addAction(QAction('50%', self.visMenu, checkable=True))
        # self.visMenu.addAction(QAction('50%', self.visMenu, checkable=True))
        # self.visMenu.addAction(QAction('50%', self.visMenu, checkable=True))
        # self.visMenu.addAction(QAction('50%', self.visMenu, checkable=True))

        #self.vis_widget = PopupWidget('Channel visibility', channels_labels, self.visible_flags, parent=self)
        #self.toolbar.addWidget(self.vis_widget)

        self.model = QtGui.QStandardItemModel()
        self.listView = QtWidgets.QListView()
        self.alglist = QtWidgets.QListView()
        self.algmodel = QtGui.QStandardItemModel()

        if channels_labels is not None:
            for i in range(len(channels_labels)):
                item = QtGui.QStandardItem(channels_labels[i])
                item.setCheckable(True)
                #check = QtCore.Qt.Checked if checked else QtCore.Qt.Unchecked
                item.setCheckState(QtCore.Qt.Checked)

                self.model.appendRow(item)
        self.listView.setModel(self.model)
        algnames = []
        # algnames = ["[preprocessing] Butterworth filter", "[analyze] SVM"]
        for i in range(len(algnames)):
            item = QtGui.QStandardItem(algnames[i])
            item.setCheckable(True)
            item.setCheckState(QtCore.Qt.Checked)
            self.algmodel.appendRow(item)

        self.alglist.setModel(self.algmodel)

        #vbox.addWidget(self.listView)
        #vbox.addStretch(1)
        self.listView.clicked.connect(self.process_vischeck)

        self.listwidget = QWidget()
        self.listlayout = QVBoxLayout(self)
        self.listlayout.addWidget(self.listView)
        self.listlayout.addWidget(self.alglist)
        ab1 = QPushButton()
        ab1.setText("Add")
        ab2 = QPushButton()
        ab2.setText("Edit")
        ab3 = QPushButton()
        ab3.setText("Delete")
        self.listlayout.addWidget(ab1)
        self.listlayout.addWidget(ab2)
        self.listlayout.addWidget(ab3)
        self.listwidget.setLayout(self.listlayout)
        self.widgets.setWidget(self.listwidget)

        self.channelsfilter_widget = PopupWidget('Spectral plots', channels_labels, self.channelsanalyze_flags, parent=self)
        self.toolbar.addWidget(self.channelsfilter_widget)

        self.rwopenAction = QAction(self)
        self.rwopenAction.setText("Rhythms")
        self.rwopenAction.triggered.connect(lambda: self.rw.show())
        self.toolbar.addAction(self.rwopenAction)

        self.twopenAction = QAction(self)
        self.twopenAction.setText("Topography")
        self.twopenAction.triggered.connect(lambda: self.tw.showMaximized())
        self.toolbar.addAction(self.twopenAction)

        #self.setLayout(vbox)

        #self.channelsfilter_widget = PopupWidget('Channels to filter', channels_labels, self.channelsfilter_flags, parent=self)
        #self.toolbar.addWidget(self.channelsfilter_widget)

        self.show_levels_flag = False
        self.overlap = overlap
        if overlap:
            self.show_levels_flag = True
            self.show_levels()

    def create_surface(self):
        N = self.n_samples_to_display
        color = np.ones((N, 4), dtype=np.float32)
        color[:, 0] = np.linspace(0, 1, N)
        color[:, 1] = color[::-1, 0]
        canvas = scene.SceneCanvas(keys='interactive', show=True)
        grid = canvas.central_widget.add_grid(spacing=0)
        self.viewbox = grid.add_view(row=0, col=1, camera='panzoom')

        x_axis = scene.AxisWidget(orientation='bottom')
        x_axis.stretch = (1, 0.1)
        grid.add_widget(x_axis, row=1, col=1)
        x_axis.link_view(self.viewbox)
        y_axis = scene.AxisWidget(orientation='left')
        y_axis.stretch = (0.1, 1)
        grid.add_widget(y_axis, row=0, col=0)
        y_axis.link_view(self.viewbox)
        self.setLayout(QVBoxLayout())
        self.layout().addWidget(canvas.native)
        for i in range(0, self.n_channels):
            pos = np.zeros((N, 2), dtype=np.float32)
            pos[:, 0] = self.x_mesh[:self.n_samples_to_display]
            if i:
                pos[:, 1] = pos[:, 1] + 50*i
            c = scene.Line(pos, color, parent=self.viewbox.scene)
            self.curves.append(pos)
            self.lines.append(c)
        self.viewbox.camera.set_range()

    def open_spectral(self, index):
        if not self.channelsanalyze_flags[index]:
            self.sw[index] = SpectralWindow(self.channels_labels[index], int(self.n_samples / 2), self.fs)
            self.channelsanalyze_flags[index] = True
        self.sw[index].set_data(np.abs(self.spectraldata)[:int(self.n_samples/2)])
        self.sw[index].show()

    def process_vischeck(self, index):
        self.visible_flags[index.row()] = self.listView.model().itemData(index)[10]
        if not self.visible_flags[index.row()]:
            self.curves[index.row()][:, 1] = np.ones(self.n_samples_to_display) + 50 * index.row()

        self.update_data()

    def update_std(self, chunk):
        if self.std is None:
            self.std = np.std(chunk)
        else:
            self.std = 0.5 * np.std(chunk) + 0.5 * self.std

    def set_chunk(self, signal, chunk):
        n_samples = len(chunk)

        if self.mode == 1:
            self.raw_buffer = np.zeros((n_samples, self.n_channels))
            self.n_samples_to_display += n_samples
            self.x_mesh = np.linspace(0, self.n_samples_to_display / self.fs, self.n_samples_to_display)
            self.n_samples = n_samples
            self.create_surface()
            # for i in range(0, self.n_channels):
            #     self.sw.append(SpectralWindow(self.channels_labels[i], int(self.n_samples/2), self.fs))
            self.rw = RhythmsWindow(int(self.n_samples/2), self.fs)

        self.raw_buffer[:-n_samples] = self.raw_buffer[n_samples:]
        self.raw_buffer[-n_samples:] = chunk
        if not self.calib_flag:
            self.calibrate()

        if signal is not None:
            # runica(self.raw_buffer, self.fs, self.channels_labels)
            #montage = Montage(ch_names32)
            # ch_names32 = ['EEG FP1', 'EEG FPZ', 'EEG FP2', 'EEG  F7', 'EEG  F3', 'EEG  FZ', 'EEG  F4', 'EEG  F8', 'EEG  T3',
            #               'EEG  C3', 'EEG  CZ', 'EEG C4', 'EEG T4', 'EEG  T5', 'EEG P3', 'EEG PZ', 'EEG P4', 'EEG T6', 'EEG O1',
            #               'EEG OZ', 'EEG O2', 'EEG PBH', 'EEG NOS', 'EEG SHL', 'EEG SHH', 'EEG A1', 'EEG A2', '', '', '', '', '']
            #montage = Montage(ch_names32)
            #print(self.channels_labels)
            #a = ch_names27.conjugate
            #print(np.shape(self.channels_labels))
            #print(np.shape(ch_names27))
            #b = self.channels_labels.conjugate

            #print(montage.get_names())
            # spatial, topo = runica(np.random.normal(size=(100000, 32)), 1000, montage.get_names(), mode='csp')
            if False:
                self.tw = runica(self.raw_buffer, self.fs, self.channels_labels, mode='csp')
                topocanvas = self.tw.table.topographies_items[0].topo
                width, height = topocanvas.get_width_height()
                print("canvas pixels size "+str(width)+"x"+str(height))
                for x in range(0, width):
                    for y in range(0, height):
                        pixel = QImage(topocanvas.buffer_rgba(), width, height, QImage.Format_ARGB32).pixel(x, y)
                        #print(QtGui.QColor(pixel).getRgb()[:-1])
                #print("topography data: ", self.tw.topographies)
                # plot_topomap(spatial, montage.get_pos())
                # plot_topomap(topo, montage.get_pos())

            self.filtered_buffer1 = self.raw_buffer.copy()
            self.filtered_buffer2 = self.raw_buffer.copy()
            self.filtered_buffer3 = self.raw_buffer.copy()
            signal.fit_model(self.filtered_buffer1, [8, 13], self.channels_labels)
            self.filtered_buffer1 = signal.apply(self.filtered_buffer1)

            signal.fit_model(self.filtered_buffer2, [4, 8], self.channels_labels)
            self.filtered_buffer2 = signal.apply(self.filtered_buffer2)

            signal.fit_model(self.filtered_buffer3, [1, 4], self.channels_labels)
            self.filtered_buffer3 = signal.apply(self.filtered_buffer3)
        self.update_data()

    def set_data(self, i, data, spectraldata):
        if self.visible_flags[i]:
            self.curves[i][:, 1] = data
            self.lines[i].set_data(pos=self.curves[i])
        self.spectraldata = spectraldata
        if self.mode == 1 and self.channelsanalyze_flags[i]:
            self.sw[i].set_data(np.abs(spectraldata)[:int(self.n_samples/2)])

    def update_data(self):
        for i in range(0, self.n_channels, 1):
            runnable = DPHandler(self.filtered_buffer3[:, i], self.calib_values[i], self.scaler, i)
            runnable.signals.result.connect(self.set_data)
            self.pool.start(runnable)
        self.rw.set_data([self.filtered_buffer1, self.filtered_buffer2, self.filtered_buffer3])
        self.ampLabel.setText("Amplitude scale: x" + str(self.scaler))

    def update_scaler(self, increase=False):
        self.step = self.stepSpinBox.value()
        self.scaler += self.step if increase else -self.step
        if self.scaler < 0:
            self.scaler = 0
        self.update_data()
        self.update_levels()

    def update_n_samples_to_display(self, increase=False):
        step = int(self.fs * 0.5)
        self.n_samples_to_display += step if increase else -step
        if self.n_samples_to_display < 10:
            self.n_samples_to_display = 10
        elif self.n_samples_to_display > self.n_samples:
            self.n_samples_to_display = self.n_samples
        self.update_data()
        self.update_levels()

    def show_levels(self):
        self.levels = {'zero': [], 'p1': [], 'm1': []}
        for i in (range(self.n_channels) if not self.overlap else [0]):
            for level, val in zip(['zero', 'p1', 'm1'], [0, 1, -1]):
                c = LSLPlotDataItem()
                c.setPen(pg.mkPen('g'))
                #self.addItem(c)
                if not self.overlap:
                    c.setPos(0, i + 1)
                #self.levels[level].append(c)
        self.update_levels()

    def update_levels(self):
        if self.show_levels_flag:
            for i in (range(self.n_channels) if not self.overlap else [0]):
                for level, val in zip(['zero', 'p1', 'm1'], [0, 1, -1]):
                    self.levels[level][i].setData(self.x_mesh[:self.n_samples_to_display],
                              self.x_mesh[:self.n_samples_to_display]*0 + val * self.scaler / (self.std or 1))

    def save_data(self):
        info = mne.create_info(ch_names=self.channels_labels,
                               ch_types=['misc'] * len(self.channels_labels),
                               sfreq=self.fs)
        raw_buffer_t = np.zeros((len(self.channels_labels), self.n_samples_to_display))
        for i in range(0, self.raw_buffer.shape[1]):
            raw_buffer_t[i, :] = self.raw_buffer[:, i]
        raw = mne.io.RawArray(raw_buffer_t, info)
        now = datetime.now()
        dt_string = now.strftime("%Y%m%d-%H_%M_%S")
        print("Timestamp: ", dt_string)
        edf_saver.write_mne_edf(raw, os.path.realpath(os.path.dirname(os.path.realpath(__file__)))+"\\..\\..\\saved_eeg\\"+dt_string+".edf", overwrite=True)

    def calibrate(self):
        valnum = len(self.raw_buffer)
        for i in range(0, self.n_channels):
            av = 0
            for j in range(0, valnum):
                av += self.raw_buffer[j][i]
            self.calib_values[i] = av / valnum
        self.calib_flag = True
        self.scaler = 9e-1 / abs(self.raw_buffer[-1, 0]) if self.raw_buffer[-1, 0] else 1
        self.step = self.scaler / 10
        self.stepSpinBox.setValue(self.step)
        print("Calibration data:")
        print(self.calib_values)
        #print(self.scaler)

if __name__ == '__main__':
    a = QtWidgets.QApplication([])
    plot_widget = RawViewer(250, ['ef', 'sf', 'qwr']*3)

    plot_widget.set_chunk(np.sin(np.arange(20)).reshape(20, 1).dot(np.ones((1, 9))))
    plot_widget.show()
    a.exec_()


from messy.lib import roles as r
from messy.lib import plate_utils

from dateutil import parser
import uuid

def setup(dbh):

    dbh.EK.bulk_update(ek_initlist, dbsession=dbh.session())
    dbh.Group.bulk_insert(messy_groups, dbsession=dbh.session())

    # add controls

    dbh.session().add(
        dbh.Institution(
            code='NOT-AVAILABLE',
            name='Not Available'
        )
    )
    dbh.session().flush()

    dbh.session().add(
        dbh.Collection(
            code='CONTROL',
            uuid=uuid.uuid4(),
            group=dbh.get_group('CollectionMgr'),
            institutions=[dbh.get_institutions_by_codes('NOT-AVAILABLE', None)[0]],
        )
    )
    dbh.session().flush()

    plates = [
        dbh.Plate(
            code='TEMPLATE-96',
            group=dbh.get_group('PlateMgr'),
            date=parser.parse('1970'),
            specimen_type='water',
            experiment_type='sample-container',
            user=dbh.get_user('system/_SYSTEM_'),
        ),
        dbh.Plate(
            code='TEMPLATE-384',
            group=dbh.get_group('PlateMgr'),
            date=parser.parse('1970'),
            specimen_type='water',
            experiment_type='sample-container',
            user=dbh.get_user('system/_SYSTEM_'),
        )
    ]
    for p in plates:
        dbh.session().add(p)
    dbh.session().flush(plates)

    d_sample = dict(
        collection=dbh.get_collections_by_codes('CONTROL', None, ignore_acl=True)[0],
        received_date=parser.parse('1970'),
        collection_date=parser.parse('1970'),
        passage='original',
        species='no-species',
        host='no-species',
        host_occupation='other',
        host_status='unknown',
        category='Z-CR',
        specimen_type='no-specimen',
        ct_method='no-ct',
        originating_institution=dbh.get_institutions_by_codes('NOT-AVAILABLE', None)[0],
        sampling_institution=dbh.get_institutions_by_codes('NOT-AVAILABLE', None)[0],
    )

    samples = [
        dbh.Sample().update(d_sample | dict(code='-')),
        dbh.Sample().update(d_sample | dict(code='NTC1')),
        dbh.Sample().update(d_sample | dict(code='NTC2')),
        dbh.Sample().update(d_sample | dict(code='NTC3')),
        dbh.Sample().update(d_sample | dict(code='NTC4')),
        dbh.Sample().update(d_sample | dict(code='*')),
    ]
    for s in samples:
        dbh.session().add(s)
    dbh.session().flush(samples)

    plate_utils.create_positions(plates[0], 96)
    plate_utils.create_positions(plates[1], 384)
    dbh.session().flush()

# add additional initial data here

messy_groups = [
    ('InstitutionMgr', [r.INSTITUTION_MANAGE]),
    ('CollectionMgr', [r.COLLECTION_MODIFY, r.COLLECTION_VIEW]),
    ('CollectionViewer', [r.COLLECTION_VIEW]),
    ('SampleMgr', [r.SAMPLE_MANAGE]),
    ('SampleModifier', [r.SAMPLE_MODIFY, r.SAMPLE_VIEW]),
    ('SampleViewer', [r.SAMPLE_VIEW]),
    ('PlateMgr', [r.PLATE_MANAGE]),
    ('PlateModifier', [r.PLATE_MODIFY, r.PLATE_VIEW]),
    ('PlateViewer', [r.PLATE_VIEW]),
    ('SequencingRunMgr', [r.SEQUENCINGRUN_MANAGE]),
    ('SequencingRunModifier', [r.SEQUENCINGRUN_MODIFY, r.SEQUENCINGRUN_VIEW]),
    ('SequencingRunViewer', [r.SEQUENCINGRUN_VIEW]),
    ('SequenceMgr', [r.SEQUENCE_MODIFY, r.SEQUENCE_VIEW]),
    ('SequenceViewer', [r.SEQUENCE_VIEW]),
    ('ProjectMgr', [r.SAMPLE_MODIFY, r.PLATE_MODIFY, r.SEQUENCINGRUN_MODIFY, r.SEQUENCE_MODIFY]),
    ('ProjectViewer', [r.SAMPLE_VIEW, r.PLATE_VIEW, r.SEQUENCINGRUN_VIEW, r.SEQUENCE_VIEW]),
    ('Collaborator', [r.SAMPLE_MODIFY, r.SEQUENCE_VIEW]),
]

ek_initlist = [
    ('@SYSNAME', 'System names',
        [
            ('messy'.upper(), 'messy'),
        ]
     ),
    ('@ROLES', None,
        [
            (r.INSTITUTION_MANAGE, 'manage institution'),
            (r.INSTITUTION_MODIFY, 'modify institution'),
            (r.COLLECTION_MANAGE, 'manage collection'),
            (r.COLLECTION_MODIFY, 'modify collection'),
            (r.COLLECTION_VIEW, 'view collection'),
            (r.SAMPLE_MANAGE, 'manage sample'),
            (r.SAMPLE_MODIFY, 'modify sample'),
            (r.SAMPLE_VIEW, 'view sample'),
            (r.PLATE_MANAGE, 'plate manager'),
            (r.PLATE_MODIFY, 'modify plate'),
            (r.PLATE_VIEW, 'view plate'),
            (r.SEQUENCINGRUN_MANAGE, 'manage sequencing run'),
            (r.SEQUENCINGRUN_MODIFY, 'modify sequencing run'),
            (r.SEQUENCINGRUN_VIEW, 'view sequencing run'),
            (r.SEQUENCE_MANAGE, 'manage sequence'),
            (r.SEQUENCE_MODIFY, 'modify sequence'),
            (r.SEQUENCE_VIEW, 'view sequence'),
        ]
     ),
    ('@SPECIES', "Species",
        [
            ('betacoronavirus', 'betacoronavirus'),
            ('betacoronavirus-ncov19', 'betacoronavirus nCoV-19'),
            ('human', 'human'),
            ('no-species', 'no-species'),
        ]
     ),
    ('@PASSAGE', 'Passage',
        [
            ('original', 'Original'),
            ('vero', 'Vero cell 1st passage'),
            ('vero+2', 'Vero cell 2nd passage'),
            ('hek293', 'HEK293 cell 1st passage'),
            ('hek293+2', 'HEK293 cell 2nd passage')
        ]
     ),
    ('@HOST_STATUS', 'Host status',
        [
            ('unknown', 'Unknown'),
            ('hospitalized', 'Hospitalized'),
            ('live', 'Live'),
            ('deceased', 'Deceased'),
            ('released', 'Released')
        ]
     ),
    ('@HOST_OCCUPATION', 'Host occupation',
        [
            ('irrelevant', 'Irrelevant'),
            ('other', 'Other occupation'),
            ('unemployed', 'Unemployed'),
            ('health', 'Health sector'),
            ('government', 'Government sector'),
            ('media', 'Media sector'),
            ('telecommunication', 'Telecommunication sector'),
            ('finance', 'Financial sector'),
            ('it', 'Information Technology sector'),
            ('teacher', 'Teacher'),
            ('lecturer', 'Lecturer'),
            ('education', 'Education sector'),
            ('student', 'Student'),
            ('retail', 'Retail sector'),
            ('domestic', 'Domestic (family) sector'),
            ('transport', 'Transportation sector'),
            ('entertainment', 'Entertainment sector'),
            ('agriculture', 'Agricultural sector'),
            ('fishery', 'Fishery sector'),
            ('mining', 'Mining sector'),
            ('foresty', 'Forestry sector'),
        ]
     ),
    ('@CATEGORY', 'Sample category',
        [
            ('R-RA', 'R - Random surveillance and tracking'),
            ('S-SE', 'S - Sentinel surveillance'),
            ('A-NE', 'A - COVID19 clinically (-) test'),
            ('B-CL', 'B - COVID19 severe clinically'),
            ('C-TR', 'C - travel history'),
            ('D-LO', 'D - long COVID19'),
            ('E-RE', 'E - reinfection'),
            ('F-PO', 'F - post-vaccination'),
            ('G-CH', 'G - child case'),
            ('H-CO', 'H - comorbid case'),
            ('Z-CR', 'Z - Control or reference sample'),
        ]
     ),
    ('@SPECIMEN_TYPE', 'Specimen type',
        [
            ('np+op', 'Nasopharyngeal and oropharyngeal swab'),
            ('np', 'Nasopharyngeal swab'),
            ('op', 'Oropharyngeal swab swab'),
            ('sputum', 'Sputum'),
            ('blood', 'Blood'),
            ('serum', 'Serum'),
            ('saliva', 'Saliva'),
            ('stool', 'Stool'),
            ('wastewater', 'Wastewater'),
            ('rna', 'RNA'),
            ('cdna', 'cDNA'),
            ('dna', 'double-strand DNA'),
            ('ssdna', 'single-strand DNA'),
            ('water', 'empty sample'),
            ('no-specimen', 'no specimen'),
        ]
     ),
    ('@CT_METHOD', 'Ct Methodology',
        [
            ('rtpcr', 'generic realtime RT-PCR'),
            ('cobas', 'Cobas'),
            ('charite-berlin', 'Charite-Berlin method'),
            ('lamp', 'LAMP'),
            ('no-ct', 'No Ct'),
        ]
     ),
    ('@EXPERIMENT_TYPE', 'Laboratory experiment type',
        [
            ('sample-container', 'container for sample storage'),
            ('rna-extraction', 'RNA extraction'),
            ('qc-qubit', 'QC with qubit'),
            ('qc-accuclear', 'QC with AccuClear'),
            ('cdna-ssiv', 'cDNA with Superscript IV'),
            ('cdna-lunascript', 'cDNA with Lunascript'),
            ('pcr-artic-idt', 'PCR with IDT ARTICv3 kit'),
            ('pcr-artic-neb', 'PCR with NEB balanced ARTICv3 kit'),
            ('pcr-artic-covidseq', 'PCR with CovidSeq ARTIC'),
            ('libprep-truseqstrandedrna', 'library prep with Truseq Stranded RNA'),
            ('libprep-truseqnano', 'library prep with Truseq Nano'),
            ('libprep-nextera', 'library prep with Nextera XT'),
            ('libprep-dnaprep', 'library prep with DNAPrep (NexteraFlex)'),
            ('libprep-covidseq', 'library prep with CovidSeq DNAPrep'),
            ('libprep-ultraii', 'library prep with UltraII'),
            ('libprep-ultraiifs', 'library prep with UltraII FS'),
            ('libprep-artic-ultraii', 'library prep with NEB Artic UltraII kit'),
            ('libprep-artic-ultraiifs', 'library prep with NEB Artic UltraIIFS kit'),
        ]
     ),
    ('@SEQUENCING_KIT', 'Sequencing kit',
        [
            ('merged-data', 'Merged data from more than one run'),
            ('MiSeq-V3-600', 'MiSeq 600-cycle v3'),
            ('MiSeq-V2-500', 'MiSeq 500-cycle v2'),
            ('MiSeq-V2-300', 'MiSeq 300-cycle v2'),
            ('MiSeq-MicroV2-300', 'MiSeq 300-cycle v2 Micro'),
            ('MiSeq-NanoV2-500', 'MiSeq Nano v2 500-cycle'),
            ('MiSeq-NanoV2-300', 'MiSeq Nano v2 300-cycle'),
            ('NextSeq-HiV2.5-300', 'NextSeq 500/550 High Output 300-cycle'),
            ('NextSeq-HiV2.5-150', 'NextSeq 500/550 High Output 150-cycle'),
            ('NextSeq-MidV2.5-300', 'NextSeq 500/550 Mid Output 300-cycle'),
            ('NextSeq-MidV2.5-150', 'NextSeq 500/550 Mid Output 150-cycle'),
            ('NovaSeq-SPv1.5-500', 'NovaSeq 6000 SP v1.5 500-cycle'),
            ('NovaSeq-SPv1.5-300', 'NovaSeq 6000 SP v1.5 300-cycle'),
            ('NovaSeq-SPv1.5-200', 'NovaSeq 6000 SP v1.5 200-cycle'),
            ('NovaSeq-SPv1.5-100', 'NovaSeq 6000 SP v1.5 100-cycle'),
            ('data-container', 'Container for imported data')
        ]
     ),
    ('@METHOD', 'Upstream analysis method',
        [
            ('ncov19-pipeline/artic', 'ncov19-pipeline ARTIC mode'),
            ('ncov19-pipeline/mapping', 'ncov19-pipeline generic mapping (minimap2) mode'),
            ('ncov19-pipeline/assembling', 'ncov19-pipeline generic denovo assembling (SPAdes) mode'),
            ('gisaid', 'direct GISAID import'),
            ('other', 'other analysis method'),
        ]
     ),
    ('@ADAPTERINDEX', 'Adapter-index for sequencing',
        [
            ('covidseq-set-1', 'CovidSeq UDI Set 1 1-96'),
            ('covidseq-set-2', 'CovidSeq UDI Set 2 97-192'),
            ('covidseq-set-3', 'CovidSeq UDI Set 3 193-288'),
            ('covidseq-set-4', 'CovidSeq UDI Set 4 289-384'),
            ('nebnext-udi-set-1', 'NebNext Multiplex Oligo UDI Set 1'),
            ('nebnext-udi-set-2', 'NebNext Multiplex Oligo UDI Set 2'),
            ('nebnext-udi-set-3', 'NebNext Multiplex Oligo UDI Set 3'),
            ('nebnext-udi-set-4', 'NebNext Multiplex Oligo UDI Set 4'),
        ]
     ),
    ('@LINEAGEMETHOD', 'Lineage classification method',
        [
            ('pangolin', 'Pangolin'),
            ('scorpio', 'Scorpio'),
            ('gisaid', 'GISAID'),
            ('nextstrain', 'NextStrain'),
        ]
     ),
]

# EOF


from tkinter import StringVar, Tk, ttk
from tkinter.messagebox import showinfo
import routes
from sqlalchemy.exc import DataError
import time
from datetime import date
import calendar

class App(Tk):
    def __init__(self) -> None:
        super().__init__()

        self.title('Employee Login')
        self.geometry('800x500')

        # Main Label
        self.label = ttk.Label(
            self, text='Enter Employee ID', font='none 30 bold')
        self.label.pack()

        # ID Entry
        self.idEntry = StringVar()
        self.idInput = ttk.Entry(
            self, textvariable=self.idEntry, font='none 30 normal')
        self.idInput.focus()
        self.idInput.pack()

        # Submit to verify user
        self.submit = ttk.Button(self, text='Submit', width=30)
        self.submit.bind('<Button-1>', lambda event: self.Submit())
        self.submit.pack()

        self.bind('<Return>', lambda event: self.Submit())

    # Displaying BUTTONS FOR BREAK
    def clock(self):
        currentDate = date.today()
        showDate = calendar.day_name[currentDate.weekday()]
        t = time.localtime()
        self.currentTime = time.strftime('%I:%M:%S %p', t)
        self.clockTime = time.strftime('%I:%M%p', t)
        self.label['text'] = f' Today is {showDate} \n Please Clock in {self.currentTime}'
        self.label.configure(justify='center')
        self.label.after(1000, self.clock)

    def clockInFunction(self):
        showinfo(title='Confirmation',
                 message=f'You have clocked in at {self.clockTime}')
        routes.updateClockIn(userId, self.clockTime)
        self.destroy()

    def clockOutFunction(self):
        showinfo(title='Confirmation',
                 message=f'You have clocked it at {self.clockTime}')
        routes.updateClockOut(userId, self.clockTime)
        self.destroy()

    def breakInFunction(self):
        showinfo(title='Confirmation',
                 message=f'You have clocked in for break at {self.clockTime}')
        routes.updateBreakIn(userId, self.clockTime)
        self.destroy()

    def breakOutFunction(self):
        showinfo(title='Confirmation',
                 message=f'You have clocked out for break at {self.clockTime}')
        routes.updateBreakOut(userId, self.clockTime)
        self.destroy()

    def Submit(self):
        global userId
        userId = self.idEntry.get()
        try:
            showinfo(title='Conformation',
                     message=f'Employee {routes.read(userId)} was Submitted')
            self.idInput.destroy()
            self.submit.destroy()
            # Clock in Entry
            self.clockIn = ttk.Button(self, text='Clock in', width=30,)
            self.clockIn['command'] = self.clockInFunction
            self.clockIn.pack()
            # Break out entry
            self.breakOut = ttk.Button(self, text='Break Out', width=30)
            self.breakOut['command'] = self.breakOutFunction
            self.breakOut.pack()
            # Break in entry
            self.breakIn = ttk.Button(self, text='Break In', width=30)
            self.breakIn['command'] = self.breakInFunction
            self.breakIn.pack()
            # Clock out Entry
            self.clockOut = ttk.Button(self, text='Clock Out', width=30)
            self.clockOut['command'] = self.clockOutFunction
            self.clockOut.pack()
            self.clock()

        except AttributeError:
            showinfo(
                title='Error', message=f'There is no user with the ID {self.idEntry.get()}')

        except DataError:
            showinfo(
                title='Error', message=f'There is no user with the ID {self.idEntry.get()}')

if __name__ == '__main__':
    app = App()
    app.mainloop()


import unittest
import mock
import sys
import os

module_path = os.path.dirname(os.path.curdir + "." + os.path.sep)
sys.path.insert(0, os.path.abspath(module_path+"/source"))

from serial_connection import serial_port_connection, serial_open
from roboapi import Atmega, Buzzer

# Creating an instance object for testing
Test_object = Atmega(9600)

class Serialtest(unittest.TestCase):
    """ Testing serial.open(), serial.write() and serail port exceptions
    """

    @mock.patch('serial_connection.serial.Serial')
    @mock.patch('serial_connection.glob')
    def Test_serial_port_connection(self, mock_glob, mock_serial_port):
        """
        test for correct serial port initialization
        checking whether serial_port_connection is
        called with the right arguments
        """
        self.portdetect = ['ttyUSB0', 'ttyUSB1']
        self.baudrate = 9600
        mock_glob('ttyUSB0')
        serial_port_connection(self.portdetect, baudrate=self.baudrate)
        mock_serial_port.assert_called_with(self.portdetect[0], self.baudrate)

    @mock.patch('serial_connection.serial.sys.exit')
    @mock.patch('serial_connection.glob')
    @mock.patch('serial_connection.serial.Serial')
    def Test_serial_open(self, mock_serial_port,
                         mock_serial_port_path, mock_sys_call):
        """test to check for serial port exception
        if a serial port is open no exception should be thrown
        """
        mock_serial_port_path('ttyUSB0')
        self.portdetect = ['ttyUSB0', 'ttyUSB1']
        self.baudrate = 9600
        mock_serial_port_path('ttyUSB0')
        serial_port_connection(self.portdetect, baudrate=self.baudrate)
        # checking whether sys.exit(0) is called or not
        mock_sys_call.assert_not_called()

    @mock.patch('serial_connection.serial.Serial')
    def Test_serial_write(self, mock_serial_port):
        """
        test for serial port write calls
        """
        self.portdetect = ['ttyUSB0', 'ttyUSB1']
        self.baudrate = 9600
        port = serial_port_connection(self.portdetect, baudrate=self.baudrate)
        # serial_write called with a dummy character
        Test_object.serial_write('p')
        """
        testing whether serial.write of the Pyserial api
        is called with right arguments
        """
        mock_serial_port.return_value.write.assert_called_once_with('p')
        data_buffer = []
        """
        calling config register and checking whether
        serial.write of Pyserial has appropriate calls
        """
        data_buffer = Test_object.config_register('DDRA', Pins=[3],
                                                  set_pins=True)
        expected = [mock.call(x) for x in data_buffer]
        port.write.assert_has_calls(expected)

    @mock.patch('serial_connection.sys')
    @mock.patch('serial_connection.serial.Serial.isOpen')
    def Test_serial_close_exception(self, mock_serial_port, mock_system_call):
        '''check for exit call when the port is closed'''
        self.portdetect = ['ttyUSB0', 'ttyUSB1']
        self.baudrate = 9600
        serial_open(self.baudrate)
        mock_system_call.exit.assert_called()


import numpy as np

np.seterr(invalid='ignore')

from configs.config import BaseEnvConfig
from envs.model.agent import *
from envs.model.mdp import JointState
from shapely.geometry import *

tmp_config = BaseEnvConfig()

def tensor_to_joint_state(state):  # 恢复原先尺度
    robot_states, human_states = state

    robot_states = robot_states.cpu().squeeze(0).data.numpy()
    robot_states = [RobotState(robot_state[0] * tmp_config.env.nlon,
                               robot_state[1] * tmp_config.env.nlat,
                               robot_state[2] * tmp_config.env.rotation_limit,
                               robot_state[3] * tmp_config.env.max_uav_energy) for robot_state in robot_states]

    human_states = human_states.cpu().squeeze(0).data.numpy()
    human_states = [HumanState(human_state[0] * tmp_config.env.nlon,
                               human_state[1] * tmp_config.env.nlat,
                               human_state[2] * tmp_config.env.rotation_limit,
                               human_state[3] * tmp_config.env.num_timestep) for human_state in human_states]

    return JointState(robot_states, human_states)

def tensor_to_robot_states(robot_state_tensor):
    robot_states = robot_state_tensor.cpu().squeeze(0).data.numpy()
    robot_states = [RobotState(robot_state[0] * tmp_config.env.nlon,
                               robot_state[1] * tmp_config.env.nlat,
                               robot_state[2] * tmp_config.env.rotation_limit,
                               robot_state[3] * tmp_config.env.max_uav_energy) for robot_state in robot_states]
    return robot_states

def get_human_position_list(selected_timestep, human_df):
    selected_timestamp = tmp_config.env.start_timestamp + selected_timestep * tmp_config.env.step_time
    selected_data = human_df[human_df.timestamp == selected_timestamp]
    selected_data = selected_data.set_index("id")

    if selected_timestep < tmp_config.env.num_timestep:
        selected_next_data = human_df[human_df.timestamp == (selected_timestamp + tmp_config.env.step_time)]
        selected_next_data = selected_next_data.set_index("id")
    else:
        selected_next_data = None

    return selected_data, selected_next_data

def get_human_position_from_list(selected_timestep, human_id, selected_data, selected_next_data):
    px, py = selected_data.loc[human_id, ["x", "y"]]

    if selected_timestep < tmp_config.env.num_timestep:
        npx, npy = selected_next_data.loc[human_id, ["x", "y"]]
        theta = get_theta(0, 0, npx - px, npy - py)
        # print(px, py, npx, npy, theta)
    else:
        theta = 0

    return px, py, theta

def judge_aoi_update(human_position, robot_position):
    should_reset = False
    for robot_id in range(tmp_config.env.robot_num):
        unit_distance = np.sqrt(np.power(robot_position[robot_id][0] - human_position[0], 2)
                                + np.power(robot_position[robot_id][1] - human_position[1], 2))
        if unit_distance <= tmp_config.env.sensing_range:
            should_reset = True
            break

    return should_reset

def inPoly(polygon, x, y):
    pt = (x, y)
    line = LineString(polygon)
    point = Point(pt)
    polygon = Polygon(line)
    return polygon.contains(point)

def iscrosses(line1, line2):
    if LineString(line1).crosses(LineString(line2)):
        return True
    return False

def crossPoly(square, x1, y1, x2, y2):
    our_line = LineString([[x1, y1], [x2, y2]])
    line1 = LineString([square[0], square[2]])
    line2 = LineString([square[1], square[3]])
    if our_line.crosses(line1) or our_line.crosses(line2):
        return True
    else:
        return False

def judge_collision(new_robot_px, new_robot_py, old_robot_px, old_robot_py):
    if tmp_config.env.no_fly_zone is None:
        return False

    for square in tmp_config.env.no_fly_zone:
        if inPoly(square, new_robot_px, new_robot_py):
            return True
        if crossPoly(square, new_robot_px, new_robot_py, old_robot_px, old_robot_py):
            return True
    return False

def get_theta(x1, y1, x2, y2):
    ang1 = np.arctan2(y1, x1)
    ang2 = np.arctan2(y2, x2)
    theta = np.rad2deg((ang1 - ang2) % (2 * np.pi))
    return theta

def consume_uav_energy(fly_time, hover_time):
    tmp_config = BaseEnvConfig()

    # configs
    Pu = 0.5  # the average transmitted power of each user, W,  e.g. mobile phone
    P0 = 79.8563  # blade profile power, W
    P1 = 88.6279  # derived power, W
    U_tips = 120  # tip speed of the rotor blade of the UAV,m/s
    v0 = 4.03  # the mean rotor induced velocity in the hovering state,m/s
    d0 = 0.6  # fuselage drag ratio
    rho = 1.225  # density of air,kg/m^3
    s0 = 0.05  # the rotor solidity
    A = 0.503  # the area of the rotor disk, m^2
    Vt = tmp_config.env.velocity  # velocity of the UAV,m/s

    Power_flying = P0 * (1 + 3 * Vt ** 2 / U_tips ** 2) + \
                   P1 * np.sqrt((np.sqrt(1 + Vt ** 4 / (4 * v0 ** 4)) - Vt ** 2 / (2 * v0 ** 2))) + \
                   0.5 * d0 * rho * s0 * A * Vt ** 3

    Power_hovering = P0 + P1

    return fly_time * Power_flying + hover_time * Power_hovering

def get_border(ur, lf):
    upper_left = [lf[0], ur[1]]
    upper_right = [ur[0], ur[1]]
    lower_right = [ur[0], lf[1]]
    lower_left = [lf[0], lf[1]]

    coordinates = [
        upper_left,
        upper_right,
        lower_right,
        lower_left,
        upper_left
    ]

    geo_json = {"type": "FeatureCollection",
                "properties": {
                    "lower_left": lower_left,
                    "upper_right": upper_right
                },
                "features": []}

    grid_feature = {
        "type": "Feature",
        "geometry": {
            "type": "Polygon",
            "coordinates": [coordinates],
        }
    }

    geo_json["features"].append(grid_feature)

    return geo_json

def traj_to_timestamped_geojson(index, trajectory, robot_num, color):
    point_gdf = trajectory.df.copy()
    point_gdf["previous_geometry"] = point_gdf["geometry"].shift()
    point_gdf["time"] = point_gdf.index
    point_gdf["previous_time"] = point_gdf["time"].shift()

    features = []

    # for Point in GeoJSON type
    for _, row in point_gdf.iterrows():
        corrent_point_coordinates = [
            row["geometry"].xy[0][0],
            row["geometry"].xy[1][0]
        ]
        current_time = [row["time"].isoformat()]

        if index < robot_num:
            radius = 8  # 125(5 units)
            opacity = 0.05
            popup_html = f'
UAV {int(row["id"])}
' + f'
raw coord: {corrent_point_coordinates}
' \
                         + f'
grid coord: ({row["x"]},{row["y"]})
' \
                         + f'
dist coord: ({row["x_distance"]}m, {row["y_distance"]}m)
' \
                         + f'
energy: {row["energy"]}J
'
        else:
            radius = 2
            opacity = 1
            popup_html = f'
Human {int(row["id"])}
' + f'
raw coord: {corrent_point_coordinates}
' \
                         + f'
grid coord: ({row["x"]},{row["y"]})
' \
                         + f'
dist coord: ({row["x_distance"]}m, {row["y_distance"]}m)
' \
                         + f'
aoi: {int(row["aoi"])}
'

        # for Point in GeoJSON type  (Temporally Deprecated)
        features.append(
            {
                "type": "Feature",
                "geometry": {
                    "type": "Point",
                    "coordinates": corrent_point_coordinates,
                },
                "properties": {
                    "times": current_time,
                    'popup': popup_html,
                    "icon": 'circle',  # point
                    "iconstyle": {
                        'fillColor': color,
                        'fillOpacity': opacity,  # 透明度
                        'stroke': 'true',
                        'radius': radius,
                        'weight': 1,
                    },

                    "style": {  # line
                        "color": color,
                    },
                    "code": 11,

                },
            }
        )
    return features

if __name__ == "__main__":
    print(judge_collision(new_robot_px=6505, new_robot_py=5130,
                          old_robot_px=6925, old_robot_py=5130))


## 0926
## week1 assignment spell checker by claire

import re
from collections import Counter
import streamlit as st

def words(text): return re.findall(r"[^\d\W]+\'*[^\d\W]+", text.lower())
## match all words without numbers
## include ' if there's any

WORDS = Counter(words(open('big.txt').read()))

def P(word, N=sum(WORDS.values())):
    "Probability of `word`."
    return WORDS[word] / N

def correction(word):
    "Most probable spelling correction for word."
    return max(candidates(word), key=P)

def candidates(word):
    "Generate possible spelling corrections for word."
    return (known([word]) | known(edits1(word)) or known(edits2(word)) or [word])

def known(words):
    "The subset of `words` that appear in the dictionary of WORDS."
    return set(w for w in words if w in WORDS)

def edits1(word):
    "All edits that are one edit away from `word`."
    letters    = "abcdefghijklmnopqrstuvwxyz'"
    splits     = [(word[:i], word[i:])    for i in range(len(word) + 1)]
    deletes    = [L + R[1:]               for L, R in splits if R]
    transposes = [L + R[1] + R[0] + R[2:] for L, R in splits if len(R)>1]
    replaces   = [L + c + R[1:]           for L, R in splits if R for c in letters]
    inserts    = [L + c + R               for L, R in splits for c in letters]
    return set(deletes + transposes + replaces + inserts)

def edits2(word):
    "All edits that are two edits away from `word`."
    return (e2 for e1 in edits1(word) for e2 in edits1(e1))

## Evaluation

#### spelltest(Testset(open('spell-testset1.txt'))) # Development set

## App

st.title("Spell Checker Demo")

origin = st.sidebar.checkbox("Show original word", value=False, key=None, help=None)

input1  = ""
input1 = st.selectbox("Choose a word", ['love','happy', 'teecher'])

input2 = ""
input2 = st.text_input("Key in a word:")

test_input = ""
if (input1 != ""):
    test_input = input1
if (input2 != ""):
    test_input = input2

if (test_input != ""):
    if (origin):
        st.text(correction(test_input))

    if (test_input == correction(test_input)):
        st.success("Correct!", icon="âœ…")
        st.balloons()
        test_input = ""
        input2 = ""
        input1 = ""
    else:
        st.error("Wrong", icon="ðŸš¨")
        test_input = ""
        input2 = ""
        input1 = ""