import numba
import numpy as np
import functools
import pandas as pd
import uproot3

def set_plotting_style():
    from matplotlib import rcParams
    rcParams["font.family"] = "sans-serif"
    rcParams["font.sans-serif"] = ["Helvetica", "Arial", "Liberation Sans", "Bitstream Vera Sans", "DejaVu Sans"]
    rcParams['legend.fontsize'] = 11
    rcParams['legend.labelspacing'] = 0.2
    rcParams['hatch.linewidth'] = 0.5  # https://stackoverflow.com/questions/29549530/how-to-change-the-linewidth-of-hatch-in-matplotlib
    rcParams['axes.xmargin'] = 0.0 # rootlike, no extra padding within x axis
    rcParams['axes.labelsize'] = 'x-large'
    rcParams['axes.formatter.use_mathtext'] = True
    rcParams['legend.framealpha'] = 0.65
    rcParams['axes.labelsize'] = 'x-large'
    rcParams['axes.titlesize'] = 'large'
    rcParams['xtick.labelsize'] = 'large'
    rcParams['ytick.labelsize'] = 'large'
    rcParams['figure.subplot.hspace'] = 0.1
    rcParams['figure.subplot.wspace'] = 0.1
    rcParams['figure.subplot.right'] = 0.96
    rcParams['figure.max_open_warning'] = 0
    rcParams['figure.dpi'] = 100
    rcParams["axes.formatter.limits"] = [-5,4] # scientific notation if log(y) outside this

@numba.njit()
def compute_bin_1d_uniform(x, bins, overflow=False):
    n = bins.shape[0] - 1
    b_min = bins[0]
    b_max = bins[-1]
    if overflow:
        if x > b_max: return n-1
        elif x < b_min: return 0
    ibin = int(n * (x - b_min) / (b_max - b_min))
    if x < b_min or x > b_max:
        return -1
    else:
        return ibin

@numba.njit()
def numba_histogram(a, bins, weights=None,overflow=False):
    db = np.ediff1d(bins)
    is_uniform_binning = np.all(db-db[0]<1e-6)
    hist = np.zeros((len(bins)-1,), dtype=np.float64)
    a = a.flat
    b_min = bins[0]
    b_max = bins[-1]
    n = bins.shape[0] - 1
    if weights is None:
        weights = np.ones(len(a),dtype=np.float64)
    if is_uniform_binning:
        for i in range(len(a)):
            ibin = compute_bin_1d_uniform(a[i], bins, overflow=overflow)
            if ibin >= 0:
                hist[ibin] += weights[i]
    else:
        ibins = np.searchsorted(bins, a, side='left')
        for i in range(len(a)):
            ibin = ibins[i]
            if overflow:
                if ibin == n+1: ibin = n
                elif ibin == 0: ibin = 1
            if ibin >= 1 and ibin <= n:
                hist[ibin-1] += weights[i]
        pass
    return hist, bins

@numba.njit()
def numba_histogram2d(ax,ay, bins_x, bins_y, weights=None,overflow=False):
    db_x = np.ediff1d(bins_x)
    db_y = np.ediff1d(bins_y)
    is_uniform_binning_x = np.all(db_x-db_x[0]<1e-6)
    is_uniform_binning_y = np.all(db_y-db_y[0]<1e-6)
    hist = np.zeros((len(bins_x)-1,len(bins_y)-1), dtype=np.float64)
    ax = ax.flat
    ay = ay.flat
    b_min_x = bins_x[0]
    b_max_x = bins_x[-1]
    n_x = bins_x.shape[0] - 1
    b_min_y = bins_y[0]
    b_max_y = bins_y[-1]
    n_y = bins_y.shape[0] - 1
    if weights is None:
        weights = np.ones(len(ax),dtype=np.float64)
    if is_uniform_binning_x and is_uniform_binning_y:
        for i in range(len(ax)):
            ibin_x = compute_bin_1d_uniform(ax[i], bins_x, overflow=overflow)
            ibin_y = compute_bin_1d_uniform(ay[i], bins_y, overflow=overflow)
            if ibin_x >= 0 and ibin_y >= 0:
                hist[ibin_x,ibin_y] += weights[i]
    else:
        ibins_x = np.searchsorted(bins_x, ax, side='left')
        ibins_y = np.searchsorted(bins_y, ay, side='left')
        for i in range(len(ax)):
            ibin_x = ibins_x[i]
            ibin_y = ibins_y[i]
            if overflow:
                if ibin_x == n_x+1: ibin_x = n_x
                elif ibin_x == 0: ibin_x = 1
                if ibin_y == n_y+1: ibin_y = n_y
                elif ibin_y == 0: ibin_y = 1
            if ibin_x >= 1 and ibin_y >= 1 and ibin_x <= n_x and ibin_y <= n_y:
                hist[ibin_x-1,ibin_y-1] += weights[i]
    return hist, bins_x, bins_y

def make_profile(tobin,toreduce,edges=None,errors=True):
    from scipy.stats import binned_statistic
    yvals = binned_statistic(tobin,toreduce, 'mean', bins=edges).statistic
    yerr = yvals*0.
    if errors:
        yerr = binned_statistic(tobin,toreduce, 'std', bins=edges).statistic/binned_statistic(tobin,toreduce, 'count', bins=edges).statistic**0.5
    from yahist import Hist1D
    h = Hist1D()
    h._counts = yvals
    h._errors = yerr
    h._edges = edges
    return h

@functools.lru_cache(maxsize=256)
def get_chunking(filelist, chunksize, treename="Events", workers=12, skip_bad_files=False, xrootd=False, client=None, use_dask=False):
    """
    Return 2-tuple of
    - chunks: triplets of (filename,entrystart,entrystop) calculated with input `chunksize` and `filelist`
    - total_nevents: total event count over `filelist`
    """
    import uproot3
    from tqdm.auto import tqdm
    import concurrent.futures

    if xrootd:
        temp = []
        for fname in filelist:
            if fname.startswith("/hadoop/cms"):
                temp.append(fname.replace("/hadoop/cms","root://redirector.t2.ucsd.edu/"))
            else:
                temp.append(fname.replace("/store/","root://cmsxrootd.fnal.gov//store/"))
        filelist = temp

    chunksize = int(chunksize)
    chunks = []
    nevents = 0

    if use_dask:
        if not client:
            from dask.distributed import get_client
            client = get_client()
        def numentries(fname):
            import uproot3
            try:
                return (fname,uproot3.numentries(fname,treename))
            except:
                return (fname,-1)
        info = client.gather(client.map(numentries, filelist))
        for fn, nentries in info:
            if nentries < 0:
                if skip_bad_files:
                    print("Skipping bad file: {}".format(fn))
                    continue
                else: raise RuntimeError("Bad file: {}".format(fn))
            nevents += nentries
            for index in range(nentries // chunksize + 1):
                chunks.append((fn, chunksize*index, min(chunksize*(index+1), nentries)))
    else:
        if skip_bad_files:
            # slightly slower (serial loop), but can skip bad files
            for fname in tqdm(filelist):
                try:
                    items = uproot3.numentries(fname, treename, total=False).items()
                except (IndexError, ValueError) as e:
                    print("Skipping bad file", fname)
                    continue
                for fn, nentries in items:
                    nevents += nentries
                    for index in range(nentries // chunksize + 1):
                        chunks.append((fn, chunksize*index, min(chunksize*(index+1), nentries)))
        else:
            executor = None if len(filelist) < 5 else concurrent.futures.ThreadPoolExecutor(min(workers, len(filelist)))
            for fn, nentries in uproot3.numentries(filelist, treename, total=False, executor=executor).items():
                nevents += nentries
                for index in range(nentries // chunksize + 1):
                    chunks.append((fn, chunksize*index, min(chunksize*(index+1), nentries)))

    return chunks, nevents

def hist2d_dask(df, x, y, bins, method=1):
    """
    np.histogram2d from dask dataframe.
    3 methods to test:
        1 - map over delayed objects from dataframe
        2 - map over partitions from dataframe
        3 - map over partitions from dataframe using numba

    Examples
    --------
    >>> bins = [np.linspace(-15,15,200),np.linspace(-15,15,200)]
    >>> hist2d_dask(df, x="DV_x", y="DV_y", bins=bins).compute()
    """
    if method == 1:
        from dask import delayed
        @delayed
        def f(df, bins):
            return np.histogram2d(df.eval(x),df.eval(y), bins=bins)[0]
        bins = delayed(bins)
        agg = delayed(sum)
        return agg(f(obj, bins) for obj in df.to_delayed())
    if method == 2:
        def f(df, bins):
            return np.histogram2d(df.eval(x), df.eval(y), bins=bins)[0].flatten()[np.newaxis]
        # FIXME, probably better way to do this
        # can specify a (x, y bins) or just 1d bins to be used for x and y
        # need to reshape differently depending on case
        if len(bins) == 2:
            shape = [bins[0].shape[0]-1, bins[1].shape[0]-1]
        else:
            shape = [bins.shape[0]-1, bins.shape[0]-1]
        return df.map_partitions(f, bins).sum(axis=0).reshape(shape)
    if method == 3:
        def f(df, bins):
            if len(bins) == 2:
                binsx, binsy = bins
            else:
                binsx, binsy = bins, bins
            return numba_histogram2d(df.eval(x).values, df.eval(y).values, binsx, binsy, None, False)[0].flatten()[np.newaxis]
        # FIXME, probably better way to do this
        # can specify a (x, y bins) or just 1d bins to be used for x and y
        # need to reshape differently depending on case
        if len(bins) == 2:
            shape = [bins[0].shape[0]-1, bins[1].shape[0]-1]
        else:
            shape = [bins.shape[0]-1, bins.shape[0]-1]
        return df.map_partitions(f, bins).sum(axis=0).reshape(shape)

def get_geometry_df(fname):
    """
    Get pixel geometry from inputs made in `geometry/`
    """
    import uproot3
    import numpy as np
    f = uproot3.open(fname)
    t = f["idToGeo"]
    df = t.pandas.df(branches=["shape","translation","matrix"],flatten=False)
    df["translation_x"] = df["translation"].str[0]
    df["translation_y"] = df["translation"].str[1]
    df["translation_z"] = df["translation"].str[2]
    df["translation_rho"] = np.hypot(df["translation_x"],df["translation_y"])
    df = df[df["shape"].apply(lambda x:x[0])==2.]
    df["endcap"] = df.eval("abs(translation_z)>25")
    # layer 1-4 for barrel, 5,7,8 for z disks
    layers = []
    for rho, z in df[["translation_rho","translation_z"]].values:
        ilayer = -1
        if abs(z)<25:
            if 0 < rho < 5: ilayer = 1
            if 5 < rho < 9: ilayer = 2
            if 9 < rho < 14: ilayer = 3
            if 14 < rho < 25: ilayer = 4
        else:
            if 25 < abs(z) < 35: ilayer = 5
            if 35 < abs(z) < 45: ilayer = 6
            if 45 < abs(z) < 70: ilayer = 7
        layers.append(ilayer)
    df["layer"] = np.array(layers)
    df = df.query("translation_rho<18") # 4 pixel layers
    return df

def plot_overlay_bpix(ax,**kwargs):
    """
    Given an axes object, overlays 2D lines for the transverse projection of the first 3 bpix layers
    Note the hardcoded geometry path
    """
    import numpy as np
    color = kwargs.pop("color","k")
    binary_triplets = np.unpackbits(np.arange(8,dtype=np.uint8)[:,np.newaxis],1)[:,-3:].astype(int)
    step_directions = binary_triplets*2-1
    geometryfile = kwargs.pop("geometryfile","/home/users/namin/2019/scouting/repo/geometry/tracker_geometry_data2018.root")
    gdf = get_geometry_df(geometryfile)
    expand_l = kwargs.pop("expand_l",0.00)
    expand_w = kwargs.pop("expand_w",0.00)
    expand_h = kwargs.pop("expand_h",0.05)
    do_expand = (expand_h > 0) or (expand_w > 0) or (expand_h)
    for irow,entry in gdf.query("0 < translation_z < 8 and translation_rho<{}".format(kwargs.pop("maxrho",14))).iterrows():
        shape = entry["shape"][1:-1].T
        if do_expand:
            newshape = np.array(shape)
            newshape[0] += expand_l
            newshape[1] += expand_w
            newshape[2] += expand_h
            shape = newshape
        translation = entry["translation"]
        matrix = entry["matrix"].reshape(3,3)
        points = shape * step_directions
        points = np.array([np.dot(matrix,point)+translation for point in points])
        points = points[np.array([6,2,1,5,6])]
        ax.plot(points[:,0],points[:,1],color=color,**kwargs)
    return ax

def plot_overlay_pixel_rhoz(ax,**kwargs):
    """
    Given an axes object, overlays 2D lines for the rho-z projection of the pixel layers
    """
    import numpy as np
    color = kwargs.pop("color","k")
    lw = kwargs.pop("lw",1.0)
    alpha = kwargs.pop("alpha",0.015)
    binary_triplets = np.unpackbits(np.arange(8,dtype=np.uint8)[:,np.newaxis],1)[:,-3:].astype(int)
    step_directions = binary_triplets*2-1
    geometryfile = kwargs.pop("geometryfile","/home/users/namin/2019/scouting/repo/geometry/tracker_geometry_data2018.root")
    gdf = get_geometry_df(geometryfile)
    expand_l = kwargs.pop("expand_l",0.00)
    expand_w = kwargs.pop("expand_w",0.00)
    expand_h = kwargs.pop("expand_h",0.00)
    do_expand = (expand_h > 0) or (expand_w > 0) or (expand_h)
#     for irow,entry in gdf.query("0 < translation_z < 8 and translation_rho<14").iterrows():
    for irow,entry in gdf.iterrows():
        shape = entry["shape"][1:-1].T
        if do_expand:
            newshape = np.array(shape)
            newshape[0] += expand_l
            newshape[1] += expand_w
            newshape[2] += expand_h
            shape = newshape
        translation = entry["translation"]
        matrix = entry["matrix"].reshape(3,3)
        points = shape * step_directions
        points = np.array([np.dot(matrix,point)+translation for point in points])
        points = points[np.array([6,2,1,5,6])]
        rho = np.hypot(points[:,0],points[:,1])
        p1,p2=points[:,2],rho
        ax.plot(p1,p2,color=color,alpha=alpha,lw=lw,**kwargs)
    return ax

def futures_widget(futures):
    """
    Takes a list of futures and returns a jupyter widget object of squares,
    one per future, which turn green or red on success or failure
    """
    import ipywidgets
    import toolz
    def make_button(future):
        button = ipywidgets.Button(
            description='',
            disabled=False,
            button_style='', # 'success', 'info', 'warning', 'danger' or ''
            tooltip='',
            layout=ipywidgets.Layout(width='20px', height='20px'),
        )
        def callback(f):
            if f.exception():
                button.button_style = "danger"
            else:
                button.button_style = "success"
        future.add_done_callback(callback)
        return button

    items = [make_button(future) for future in futures]

    box = ipywidgets.VBox([ipywidgets.HBox(row) for row in toolz.itertoolz.partition(30, items)])
    return box

@pd.api.extensions.register_dataframe_accessor("tree")
class TreeLikeAccessor:
    def __init__(self, pandas_obj):
        self._obj = pandas_obj

    def draw(self, varexp, sel, bins=None, overflow=True, fast=True, label=None):
        try:
            from yahist import Hist1D
        except:
            raise Exception("Need Hist1D object from the yahist package")

        df = self._obj

        if fast:
            return Hist1D(hacky_query_eval(df, varexp, sel), bins=bins, label=label)

        weights = df.eval(sel)
        mask = np.zeros(len(df), dtype=bool)
        extra = dict()
        extra["overflow"] = overflow
        if (weights.dtype in [int, np.int32]):
            mask = weights != 0
            extra["weights"] = weights[mask]
        if (weights.dtype == bool):
            mask = weights > 0.5
            # no weights for bools
        if (weights.dtype == float):
            mask = weights != 0.
            extra["weights"] = weights[mask]
        vals = df[mask].eval(varexp)
        if bins is not None:
            # if type(bins) in [str]:
            #     raise NotImplementedError()
            extra["bins"] = bins
        return Hist1D(vals, **extra)

@pd.api.extensions.register_dataframe_accessor("vec")
class LorentzVectorAccessor:
    def __init__(self, pandas_obj):
        self._validate(pandas_obj)
        self._obj = pandas_obj

    @staticmethod
    def _validate(obj):
        missing_columns = set(["Muon1_pt","Muon1_eta","Muon1_phi","Muon2_pt","Muon2_eta","Muon2_phi"])-set(obj.columns)
        if len(missing_columns):
            raise AttributeError("Missing columns: {}".format(missing_columns))

    @property
    def mu1(self):
        import uproot3_methods
        LV = uproot3_methods.TLorentzVectorArray.from_ptetaphim(
            self._obj["Muon1_pt"],self._obj["Muon1_eta"],self._obj["Muon1_phi"],0.10566,
        )
        return LV

    @property
    def mu2(self):
        import uproot3_methods
        LV = uproot3_methods.TLorentzVectorArray.from_ptetaphim(
            self._obj["Muon2_pt"],self._obj["Muon2_eta"],self._obj["Muon2_phi"],0.10566,
        )
        return LV

    @property
    def dimu(self):
        return self.mu1 + self.mu2

smaller_dtypes = [
    ["dimuon_mass","float32"],
    ["dimuon_pt","float32"],
    ["DV_chi2prob","float32"],
    ["DV_ndof","int8"],
    ["DV_redchi2","float32"],
    ["DV_layerPixel","int8"],
    ["Muon1_charge","int8"],
    ["Muon1_excesshits","int8"],
    ["Muon1_m","float32"],
    ["Muon1_mass","float32"],
    ["Muon1_nExcessPixelHits","int8"],
    ["Muon1_nExpectedPixelHits","int8"],
    ["Muon1_nMatchedStations","int8"],
    ["Muon1_nTrackerLayersWithMeasurement","int8"],
    ["Muon1_nValidMuonHits","int8"],
    ["Muon1_nValidPixelHits","int8"],
    ["Muon1_nValidStripHits","int8"],
    ["Muon2_charge","int8"],
    ["Muon2_excesshits","int8"],
    ["Muon2_m","float32"],
    ["Muon2_mass","float32"],
    ["Muon2_nExcessPixelHits","int8"],
    ["Muon2_nExpectedPixelHits","int8"],
    ["Muon2_nMatchedStations","int8"],
    ["Muon2_nTrackerLayersWithMeasurement","int8"],
    ["Muon2_nValidMuonHits","int8"],
    ["Muon2_nValidPixelHits","int8"],
    ["Muon2_nValidStripHits","int8"],
    ["categ","int8"],
    ["luminosityBlock","int32"],
    ["nDV","int8"],
    ["nDV_good","int8"],
    ["nDV_raw","int8"],
    ["nGenMuon","int8"],
    ["nGenPart","int16"],
    ["nJet","int8"],
    ["nMuon","int8"],
    ["nMuon_good","int8"],
    ["nMuon_raw","int8"],
    ["nPV","int8"],
    ["nPVM","int8"],
    ["run","int32"],
]

def make_df(
    path,
    branches = ["dimuon_mass", "pass_*"],
    cut = "pass_baseline_iso",
    chunksize = 500e3,
    xrootd = False,
    persist = True,
    client = None,
    func = None,
    partition_size = None,
    npartitions = None,
    use_dask = False,
    skip_bad_files = False,
    nthreads = 6,
    progress = True,
):
    """
    Returns dataframe from input ROOT files containing given branches
    for events passing a given cut. If `use_dask=True`, returns dask dataframe.

    path: file path(s) or glob string(s)
    branches: list of branches/glob strings/regex for branches to read
    cut: selection string input to `df.query()`
    chunksize: events per task
    xrootd: use xrootd for input files
    persist: whether to return persisted dask dataframe or not
    func: override reading function (must read a chunk and return a DataFrame)
    partition_size: if not None, passed into df.repartition() before persisting - NOTE might be duplicating the reading. use `npartitions`.
    npartitions: if not None, passed into df.repartition() before persisting
    skip_bad_files: whether to skip bad files according to failure of uproot.numentries
    nthreads: number of ThreadPoolExecutor threads when not using dask (default 6)
    progress: show progress bar
    """
    import dask.dataframe as dd
    from dask import delayed
    import uproot3
    import pandas as pd
    from tqdm.auto import tqdm
    import concurrent.futures
    if isinstance(path, (str, bytes)):
        paths = uproot3.tree._filename_explode(path)
    else:
        paths = [y for x in path for y in uproot3.tree._filename_explode(x)]

    if not func:
        def func(fname, entrystart = None, entrystop = None):
            t = uproot3.open(fname)["Events"]
            arrs = t.arrays(
                branches,
                outputtype = dict,
                namedecode = "ascii",
                entrystart = entrystart,
                entrystop = entrystop,
            )
            sel = slice(None,None)
            df = pd.DataFrame()
            for k in arrs.keys():
                v = arrs[k][sel]
                is_jagged = ("Jagged" in str(type(v)))
                if is_jagged:
                    if k.startswith("DV_"):
                        df[k] = v[:,0]
                    elif k.startswith("PVM_"):
                        df[k] = v[:,0]
                    elif k.startswith("Muon_"):
                        df[k.replace("Muon_","Muon1_")] = v[:,0]
                        df[k.replace("Muon_","Muon2_")] = v[:,1]
                    elif k.startswith("Jet_"):
                        if "ht" not in df.columns:
                            df["ht"] = v.sum()
                        v = v.pad(2).fillna(-1)
                        df[k.replace("Jet_","Jet1_")] = v[:,0]
                        df[k.replace("Jet_","Jet2_")] = v[:,1]
                    else:
                        df[k] = v
                else:
                    df[k] = v
            for name,dtype in smaller_dtypes:
                if name not in df.columns: continue
                df[name] = df[name].astype(dtype, copy=False)
            df = df.query(cut)
            return df

    chunks, total_events = get_chunking(tuple(paths), chunksize, client=client, xrootd=xrootd, use_dask=use_dask, skip_bad_files=skip_bad_files)

    smallchunk_nevents = int(max(chunks[0][1] + (chunks[0][2]-chunks[0][1])//10, 1))
    smallchunk = (chunks[0][0], chunks[0][1], smallchunk_nevents)
    meta = func(*smallchunk)

    if not use_dask:
        smallchunk_mb = meta.memory_usage().sum()/1e6
        estimated_mb = smallchunk_mb * total_events / smallchunk_nevents
        if estimated_mb > 15e3:
            raise RuntimeError("This dataframe would take approx. {:.1f}GB of RAM. Reduce the input size.".format(estimated_mb*1e-3))

        executor = concurrent.futures.ThreadPoolExecutor(nthreads)
        futures = [executor.submit(func, *chunk) for chunk in chunks]
        def wrapper(x):
            if progress:
                return tqdm(x)
            return x
        ddf = pd.concat((future.result() for future in wrapper(futures)), sort=True, ignore_index=True, copy=False)
        del executor
    else:
        delayed_func = delayed(func)
        ddf = dd.from_delayed((delayed_func(*chunk) for chunk in chunks), meta=meta).reset_index(drop=True)
        if partition_size:
            ddf = ddf.repartition(partition_size=partition_size)
        if npartitions:
            ddf = ddf.repartition(npartitions=npartitions)
        if persist:
            ddf = ddf.persist()

    return ddf

def dataframe_to_ttree(df, filename, treename="t", chunksize=1e6, compression=uproot3.LZ4(1), progress=True):
    """
    Writes ROOT file containing one TTree with the input pandas DataFrame.

    filename: name of output file
    treename: name of output TTree
    chunksize: number of rows per basket
    compression: uproot compression object (LZ4, ZLIB, LZMA, or None)
    progress: show tqdm progress bar?
    """
    t = uproot3.newtree(df.dtypes)
    with uproot3.recreate(filename, compression=compression) as f:
        f[treename] = t
        chunksize = int(chunksize)
        iterable = range(0, len(df), chunksize)
        if progress:
            from tqdm.auto import tqdm
            iterable = tqdm(iterable)
        for i in iterable:
            chunk = df.iloc[i:i+chunksize]
            f[treename].extend({ k:chunk[k].values for k in chunk.columns })

def ttree_to_dataframe(filename, treename="t", branches=None, progress=True, **kwargs):
    """
    Read ROOT file containing one TTree into pandas DataFrame.
    Thin wrapper around `uproot.iterate`.

    filename: filename(s)/file pattern(s)
    treename: name of input TTree
    progress: show tqdm progress bar?
    branches: list of branches to read (default of `None` reads all)
    **kwargs: extra kwargs to pass to `uproot.iterate`
    """
    # entrysteps of None iterates by basket to match `dataframe_to_ttree`
    iterable = uproot3.iterate(filename, treename, branches,
                              entrysteps=kwargs.pop("entrysteps",None), outputtype=pd.DataFrame,
                              namedecode="ascii", **kwargs
                             )
    if progress:
        from tqdm.auto import tqdm
        iterable = tqdm(iterable)
    # df = pd.concat(iterable, ignore_index=True, sort=True, copy=False)
    df = pd.concat(iterable, ignore_index=True, sort=True)
    return df

def hacky_query_eval(df, varstr, selstr="", verbose=False):
    """
    Please don't read/use. This is dangerous and stupid, kind of like
    integrating a function by printing out a plot, coloring the area under it in red,
    faxing it to yourself, then counting red pixels to get the area.

    Basically I wanted some way to convert

        df.query("dimuon_mass > 5 and pass_baseline_iso").eval("dimuon_mass").mean()

    into

        df["dimuon_mass"][ (df["dimuon_mass"] > 5) & (df["pass_baseline_iso"]) ].mean()

    because the latter doesn't make an intermediate copy of all the columns with query(),
    and it also doesn't do jitting with numexpr. In principle, this is much faster to execute.

    Usage:

        arr = hacky_query_eval(
            df_data,
            varstr = "dimuon_mass",
            selstr = "pass_baseline_iso and 0<logabsetaphi<1.25",
        )
        print(arr.mean())
    """
    from pandas.core.computation.expr import Expr
    from pandas.core.computation.scope import Scope
    env = Scope(
        1,
        global_dict=globals(),
        local_dict=locals(),
        resolvers=[df],
        target=None,
    )
    def inject_df(s):
        """
        convert expression string like (a > 1) to (df["a"] > 1)
        so that it can be eval'd later
        """
        expr = Expr(s, env=env, parser="pandas")
        self = expr._visitor
        def visit_Name_hack(node, **kwargs):
            result = self.term_type(node.id, self.env, **kwargs)
            result._name = f'df["{result._name}"]'
            return result
        def _maybe_downcast_constants_hack(left, right):
            return left, right
        expr._visitor.visit_Name = visit_Name_hack
        expr._visitor._maybe_downcast_constants = _maybe_downcast_constants_hack
        expr.terms = expr.parse()
        return str(expr)
    varexpr = inject_df(varstr)
    toeval = f"({varexpr})"
    if selstr:
        selexpr = str(inject_df(selstr))
        toeval += f"[{selexpr}]"
    if verbose: print(f"Evaluating string: {toeval}")
    result = eval(toeval)
    return result

def query_dis(query, typ="basic", return_raw=False):
    import requests
    endpoint_url = "http://uaf-7.t2.ucsd.edu:50010/dis/serve"
    short = "short=true" if typ in ["basic", "sites"] else ""
    url = f"{endpoint_url}?type={typ}&query={query}&{short}"
    js = requests.get(url).json()
    if not return_raw:
        js = js["payload"]
    return js


# -*- coding: utf-8 -*-

# 导入库
import copy  # 复制库

import numpy as np  # numpy库
import pandas as pd  # pandas库
from imblearn.over_sampling import SMOTE  # 过抽样处理库SMOTE
from sklearn.ensemble import AdaBoostClassifier, VotingClassifier, GradientBoostingClassifier, \
    RandomForestClassifier,  BaggingClassifier  # 四种集成分类库和投票方法库
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold, cross_val_score  # 导入交叉检验算法
from sklearn.preprocessing import LabelEncoder  # 字符串转数值
from sklearn.metrics import accuracy_score,confusion_matrix

# 数据审查和预处理函数
# 基本状态查看
def set_summary(df):
    '''
    查看数据集后2条数据、数据类型、描述性统计
    :param df: 数据框
    :return: 无
    '''
    print('{:*^60}'.format('Data overview:'))
    print(df.tail(2).T)  # 打印原始数据后2条
    print('{:*^60}'.format('Data dtypes:'))
    print(df.dtypes)  # 打印数据类型
    print('{:*^60}'.format('Data DESC:'))
    print(df.describe().round(2).T)  # 打印原始数据基本描述性信息

# 缺失值审查
def na_summary(df):
    '''
    查看数据集的缺失数据列、行记录数
    :param df: 数据框
    :return: 无
    '''
    na_cols = df.isnull().any(axis=0)  # 查看每一列是否具有缺失值
    print('{:*^60}'.format('NA Cols:'))
    print(na_cols)  # 查看具有缺失值的列
    print('Total number of NA lines is: {0}'.format(df.isnull().any(axis=1).sum()))  # 查看具有缺失值的行总记录数

# 类样本均衡审查
def label_samples_summary(df):
    '''
    查看每个类的样本量分布
    :param df: 数据框
    :return: 无
    '''
    print('{:*^60}'.format('Labesl samples count:'))
    print(df.iloc[:, 0].groupby(df.iloc[:, -1]).count())

# 字符串分类转数值分类
def label_encoder(data, model_list=None, train=True):
    '''
    将特征中的字符串分类转换为数值分类
    :param data: 输入数据集
    :param model_list: LabelEncoder对象列表，在训练阶段产生
    :param train: 是否为训练阶段
    :return: 训练阶段产生训练后的LabelEncoder对象列表和转换后的数据，预测阶段产生转换后的数据
    '''
    convert_cols = ['cat', 'attribution', 'pro_id', 'pro_brand', 'order_source', 'pay_type',
                    'use_id','city']  # 定义要转换的列
    value_list = []  # 存放转换后的数据
    if train:
        model_list = []  # 存放每个特征转换的实例对象
        model_label_encoder = LabelEncoder()
        for i in convert_cols:
            model_label_encoder.fit(data[i])
            value_list.append(model_label_encoder.transform(data[i]))
            model_list.append(copy.copy(model_label_encoder))
        convert_matrix = np.array(value_list).T
        return model_list, convert_matrix
    else:
        for ind, j in enumerate(convert_cols):
            value_list.append(model_list[ind].transform(data[j]))
        convert_matrix = np.array(value_list).T
        return convert_matrix

# 时间属性拓展
def datetime2int(data):
    '''
    将日期和时间数据拓展出其他属性，例如星期几、周几、小时、分钟等。
    :param data: 数据集
    :return: 拓展后的属性矩阵
    '''
    date_set = [pd.datetime.strptime(dates, '%Y-%m-%d') for dates in
                data['order_date']]  # 将data中的order_date列转换为特定日期格式
    '''date_set = [pd.to_datetime(dates, format='%Y-%m-%d') for dates in
                data['order_date']] '''

    weekday_data = [data.weekday() for data in date_set]  # 周几
    daysinmonth_data = [data.day for data in date_set]  # 当月几号
    month_data = [data.month for data in date_set]  # 月份

    time_set = [pd.datetime.strptime(times, '%H:%M:%S') for times in
                data['order_time']]  # 将data中的order_time列转换为特定时间格式
    second_data = [data.second for data in time_set]  # 秒
    minute_data = [data.minute for data in time_set]  # 分钟
    hour_data = [data.hour for data in time_set]  # 小时

    final_set = [weekday_data, daysinmonth_data, month_data, second_data, minute_data,
                 hour_data]  # 将属性列表批量组合
    final_matrix = np.array(final_set).T  # 转换为数组并转置
    return final_matrix

# 样本均衡
def sample_balance(X, y):
    '''
    使用SMOTE方法对不均衡样本做过抽样处理
    :param X: 输入特征变量X
    :param y: 目标变量y
    :return: 均衡后的X和y
    '''
    model_smote = SMOTE()  # 建立SMOTE模型对象
    x_smote_resampled, y_smote_resampled = model_smote.fit_resample(X, y)  # 输入数据并作过抽样处理
    return x_smote_resampled, y_smote_resampled

# 数据应用
# 定义特殊字段数据格式
dtypes = {'order_id': np.object,
          'pro_id': np.object,
          'use_id': np.object}
raw_data = pd.read_table('./data/abnormal_orders.txt', delimiter=',', dtype=dtypes)  # 读取数据集

# 数据审查
set_summary(raw_data)  # 基本状态查看
na_summary(raw_data)  # 缺失值审查
label_samples_summary(raw_data)  # 类样本分布审查

# 数据预处理
drop_na_set = raw_data.dropna()  # 丢弃带有NA值的数据行
X_raw, y_raw = drop_na_set.iloc[:, 1:-1], drop_na_set.iloc[:, -1]  # 分割输入变量X和y
model_list, convert_matrix = label_encoder(X_raw)  # 字符串分类转整数型分类
#print(convert_matrix)
datetime2int_data = datetime2int(X_raw)  # 拓展日期时间属性
#print(datetime2int_data)
combine_set = np.hstack((convert_matrix, datetime2int_data))  # 合并转换后的分类和拓展后的日期数据集
constant_set = X_raw[['total_money', 'total_quantity']]  # 原始连续数据变量
X_combine = np.hstack((combine_set, constant_set))  # 再次合并数据集
# 相关性分析
X_combine=pd.DataFrame(X_combine)
print('{:*^60}'.format('Correlation Analyze:'))
print(X_combine.corr().round(2))  # 输出所有输入特征变量以及预测变量的相关性矩阵

X, y = sample_balance(X_combine, y_raw)  # 样本均衡处理
'''X=X_combine
y=y_raw'''

# 组合分类模型交叉检验
model_rf = RandomForestClassifier(n_estimators=20, random_state=0)  # 随机森林分类模型对象
model_lr = LogisticRegression(max_iter=1000,random_state=0)  # 逻辑回归模型对象
#model_BagC = BaggingClassifier(n_estimators=20,random_state=0)  # Bagging分类模型对象
#model_gdbc = GradientBoostingClassifier(max_features=0.8, random_state=0)  # GradientBoosting分类模型对象
estimators = [('randomforest', model_rf), ('Logistic', model_lr)]  # 建立组合评估器列表
model_vot = VotingClassifier(estimators=estimators, voting='soft', weights=[0.5, 0.5],
                             n_jobs=-1)  # 建立组合评估模型
cv = StratifiedKFold(2)  # 设置交叉检验方法
cv_score = cross_val_score(model_vot, X, y, cv=cv)  # 交叉检验
print('{:*^60}'.format('Cross val scores:'))
print(cv_score)  # 打印每次交叉检验得分
print('Mean scores is: %.2f' % cv_score.mean())  # 打印平均交叉检验得分
model_vot.fit(X, y)  # 模型训练
y_predicttrain = model_vot.predict(X)  # 预测结果
print('混淆矩阵：')
print(confusion_matrix(y,y_predicttrain))

# 新数据集做预测
X_raw_data = pd.read_csv('./data/new_abnormal_orders.csv', dtype=dtypes)  # 读取要预测的数据集
X_raw_new = X_raw_data.iloc[:, 1:]  # 分割输入变量X，并丢弃订单ID列
convert_matrix_new = label_encoder(X_raw_new, model_list, False)  # 字符串分类转整数型分类
datetime2int_data_new = datetime2int(X_raw_new)  # 日期时间转换
combine_set_new = np.hstack((convert_matrix_new, datetime2int_data_new))  # 合并转换后的分类和拓展后的日期数据集
constant_set_new = X_raw_new[['total_money', 'total_quantity']]  # 原始连续数据变量
X_combine_new = np.hstack((combine_set_new, constant_set_new))  # 再次合并数据集
y_predict = model_vot.predict(X_combine_new)  # 预测结果
print('{:*^60}'.format('Predicted Labesls:'))
print(y_predict)  # 打印预测值


"""Set of methods to create custom architecture."""

from typing import Any, Dict, List, Tuple, Union

import torch

from espnet.nets.pytorch_backend.conformer.convolution import ConvolutionModule
from espnet.nets.pytorch_backend.conformer.encoder_layer import (
    EncoderLayer as ConformerEncoderLayer,
)
from espnet.nets.pytorch_backend.nets_utils import get_activation
from espnet.nets.pytorch_backend.transducer.conv1d_nets import CausalConv1d, Conv1d
from espnet.nets.pytorch_backend.transducer.transformer_decoder_layer import (
    TransformerDecoderLayer,
)
from espnet.nets.pytorch_backend.transducer.vgg2l import VGG2L
from espnet.nets.pytorch_backend.transformer.attention import (
    MultiHeadedAttention,
    RelPositionMultiHeadedAttention,
)
from espnet.nets.pytorch_backend.transformer.embedding import (
    PositionalEncoding,
    RelPositionalEncoding,
    ScaledPositionalEncoding,
)
from espnet.nets.pytorch_backend.transformer.encoder_layer import EncoderLayer
from espnet.nets.pytorch_backend.transformer.positionwise_feed_forward import (
    PositionwiseFeedForward,
)
from espnet.nets.pytorch_backend.transformer.repeat import MultiSequential
from espnet.nets.pytorch_backend.transformer.subsampling import Conv2dSubsampling

def verify_block_arguments(
    net_part: str,
    block: Dict[str, Any],
    num_block: int,
) -> Tuple[int, int]:
    """Verify block arguments are valid.

    Args:
        net_part: Network part, either 'encoder' or 'decoder'.
        block: Block parameters.
        num_block: Block ID.

    Return:
        block_io: Input and output dimension of the block.

    """
    block_type = block.get("type")

    if block_type is None:
        raise ValueError(
            "Block %d in %s doesn't a type assigned.", (num_block, net_part)
        )

    if block_type == "transformer":
        arguments = {"d_hidden", "d_ff", "heads"}
    elif block_type == "conformer":
        arguments = {
            "d_hidden",
            "d_ff",
            "heads",
            "macaron_style",
            "use_conv_mod",
        }

        if net_part == "decoder":
            raise ValueError("Decoder does not support 'conformer'.")

        if block.get("use_conv_mod", None) is True and "conv_mod_kernel" not in block:
            raise ValueError(
                "Block %d: 'use_conv_mod' is True but "
                " 'conv_mod_kernel' is not specified" % num_block
            )
    elif block_type == "causal-conv1d":
        arguments = {"idim", "odim", "kernel_size"}

        if net_part == "encoder":
            raise ValueError("Encoder does not support 'causal-conv1d'.")

    elif block_type == "conv1d":
        arguments = {"idim", "odim", "kernel_size"}

        if net_part == "decoder":
            raise ValueError("Decoder does not support 'conv1d.'")
    else:
        raise NotImplementedError(
            "Wrong type. Currently supported: "
            "causal-conv1d, conformer, conv-nd or transformer."
        )

    if not arguments.issubset(block):
        raise ValueError(
            "%s in %s in position %d: Expected block arguments : %s."
            " See tutorial page for more information."
            % (block_type, net_part, num_block, arguments)
        )

    if block_type in ("transformer", "conformer"):
        block_io = (block["d_hidden"], block["d_hidden"])
    else:
        block_io = (block["idim"], block["odim"])

    return block_io

def prepare_input_layer(
    input_layer_type: str,
    feats_dim: int,
    blocks: List[Dict[str, Any]],
    dropout_rate: float,
    pos_enc_dropout_rate: float,
) -> Dict[str, Any]:
    """Prepare input layer arguments.

    Args:
        input_layer_type: Input layer type.
        feats_dim: Dimension of input features.
        blocks: Blocks parameters for network part.
        dropout_rate: Dropout rate for input layer.
        pos_enc_dropout_rate: Dropout rate for input layer pos. enc.

    Return:
        input_block: Input block parameters.

    """
    input_block = {}
    first_block_type = blocks[0].get("type", None)

    if first_block_type == "causal-conv1d":
        input_block["type"] = "c-embed"
    else:
        input_block["type"] = input_layer_type

    input_block["dropout-rate"] = dropout_rate
    input_block["pos-dropout-rate"] = pos_enc_dropout_rate

    input_block["idim"] = feats_dim

    if first_block_type in ("transformer", "conformer"):
        input_block["odim"] = blocks[0].get("d_hidden", 0)
    else:
        input_block["odim"] = blocks[0].get("idim", 0)

    return input_block

def prepare_body_model(
    net_part: str,
    blocks: List[Dict[str, Any]],
) -> Tuple[int]:
    """Prepare model body blocks.

    Args:
        net_part: Network part, either 'encoder' or 'decoder'.
        blocks: Blocks parameters for network part.

    Return:
        : Network output dimension.

    """
    cmp_io = [
        verify_block_arguments(net_part, b, (i + 1)) for i, b in enumerate(blocks)
    ]

    if {"transformer", "conformer"} <= {b["type"] for b in blocks}:
        raise NotImplementedError(
            net_part + ": transformer and conformer blocks "
            "can't be used together in the same net part."
        )

    for i in range(1, len(cmp_io)):
        if cmp_io[(i - 1)][1] != cmp_io[i][0]:
            raise ValueError(
                "Output/Input mismatch between blocks %d and %d in %s"
                % (i, (i + 1), net_part)
            )

    return cmp_io[-1][1]

def get_pos_enc_and_att_class(
    net_part: str, pos_enc_type: str, self_attn_type: str
) -> Tuple[
    Union[PositionalEncoding, ScaledPositionalEncoding, RelPositionalEncoding],
    Union[MultiHeadedAttention, RelPositionMultiHeadedAttention],
]:
    """Get positional encoding and self attention module class.

    Args:
        net_part: Network part, either 'encoder' or 'decoder'.
        pos_enc_type: Positional encoding type.
        self_attn_type: Self-attention type.

    Return:
        pos_enc_class: Positional encoding class.
        self_attn_class: Self-attention class.

    """
    if pos_enc_type == "abs_pos":
        pos_enc_class = PositionalEncoding
    elif pos_enc_type == "scaled_abs_pos":
        pos_enc_class = ScaledPositionalEncoding
    elif pos_enc_type == "rel_pos":
        if net_part == "encoder" and self_attn_type != "rel_self_attn":
            raise ValueError("'rel_pos' is only compatible with 'rel_self_attn'")
        pos_enc_class = RelPositionalEncoding
    else:
        raise NotImplementedError(
            "pos_enc_type should be either 'abs_pos', 'scaled_abs_pos' or 'rel_pos'"
        )

    if self_attn_type == "rel_self_attn":
        self_attn_class = RelPositionMultiHeadedAttention
    else:
        self_attn_class = MultiHeadedAttention

    return pos_enc_class, self_attn_class

def build_input_layer(
    block: Dict[str, Any],
    pos_enc_class: torch.nn.Module,
    padding_idx: int,
) -> Tuple[Union[Conv2dSubsampling, VGG2L, torch.nn.Sequential], int]:
    """Build input layer.

    Args:
        block: Architecture definition of input layer.
        pos_enc_class: Positional encoding class.
        padding_idx: Padding symbol ID for embedding layer (if provided).

    Returns:
        : Input layer module.
        subsampling_factor: Subsampling factor.

    """
    input_type = block["type"]

    idim = block["idim"]
    odim = block["odim"]

    dropout_rate = block["dropout-rate"]
    pos_dropout_rate = block["pos-dropout-rate"]

    if pos_enc_class.__name__ == "RelPositionalEncoding":
        pos_enc_class_subsampling = pos_enc_class(odim, pos_dropout_rate)
    else:
        pos_enc_class_subsampling = None

    if input_type == "linear":
        return (
            torch.nn.Sequential(
                torch.nn.Linear(idim, odim),
                torch.nn.LayerNorm(odim),
                torch.nn.Dropout(dropout_rate),
                torch.nn.ReLU(),
                pos_enc_class(odim, pos_dropout_rate),
            ),
            1,
        )
    elif input_type == "conv2d":
        return Conv2dSubsampling(idim, odim, dropout_rate, pos_enc_class_subsampling), 4
    elif input_type == "vgg2l":
        return VGG2L(idim, odim, pos_enc_class_subsampling), 4
    elif input_type == "embed":
        return (
            torch.nn.Sequential(
                torch.nn.Embedding(idim, odim, padding_idx=padding_idx),
                pos_enc_class(odim, pos_dropout_rate),
            ),
            1,
        )
    elif input_type == "c-embed":
        return (
            torch.nn.Sequential(
                torch.nn.Embedding(idim, odim, padding_idx=padding_idx),
                torch.nn.Dropout(dropout_rate),
            ),
            1,
        )
    else:
        raise NotImplementedError(
            "Invalid input layer: %s. Supported: linear, conv2d, vgg2l and embed"
            % input_type
        )

def build_transformer_block(
    net_part: str,
    block: Dict[str, Any],
    pw_layer_type: str,
    pw_activation_type: str,
) -> Union[EncoderLayer, TransformerDecoderLayer]:
    """Build function for transformer block.

    Args:
        net_part: Network part, either 'encoder' or 'decoder'.
        block: Transformer block parameters.
        pw_layer_type: Positionwise layer type.
        pw_activation_type: Positionwise activation type.

    Returns:
        : Function to create transformer (encoder or decoder) block.

    """
    d_hidden = block["d_hidden"]

    dropout_rate = block.get("dropout-rate", 0.0)
    pos_dropout_rate = block.get("pos-dropout-rate", 0.0)
    att_dropout_rate = block.get("att-dropout-rate", 0.0)

    if pw_layer_type != "linear":
        raise NotImplementedError(
            "Transformer block only supports linear pointwise layer."
        )

    if net_part == "encoder":
        transformer_layer_class = EncoderLayer
    elif net_part == "decoder":
        transformer_layer_class = TransformerDecoderLayer

    return lambda: transformer_layer_class(
        d_hidden,
        MultiHeadedAttention(block["heads"], d_hidden, att_dropout_rate),
        PositionwiseFeedForward(
            d_hidden,
            block["d_ff"],
            pos_dropout_rate,
            get_activation(pw_activation_type),
        ),
        dropout_rate,
    )

def build_conformer_block(
    block: Dict[str, Any],
    self_attn_class: str,
    pw_layer_type: str,
    pw_activation_type: str,
    conv_mod_activation_type: str,
) -> ConformerEncoderLayer:
    """Build function for conformer block.

    Args:
        block: Conformer block parameters.
        self_attn_type: Self-attention module type.
        pw_layer_type: Positionwise layer type.
        pw_activation_type: Positionwise activation type.
        conv_mod_activation_type: Convolutional module activation type.

    Returns:
        : Function to create conformer (encoder) block.

    """
    d_hidden = block["d_hidden"]
    d_ff = block["d_ff"]

    dropout_rate = block.get("dropout-rate", 0.0)
    pos_dropout_rate = block.get("pos-dropout-rate", 0.0)
    att_dropout_rate = block.get("att-dropout-rate", 0.0)

    macaron_style = block["macaron_style"]
    use_conv_mod = block["use_conv_mod"]

    if pw_layer_type == "linear":
        pw_layer = PositionwiseFeedForward
        pw_layer_args = (
            d_hidden,
            d_ff,
            pos_dropout_rate,
            get_activation(pw_activation_type),
        )
    else:
        raise NotImplementedError("Conformer block only supports linear yet.")

    if macaron_style:
        macaron_net = PositionwiseFeedForward
        macaron_net_args = (
            d_hidden,
            d_ff,
            pos_dropout_rate,
            get_activation(pw_activation_type),
        )

    if use_conv_mod:
        conv_mod = ConvolutionModule
        conv_mod_args = (
            d_hidden,
            block["conv_mod_kernel"],
            get_activation(conv_mod_activation_type),
        )

    return lambda: ConformerEncoderLayer(
        d_hidden,
        self_attn_class(block["heads"], d_hidden, att_dropout_rate),
        pw_layer(*pw_layer_args),
        macaron_net(*macaron_net_args) if macaron_style else None,
        conv_mod(*conv_mod_args) if use_conv_mod else None,
        dropout_rate,
    )

def build_conv1d_block(block: Dict[str, Any], block_type: str) -> CausalConv1d:
    """Build function for causal conv1d block.

    Args:
        block: CausalConv1d or Conv1D block parameters.

    Returns:
        : Function to create conv1d (encoder) or causal conv1d (decoder) block.

    """
    if block_type == "conv1d":
        conv_class = Conv1d
    else:
        conv_class = CausalConv1d

    stride = block.get("stride", 1)
    dilation = block.get("dilation", 1)
    groups = block.get("groups", 1)
    bias = block.get("bias", True)

    use_batch_norm = block.get("use-batch-norm", False)
    use_relu = block.get("use-relu", False)
    dropout_rate = block.get("dropout-rate", 0.0)

    return lambda: conv_class(
        block["idim"],
        block["odim"],
        block["kernel_size"],
        stride=stride,
        dilation=dilation,
        groups=groups,
        bias=bias,
        relu=use_relu,
        batch_norm=use_batch_norm,
        dropout_rate=dropout_rate,
    )

def build_blocks(
    net_part: str,
    idim: int,
    input_layer_type: str,
    blocks: List[Dict[str, Any]],
    repeat_block: int = 0,
    self_attn_type: str = "self_attn",
    positional_encoding_type: str = "abs_pos",
    positionwise_layer_type: str = "linear",
    positionwise_activation_type: str = "relu",
    conv_mod_activation_type: str = "relu",
    input_layer_dropout_rate: float = 0.0,
    input_layer_pos_enc_dropout_rate: float = 0.0,
    padding_idx: int = -1,
) -> Tuple[
    Union[Conv2dSubsampling, VGG2L, torch.nn.Sequential], MultiSequential, int, int
]:
    """Build custom model blocks.

    Args:
        net_part: Network part, either 'encoder' or 'decoder'.
        idim: Input dimension.
        input_layer: Input layer type.
        blocks: Blocks parameters for network part.
        repeat_block: Number of times provided blocks are repeated.
        positional_encoding_type: Positional encoding layer type.
        positionwise_layer_type: Positionwise layer type.
        positionwise_activation_type: Positionwise activation type.
        conv_mod_activation_type: Convolutional module activation type.
        input_layer_dropout_rate: Dropout rate for input layer.
        input_layer_pos_enc_dropout_rate: Dropout rate for input layer pos. enc.
        padding_idx: Padding symbol ID for embedding layer.

    Returns:
        in_layer: Input layer
        all_blocks: Encoder/Decoder network.
        out_dim: Network output dimension.
        conv_subsampling_factor: Subsampling factor in frontend CNN.

    """
    fn_modules = []

    pos_enc_class, self_attn_class = get_pos_enc_and_att_class(
        net_part, positional_encoding_type, self_attn_type
    )

    input_block = prepare_input_layer(
        input_layer_type,
        idim,
        blocks,
        input_layer_dropout_rate,
        input_layer_pos_enc_dropout_rate,
    )

    out_dim = prepare_body_model(net_part, blocks)

    input_layer, conv_subsampling_factor = build_input_layer(
        input_block,
        pos_enc_class,
        padding_idx,
    )

    for i in range(len(blocks)):
        block_type = blocks[i]["type"]

        if block_type in ("causal-conv1d", "conv1d"):
            module = build_conv1d_block(blocks[i], block_type)
        elif block_type == "conformer":
            module = build_conformer_block(
                blocks[i],
                self_attn_class,
                positionwise_layer_type,
                positionwise_activation_type,
                conv_mod_activation_type,
            )
        elif block_type == "transformer":
            module = build_transformer_block(
                net_part,
                blocks[i],
                positionwise_layer_type,
                positionwise_activation_type,
            )

        fn_modules.append(module)

    if repeat_block > 1:
        fn_modules = fn_modules * repeat_block

    return (
        input_layer,
        MultiSequential(*[fn() for fn in fn_modules]),
        out_dim,
        conv_subsampling_factor,
    )


from bson import SON

RESOURCE_COLLECTION = 'competences'

_schema = {
    'id': {'type': 'integer',
           'required': True,
           'unique': True},
    'approved_by_org_id': {'type': 'integer'},
    'approved_by_person_id': {'type': 'integer'},
    'course_id': {'type': 'integer'},
    'passed': {'type': 'boolean'},
    'person_id': {'type': 'integer'},
    'sald': {'type': 'integer'},
    'title': {'type': 'string'},
    'type_id': {'type': 'integer'},
    'additional_title': {'type': 'string'},
    'paid_date': {'type': 'datetime'},
    'date': {'type': 'datetime'},
    'valid_until': {'type': 'datetime'},
    'grade': {'type': 'string'},
    'comments': {'type': 'string'},
    '_section': {'type': 'string'},
    '_code': {'type': 'string'},
    '_name': {'type': 'string'}
}

definition = {
    'url': 'competences',
    'item_title': 'Competences',
    'datasource': {'source': RESOURCE_COLLECTION,
                   },
    'additional_lookup': {
        'url': 'regex("[\d{1,9}]+")',
        'field': 'id',
    },
    'extra_response_fields': ['id'],
    'versioning': False,
    'resource_methods': ['GET'],
    'item_methods': ['GET'],
    'mongo_indexes': {'competence_id': ([('id', 1)], {'background': True}),
                      'person_id': ([('person_id', 1)], {'background': True}),
                      'type_id': ([('type_id', 1)], {'background': True}),
                      'approved': ([('approved_by_person_id', 1)], {'background': True}),
                      'valid_until': ([('valid_until', 1)], {'background': True}),
                      'section': ([('_section', 1)], {'background': True}),
                      'code': ([('_code', 1)], {'background': True}),
                      'title': ([('title', 'text')], {'background': True})
                      },
    'schema': _schema
}

# Process resource without data_relations
_schema_process = _schema.copy()
_schema_process['active_in_org_id'] = {'type': 'integer'}
_schema_process['person_id'] = {'type': 'integer'}

process_definition = {
    'url': 'competences/process',
    'item_title': 'Competences Process',
    'datasource': {'source': RESOURCE_COLLECTION,
                   },
    'additional_lookup': {
        'url': 'regex("[\d{1,9}]+")',
        'field': 'id',
    },
    'pagination': False,
    'extra_response_fields': ['id'],
    'versioning': False,
    'resource_methods': ['GET', 'POST', 'DELETE'],
    'item_methods': ['GET', 'PATCH', 'PUT'],
    'schema': _schema_process
}

# Aggregations
agg_count_codes = {
    'url': 'competences/codes',
    'item_title': 'Competences Codes',
    'pagination': False,
    'datasource': {
        'source': RESOURCE_COLLECTION,
        'aggregation': {
            'pipeline': [
                {"$group": {"_id": "$title", "count": {"$sum": 1}}},
                {"$sort": SON([("count", -1), ("_id", -1)])}
            ]
        }
    }
}

# Aggregations
agg_count_types = {
    'url': 'competences/types/count',
    'item_title': 'Competences Types',
    'pagination': False,
    'datasource': {
        'source': RESOURCE_COLLECTION,
        'aggregation': {
            'pipeline': [
                {"$group": {"_id": "$type_id", "count": {"$sum": 1}}},
                {"$sort": SON([("count", -1), ("_id", -1)])}
            ]
        }
    }
}

# Aggregations
agg_count_persons = {
    'url': 'competences/persons/count',
    'item_title': 'Competences persons',
    'pagination': False,
    'datasource': {
        'source': RESOURCE_COLLECTION,
        'aggregation': {
            'pipeline': [
                {"$group": {"_id": "$person_id", "count": {"$sum": 1}}},
                {"$sort": SON([("count", -1), ("_id", -1)])}
            ]
        }
    }
}


import  pygame
from Settings import  Setting
from Ship import Ship
import  game_function as gf

def runGame():

    # Initialize game and create a screen object.

    pygame.init()

    gameSetting = Setting()

    screen = pygame.display.set_mode((gameSetting.screen_width , gameSetting.screen_height))

    pygame.display.set_caption("Alien Invasion")

    ship = Ship( gameSetting,screen)

    # Start the main loop for the game.

    while True:

    # Watch for keyboard and mouse events.

     gf.check_event(ship)
     ship.update()
     gf.update_screen(gameSetting , screen , ship)

runGame()


# -*- coding: utf-8 -*-
"""
Created on Wed Jan 24 10:31:40 2018

@author: zhouying
"""

import numpy as np
from sklearn.neighbors import NearestNeighbors
import scipy.io as scio

mydata = scio.loadmat('MNIST_data\\UCI\\sonar.mat')
data = np.array(mydata['data'])
label = np.transpose(mydata['label'])
label = label[:,0]

nbrs = NearestNeighbors(n_neighbors=6, algorithm="auto").fit(data)
_,indices = nbrs.kneighbors(data)

positvie =1
negative = 2

nn = indices[:,1:]
Tplas = 0.
Tminas=0.
for i in range(label.shape[0]):
    if label[i] == positvie :
        for value in label[nn[i]]:
            if value == positvie:
                Tplas+=1
    elif label[i]==negative:
        for value in label[nn[i]]:
            if value == negative:
                Tminas+=1

Tplas/=label[label==positvie].shape[0]*nn.shape[1]
Tminas/=label[label==negative].shape[0]*nn.shape[1]

print(Tplas,Tminas,Tminas-Tplas)


import tensorflow as tf
import numpy as np
import os
import time
import cv2
import numpy as np
from tqdm import tqdm

def load_graph(frozen_graph_filename):
    with tf.gfile.GFile(frozen_graph_filename, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(
            graph_def,
            input_map=None,
            return_elements=None,
            name="prefix",
            op_dict=None,
            producer_op_list=None
        )

    input_name = graph.get_operations()[0].name+':0'
    output_name = graph.get_operations()[-1].name+':0'
    return graph, input_name, output_name

def prepare_data_x(img_dir,label=False,invert_color=False):
    inicio = time.time()
    tipos = ['cruzamento', 'curva', 'parada', 'reta', 'zigzag']
    aux = []
#     resultado = []

#     print(img_dir)
    frame = cv2.imread(img_dir,0)
    img = frame
#     img = cv2.resize(frame, (150,150))
    if invert_color:
        img = (255-img)
    aux.append(img)
    img = np.array(aux)
    img = img/255
    img = img.reshape(img.shape[0], img.shape[1], img.shape[2], 1).astype('float32')
#     print("tempo antes pred:",time.time()-inicio)
#     resultado.append(img)
#     return np.array(resultado)
    return img
#     pred = model.predict(img)
#     print("tempo pos pred:",time.time()-inicio)
#     pred = np.around(pred, decimals = 2, out = None)
#     print("tempo TOTAL:",time.time()-inicio)
#     return tipos[np.argmax(pred)],pred,frame

def predict(model_path,img_path):
    # load tf graph
    tf_model,tf_input,tf_output = load_graph(model_path)
    print(tf_input,tf_output)
    # Create tensors for model input and output
    x = tf_model.get_tensor_by_name(tf_input)
    y = tf_model.get_tensor_by_name(tf_output)
    num_outputs = y.shape.as_list()[0]

    fotos = (os.listdir(folder))
    predictions = np.zeros((len(fotos),len(tipos)))

    with tf.Session(graph=tf_model) as sess:
        for foto in tqdm(fotos):
            tempo1 = time.time()
            y_out = sess.run(y, feed_dict={x: prepare_data_x(folder+foto)})
            print("tempo pred:",time.time()-tempo1)
            predictions = y_out[0]
            pred = np.around(predictions, decimals = 2, out = None)
            print(tipos[np.argmax(pred)])
        #     return tipos[np.argmax(pred)],pred,input_data
#     return predictions

if __name__ == '__main__':
    model_path = "/home/junior/Documents/seguidor_linha/modelo_9935.pb"
    load_graph(model_path)
    predict(model_path,)

    main()


import xml.etree.ElementTree as ET
import nltk
from nltk.tokenize import sent_tokenize, word_tokenize
from nltk.tag.stanford import StanfordPOSTagger
from optparse import OptionParser
#from preprocesstwitter import tokenize
import utils
import random
import os.path

usage = "BioScopeReader.py --dataset bioscope"

parser = OptionParser(usage=usage)
parser.add_option("--dataset", type="string", help="dataset", default="bioscope", dest="dataset")
parser.add_option("--cuetype", type="string", help="cuetype", default="negation", dest="cuetype")
#
#parser.add_option("--outputpath", type="string", help="outputpath", default="", dest="negation")

(options, args) = parser.parse_args()

BioScope_filenames = ['full_papers.xml', 'clinical_records.xml', 'abstracts.xml']#'

CNeSP_filenames = ['Financial_article.xml', 'Scientific_literature.xml', 'Product_review.xml.txt' ]

filenames = {'bioscope':BioScope_filenames, 'cnesp':CNeSP_filenames}

langs = {'bioscope':'en', 'cnesp':'cn'}

def pos_tag(sentence, lang, tool):
    if tool == 'stanford':
        return pos_tag_stanford(sentence, lang)
    elif tool == 'nltk':
        return pos_tag_nltk(sentence, lang)
    else:
        return None

def pos_tag_nltk(sentence, lang):
    return None

def pos_tag_stanford(sentence, lang):
    path_to_model = "stanford-postagger-2015-12-09/models/english-bidirectional-distsim.tagger"
    path_to_jar = "stanford-postagger-2015-12-09/stanford-postagger.jar"
    tagger = StanfordPOSTagger(path_to_model, path_to_jar)
    tagger.java_options = '-mx4096m'  ### Setting higher memory limit for long sentences
    sentence = 'This is testing'
    print(tagger.tag(sentence.split()))

def token_count(text, lang = 'en'):
    if lang == 'en':
        textList = word_tokenize(text.strip())
        return len(textList)
    else:
        return None

for filename in filenames[options.dataset]:
    print('Loading ', filename, ' ... ')

    tree = ET.parse(filename)
    root = tree.getroot()

    def xscope2str(scope, xscopeList, pos, depth):
        ret = ''
        xscope = [str(pos), '']
        if scope.text != None:
            ret += scope.text
            pos += token_count(scope.text)

        # print('after scope.text',depth)
        for item in scope:
            if item.tag == 'cue':

                cueType = item.get('type')

                if item.text != None:
                    fromIdx = str(pos)
                    ret += item.text
                    pos += token_count(item.text)
                    toIdx = str(pos)

                    if cueType == options.cuetype:
                        xscope += [cueType, fromIdx, toIdx]

                if item.tail != None:
                    ret += item.tail
                    pos += token_count(item.tail)

            elif item.tag == 'xcope':

                ret0 = xscope2str(item, xscopeList, pos, depth + 1)
                ret += ret0
                pos += token_count(ret0)

                if item.tail != None:
                    ret += item.tail
                    pos += token_count(item.tail)

            # if scope.tail != None:
            #                 textList = word_tokenize(scope.tail)
            #                 ret += scope.tail
            #                 pos += len(textList)
            #                 print('scope.tail:',scope.tail)

            else:
                print('ERROR!!!!!!!')
                exit()

        xscope[1] = str(pos)

        if len(xscope) > 2:
            print(xscope, 'depth=', depth)
            xscopeList.append(xscope)

        return ret

    def Inst2Str(inst):
        tokens, postags, xscopeList = inst
        ret = ' '.join(tokens) + ' ||| ' + ' '.join(postags) + ' ||| '
        ret += ' ||| '.join([' '.join(xscope) for xscope in xscopeList])
        return ret

    def writeInstance(instances, path, filename, outputtag=True, encoding='utf-8'):

        filename =  os.path.join(path, filename)
        with open(filename, 'w', encoding=encoding) as f:
            for inst in instances:
                f.write(Inst2Str(inst))
                f.write('\n')

            f.close()

        print('|\t', filename + '    ', '\t|\t', len(instances), '\t|')

    Instances = []

    for A in root:
        for DS in A:
            for D in DS:
                for sent in D:
                    print()
                    print(sent.get('id'))
                    # sentInst = xmlsent2seq(sent)

                    ret = []
                    pos = 0
                    xscopeList = []

                    print('len(sent):', len(sent))

                    if sent.text != None:
                        textList = word_tokenize(sent.text)
                        ret += textList
                        pos += len(textList)

                    for i in range(0, len(sent)):
                        scope = sent[i]

                        if scope.text != None:
                            textList = word_tokenize(scope.text)
                            ret += textList
                            pos += len(textList)

                        ret0 = xscope2str(scope, xscopeList, pos, 1)
                        textList = word_tokenize(ret0)
                        ret += textList
                        pos += len(textList)

                        if scope.tail != None:
                            textList = word_tokenize(scope.tail)
                            ret += textList
                            pos += len(textList)

                    print([ret[i] + '/' + str(i) for i in range(0, len(ret))])
                    print('#scope:', len(xscopeList), '\t', xscopeList)

                    retPOSTags = nltk.pos_tag(ret)
                    print(retPOSTags)
                    tokens, postags = zip(*retPOSTags)

                    if len(ret) != len(tokens):
                        print('What a fuck!!!!')
                        print(ret)
                        print(tokens)
                        exit(-1)

                    if len(xscopeList) > 0:
                        Instances.append((tokens, postags, xscopeList))
                        # f.write(' '.join(tokens))
                        # f.write(' ||| ' + ' '.join(postags))
                        #
                        # f.write(' ||| ' + ' ||| '.join([' '.join(xscope) for xscope in xscopeList]))
                        # f.write('\n')

    #random.shuffle(Instances)

    size = len(Instances)
    random.shuffle(Instances)
    outputPath = options.cuetype
    os.makedirs(outputPath, exist_ok=True)
    writeInstance(Instances, outputPath, filename.replace('xml', 'txt'))

    if filename.startswith('abstracts'):
        numFold = 10
        folds_size = size // numFold
        folds = [[] for i in range(0, numFold)]
        lb = 0
        for i in range(0, numFold):
            rb = size if i == numFold - 1 else lb + folds_size
            folds[i] = Instances[lb: rb]
            lb = rb

        for i in range(0, numFold):

            trainInstances = []
            for j in range(0, numFold):
                if j != i:
                    trainInstances += folds[j]

            testInstances = folds[i]

            writeInstance(trainInstances, outputPath, filename.replace('xml', 'train.' + str(i) +'.txt'))
            writeInstance(testInstances, outputPath, filename.replace('xml', 'test.' + str(i) + '.txt'))

    # lb = 0
    # rb = int(size * 0.7)
    # trainInstances = Instances[lb: rb]
    #
    # lb = rb + 1
    # rb = int(size * 0.85)
    # devInstances = Instances[lb: rb]
    #
    # lb = rb + 1
    # rb = int(size * 1)
    # testInstances = Instances[lb:rb]
    #
    # outputPath = ''
    #
    # writeInstance(trainInstances, outputPath, filename.replace('xml', 'train.txt'))
    # writeInstance(devInstances, outputPath, filename.replace('xml', 'dev.txt'))
    # writeInstance(testInstances, outputPath, filename.replace('xml', 'test.txt'))


import serial
import time
from serial.tools import list_ports
print("INFO: Connecting GPS Port..")

device_signature= '2c7c:0125'
candidates=list(list_ports.grep(device_signature))
candidates.reverse()

print([candidate.device for candidate in candidates])

try:
    print("writing to "+ candidates[2].device)
    serw = serial.Serial(candidates[2].device, baudrate = 115200, timeout = 1,rtscts=True, dsrdtr=True)
    serw.write('AT+QGPS=1\r'.encode())
    serw.close()
    time.sleep(1)
except Exception as e:
    print("ERROR: Serial port connection failed.")
    print(e)
exit(int(candidates[1].device[-1]))


import pickle
import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from torch.utils.data import Dataset

class CIFAR10(Dataset):

    def __init__(self,root,train=True,transforms=None):
        self.root = root
        self.transforms = transforms
        self.split = train

        self.data = []
        self.targets = []
        self.train_data = [file for file in os.listdir(root) if "data_batch" in file]
        self.test_data = [file for file in os.listdir(root) if "test_batch" in file]

        data_split = self.train_data if self.split else self.test_data

        for files in data_split:
            entry = self.extract(os.path.join(root,files))
            self.data.append(entry["data"])
            self.targets.extend(entry["labels"])

        self.data = np.vstack(self.data).reshape(-1, 3, 32, 32)
        self.data = self.data.transpose((0, 2, 3, 1))
        self.load_meta()

    def extract(self,filename):
        with open(filename,"rb") as f:
            batch_data = pickle.load(f,encoding="latin1")
        return batch_data

    def load_meta(self):
        path = os.path.join(self.root,"batches.meta")
        with open(path,"rb") as infile:
            data = pickle.load(infile,encoding="latin1")
            self.classes = data["label_names"]
            self.classes_to_idx = {_class:i for i,_class in enumerate(self.classes)}

    def plot(self,image,target=None):
        if target is not None:
            print(f"Target :{target} class :{self.classes[target]}")
        plt.figure(figsize=(2,2))
        plt.imshow(image.permute(1,2,0))
        plt.show()

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

    def __getitem__(self,idx):
        image,target = self.data[idx],self.targets[idx]
        image = Image.fromarray(image)

        if self.transforms:
            image = self.transforms(image)

        return image,target


# spherical Gaussian (each dim os uncorrelated and independent to the other)
from scipy.stats import norm
import numpy as np
import matplotlib.pyplot as plt

r = np.random.randn(10000, 2)

# spread out in a circular cloud
plt.scatter(r[:,0], r[:,1])
plt.show()

# elliptical Gaussian
r[:,1] = 5*r[:,1] + 2
plt.scatter(r[:,0], r[:,1])
plt.axis('equal') # for data with each axis equal together
plt.show()


from __future__ import annotations

from typing import TYPE_CHECKING

import quarrel
from src.bot import bot

from .interaction import MockInteraction

__all__ = ("mock", "build_mock", "build_application_command_mock")

if TYPE_CHECKING:
    from typing import Any

    from quarrel import Missing
    from quarrel.types.interactions import (
        Interaction,
        InteractionData,
        InteractionDataApplicationCommandOption,
    )
    from quarrel.types.member import MemberWithUser
    from quarrel.types.message import Message
    from quarrel.types.snowflake import Snowflake
    from quarrel.types.user import User

def mock(data: Interaction) -> MockInteraction:
    return MockInteraction(data, bot.state)

def build_mock(
    type: quarrel.InteractionType,
    *,
    data: InteractionData,
    guild_id: Missing[Snowflake] = quarrel.MISSING,
    channel_id: Missing[Snowflake] = quarrel.MISSING,
    member: Missing[MemberWithUser] = quarrel.MISSING,
    user: Missing[User] = quarrel.MISSING,
    message: Missing[Message] = quarrel.MISSING,
) -> MockInteraction:
    payload: Interaction = {
        "id": "0",
        "application_id": "1",
        "type": type.value,
        "token": "token123",
        "version": 1,
        "guild_id": guild_id if guild_id is not quarrel.MISSING else DEFAULT_GUILD_ID,
        "channel_id": channel_id
        if channel_id is not quarrel.MISSING
        else DEFAULT_CHANNEL_ID,
        "member": member
        if member is not quarrel.MISSING
        else DEFAULT_INTERACTION_MEMBER,
    }
    if user is not quarrel.MISSING:
        payload["user"] = user
    if message is not quarrel.MISSING:
        payload["message"] = message
    payload["data"] = data
    return mock(payload)

def build_application_command_mock(
    *,
    guild_id: Missing[Snowflake] = quarrel.MISSING,
    channel_id: Missing[Snowflake] = quarrel.MISSING,
    member: Missing[MemberWithUser] = quarrel.MISSING,
    user: Missing[User] = quarrel.MISSING,
    message: Missing[Message] = quarrel.MISSING,
) -> MockInteraction:
    return build_mock(
        quarrel.InteractionType.APPLICATION_COMMAND,
        data={},  # type: ignore
        guild_id=guild_id,
        channel_id=channel_id,
        member=member,
        user=user,
        message=message,
    )

def command_to_options(
    command: quarrel.SlashCommand[Any], **options: Any
) -> list[InteractionDataApplicationCommandOption]:
    payload = command.to_payload().get("options", [])
    for i in payload:
        i["value"] = options[i["name"]]  # type: ignore
    return payload  # type: ignore

DEFAULT_GUILD_ID = "1234"
DEFAULT_CHANNEL_ID = "5678"
DEFAULT_INTERACTION_MEMBER: MemberWithUser = {
    "roles": [],
    "joined_at": "2020-01-01T00:00:00.000Z",
    "deaf": False,
    "mute": False,
    "nick": None,
    "avatar": None,
    "user": {
        "id": "1234",
        "username": "user",
        "discriminator": "1234",
        "avatar": None,
        "bot": False,
        "system": False,
    },
}


from asgiref.sync import sync_to_async

from scts.tracking.domain.models import TrackingEvent
from scts.tracking.domain.repositories import TrackingEventRepository

class FakeTrackingEventRepository(TrackingEventRepository):

    def __init__(self) -> None:
        self.tracking_events = []

    @sync_to_async
    def save(self, trackingEvent: TrackingEvent):
        self.tracking_events.append(
            {
                "code": trackingEvent.code,
                "date": trackingEvent.date,
                "description": trackingEvent.description,
                "location": trackingEvent.location,
            }
        )