from __future__ import division
# <NAME>!
__author__ = '<NAME>'

# tpcf(dat, datR=None, randcatsize=2, bins,**kwargs)
#
# **kwargs for choosing geometry - metric 'flat' 'open' 'close'
# **kwargs for choosing xi estimator - 'simple' 'ls' '...'
# import fileios
# from metrics.metrics import * # for pycharm remove comment - for pip installation etc. comment this line
from tqdm import *
from datprep import *
import numpy as np
from metrics import *
from multiprocessing import Process
from multiprocessing.queues import Queue
from sklearn.neighbors import BallTree
from scipy.spatial import distance as dist
import matplotlib.pyplot as plt

# pcpus = cpu_count()-1

def tpcf(datfile, bins, **kwargs):
    """Main function to calculate 2pCF. Takes multiple arguments such as randfile, maskfile, calculation estimator etc. for different geometry, cosmology models
    Usage of the package is given in jupyter notebook "Using correlcalc example.nb" and in `main.py`

    All the methods in correlcalc can be imported using the following command

    `from correlcalc import *`

    We first need to define bins (in $c/H_0$ units) to calculate 2pCF. For e.g. to calculate correlation between 0-180Mpc in steps of 6Mpc, we say

    `bins=np.arange(0.002,0.06,0.002)`

    To calculate 2pCF using input data file (both ascii and fits files are supported), use `tpcf` method as follows

    `correl, poserr=tpcf('/path/to/datfile.dat',bins, randfile='/path/to/randomfile.dat', weights='eq')`

    If random file is not available or not provided, we can generate random catalog by providing the mangle mask file in `.ply` format along with specifying the size of the catalog in multiples of size of data catalog (default 2x size). To do this

    `correl, poserr=tpcf('/path/to/datfile.dat', bins, maskfile='/path/to/maskfile.ply', weights=True, randfact=3)`

    This returns `correl` and `poserr` `numpy` arrays corresponding to Two-point correlation and Poisson error

    ### Keyword Arguments
    The following keyword arguments can be included as needed

    #### Data file (Mandatory)

    Data file of galaxy/quasar redshift survey must be passed as the first argument to both `tpcf` and `atpcf` methods.

    **Supported filetypes**: ascii text files with columns, csv files or fits files are all supported. Most files provided by SDSS Value added catalogs should be directly usable.

     **To contain**: Any type of file provided must at least have columns named **Z** (redshift), **RA** (Right Ascension), **DEC** (Declination). These column names can be in any case.

     If one intends to use `weights=True` option (must to obtain accurate results) the data file must also contain radial weights with column title **radial_weight** or **WEIGHT_SYSTOT**

    #### bins (Mandatory)

    A numpy array with ascending values in $c/H_0$ units must be provided as the second argument to both `tpcf` and `atpcf` methods. In case of `atpcf` it automatically creates 2D bins as `bins2d=(bins,bins)` from provided 1D `bins`

    #### `randfile=` Path to random file (semi-Optional)

    If not provided, `maskfile=` argument must be given `.ply` file.

    **Supported filetypes**: ascii text files with columns, csv files or fits files are all supported. Most files provided by SDSS Value added catalogs should be directly usable.

     **To contain**: Any type of file provided must at least have columns named **Z** (redshift), **RA** (Right Ascension), **DEC** (Declination). These column names can be in any case.

     If one intends to use `weights=True` option (must to obtain accurate results) the data file must also contain radial weights with column title **radial_weight** or **WEIGHT_SYSTOT**

     **Beta Testing:** Beta support for other column titles for weights is added.

     Also added is calculation of weights from n(z) during random catalog generation.

    #### `mask=` Path to mangle polygon file (semi-Optional)

    If not provided, `randfile=` argument must be provided.

    **Supported filetypes**: `.ply` file containing Mangle polygons describing survey geometry in the standard format. Most files provided by SDSS Value added catalogs should be directly usable.

    #### `randfact=` (Optional)

    Size of the random catalog in integer multiples of size of data catalog if random catalog file is not provided. Default value is `2`

    #### `weights=` (Optional)

    It is highly recommended to use weights argument by providing `weights=True` or `weights='eq'` to obtain accurate two-point correlation calculations. This picks up radial weights in the prescribed format (with column title **radial_weight** or **WEIGHT_SYSTOT** ) from the data and random files provided.

    `weights=`eq'` sets equal weights and hence adds *+1* - This implementation is parallelized and is faster than `weights=False` implementation on most machines

    If `weights=False`, by default *+1* will be added for each galaxy/random pair found within the bin instead of adding total weight. For more details on weights and references, see http://www.sdss3.org/dr9/tutorials/lss_galaxy.php

    #### `geometry='flat'` (Optional)

    **Available options**:

    `'flat'`(default) - for flat geometry of the Universe

     `'open'` - for Open Universe models like Milne

     `'close'` - for Closed Universe

     **Customization**

     Formulae for calculation of distances between two points (Z1, RA1, DEC1) and (Z2, RA2, DEC2) is taken from *<NAME>, Correlation function in deep redshift space as a cosmological probe, The Astrophysical Journal 615 (2) (2004) 573*. Using the formulae in this paper, distances squares (to reduce additional computational time distance squares are calculated to avoid using expensive `sqrt` function every time) are computed in the `metrics.pyx` file for all the above mentioned geometries. `Cython` is chosen for implementation to obtain faster results in building `BallTree`s  calculating `cdist` and to reduce `query` time.

     One can customize metric definitions as per one's need by editing this file. Also **K** (curvature parameter) in the formulae given in this reference need to be manually changed in the `metrics.pyx` for closed and open cases as per the model. After changing this compile it using `python metricsetup.py build_ext --inplace`

    #### `cosmology='lcdm'` (Optional)

    Used to calculate co-moving distances from redshifts.

    **Available options**:

    `'lcdm'` (default)- for Lambda CDM model

     `'lc'` - for $R_h=ct$ and linear coasting models

     **To add**: `wcdm` and other popular cosmology models soon

    #### `estimator=` (Optional)

    **Available options**:

    `'dp'` - Davis - Peebles estimator (default - fastest)

    `'ls'`- Landy - Szalay estimator

    `'ph'` - Peebles- Hauser estimator

    `'hew'` - Hewitt estimator

    `'h'` - Hamilton estimator

    For more details on estimator formulae see https://arxiv.org/pdf/1211.6211.pdf
    """
    # Default function arguments
    # weights = np.array([])
    global dat
    global datR
    DD = DR = RD = RR = np.zeros(len(bins)-1)
    weightsflag = True
    useones = True
    cosmology = 'lcdm'
    geometry = 'flat'
    metric = flatdistsq
    randcatfact = 2
    estimator = 'dp'
    binsq = bins**2
    randfile = None
    maskfile = None

    # Options for correl calculation estimators and cosmology models
    elist = ['dp', 'ls', 'ph', 'hew', 'h']
    clist = ['lcdm', 'lc'] # add wcdm

    if kwargs is not None:
        for key, value in kwargs.items():
            # print (key, value)
            # print (key, value)
            if key.lower() == 'randfile':
                randfile = value
            elif key.lower() == 'randfact':
                randcatfact = value
            elif key.lower() == 'geometry':
                if value.lower() == 'flat':
                    geometry = 'flat'
                    metric = flatdistsq
                elif value.lower() == 'open':
                    geometry = 'open'
                    metric = opendistsq
                elif value.lower() == 'close':
                    geometry = 'close'
                    metric = closedistsq
                else:
                    print("Incorrect geometry argument provided! Using flat geometry")
            elif key.lower() == 'estimator':
                if value.lower() in elist:
                    estimator = value.lower()
                    print estimator
                else:
                    print("Incorrect estimator provided! Using 'dp' as default")
            elif key.lower() == 'cosmology':
                if value.lower() in clist:
                    cosmology = value.lower()
                else:
                    print("Incorrect Cosmology provided! Using 'lcdm' as default")
            elif key.lower() == 'mask':
                maskfile = value
            elif key.lower() == 'weights':
                if value is True:
                    weightsflag = True
                    useones = False

                    # fdat=readinfile(datfile,ftype='internal')
                    # weights=1.0/(1.0+4.0*np.array(fdat['nz']))
                    # weights=weights/np.mean(weights)
                    # print (weights)
                elif isinstance(value, str):
                    if value.lower() == 'eq':
                        weightsflag = True
                        useones = True
                else:
                    weightsflag = False
            else:
                print ("key argument `%s` not valid" % key)
    else:
        print ("Refer documentation to enter valid keyword arguments")

    print("Calculating Correlation function with the following parameters")
    print ("data file=")
    print(datfile)
    print("random file=")
    print(randfile)
    print("Random catalog size factor(if random file is None)=")
    print(randcatfact)
    print("mask/window file=")
    print(maskfile)
    print ("Cosmology=")
    print(cosmology)
    print("Geometry=")
    print(geometry)
    print("Correl estimator=")
    print(estimator)
    print ("Weights=")
    print (weightsflag)
    print ("Using ones as weights?=")
    print (useones)
    print("-----------------------------------------")
    # Prepare dat from data file
    dat, weights = datprep(datfile, 'data', cosmology)
    global Nd
    Nd = len(dat)
    # Nd = len(dat)
    # print (weights)
    # Prepare datR from random file or generate a random catalog
    if randfile is None:
        randcatsize = randcatfact*Nd
        if maskfile is None:
            print ("Mask file compulsory. Please provide mask='maskfilepath.ply'")
        else:
            # rweights = np.array([])
            datR, rweights = randcatprep(datfile, randcatsize, maskfile, cosmology)
            # randfile='./randcat.dat'
            # datR, rweights=datprep(randfile,'random',cosmology)
    else:
        datR, rweights = datprep(randfile, 'random', cosmology)

    # if len(weights)!=0:
        # rfdat=readinfile(randfile,ftype='internal')
        # rweights=1.0/(1.0+4.0*np.array(rfdat['nz']))
        # rweights=rweights/np.mean(rweights)
        # print (rweights)
    # Nr=len(datR)

    global Nr
    Nr = len(datR)

    fact = (1.0*Nr)/Nd

    # Creating module-wise global balltrees so that they don't have to be created many times.

    global dbt
    global rbt

    print ("Creating BallTree for data points using metric=")
    print (metric)
    dbt = BallTree(dat, metric='pyfunc', func=metric)

    print ("Creating BallTree for random points using metric=")
    print (metric)
    rbt = BallTree(datR, metric='pyfunc', func=metric)

    print ("Calculating 2pCF...")
    # f=(1.0*Nrd)/N
    # print (weights)
    # Reference: arXiv: 1211.6211
    if estimator == 'dp':
        if weightsflag is False: # or len(weights) != Nd:
            # print (weightsflag)
            # print(len(weights))
            # print(len(datR))
            DD = DDcalc(dat, binsq)
            DR = DRcalc(dat, binsq)
            RD = RDcalc(datR, binsq)
        else:
            if useones is True or len(weights) != Nd:
                weights = np.ones(Nd)
                rweights = np.ones(Nr)
            # if len(rweights)!=len(datR):
            DD = DDwcalc(dat, binsq, metric, weights)
            DR = DRwcalc(dat, datR, binsq, metric, rweights)
            RD = RDwcalc(dat, datR, binsq, metric, weights)
            # else:
            #     DD=DDwcalc(dat,binsq,metric,weights)
            #     DR=DRwcalc(dat,datR,binsq,metric,rweights)
        print ("Using Davis-Peebles estimator")
        correl = fact*(DD*2.0/(DR+RD))-1.0

    elif estimator == 'ph':
        if weightsflag is False: # or len(weights) != Nd:
            DD = DDcalc(dat, binsq)
            RR = RRcalc(datR, binsq)
        else:
            if useones is True or len(weights) != Nd:
                weights = np.ones(Nd)
                rweights = np.ones(Nr)
            DD = DDwcalc(dat, binsq, metric, weights)
            if len(rweights) != Nr:
                RR = RRcalc(datR, binsq)
            else:
                RR = RRwcalc(datR, binsq, metric, rweights)
        print ("Using Peebles-Hauser estimator")
        correl = fact**2*(DD/RR)-1.0
    else:
        if weightsflag is False: # or len(weights) != Nd:
            DD = DDcalc(dat, binsq)
            RR = RRcalc(datR, binsq)
            DR = DRcalc(dat, binsq)
            RD = RDcalc(datR, binsq)
        else:
            if useones is True or len(weights) != Nd:
                weights = np.ones(Nd)
                rweights = np.ones(Nr)
            DD = DDwcalc(dat, binsq, metric, weights)
            DR = DRwcalc(dat, datR, binsq, metric, rweights)
            RD = RDwcalc(dat, datR, binsq, metric, weights)
            RR = RRwcalc(datR, binsq, metric, rweights)
        if estimator == 'ls':
            print ("Using Landy-Szalay estimator")
            # correl = (DD-2.0*DR+RR)/RR
            correl = fact**2*(DD/RR)-fact*(DR+RD)/RR+1.0
        elif estimator == 'hew':
            print ("Using Hewett estimator")
            correl = fact**2*(DD/RR)-fact*0.5*(DR+RD)/RR
        elif estimator == 'h':
            print ("Using Hamilton estimator")
            correl = (4.0*DD*RR)/(DR+RD)**2 - 1.0
    correlerr = poserr(correl, DD)
    print("Two-point correlation=")
    np.savetxt("DD_"+str(cosmology)+"_"+str(geometry)+"_"+str(estimator)+".txt", DD)
    np.savetxt("DR_"+str(cosmology)+"_"+str(geometry)+"_"+str(estimator)+".txt", DR)
    np.savetxt("RD_"+str(cosmology)+"_"+str(geometry)+"_"+str(estimator)+".txt", RD)
    np.savetxt("RR_"+str(cosmology)+"_"+str(geometry)+"_"+str(estimator)+".txt", RR)
    np.savetxt("bins_"+str(cosmology)+"_"+str(geometry)+"_"+str(estimator)+".txt", bins)
    np.savetxt("tpcf_"+str(cosmology)+"_"+str(geometry)+"_"+str(estimator)+".txt", (correl, correlerr))
    print (correl, correlerr)
    return correl, correlerr

def DDcalc(dat, bins):
    print ("Calculating DD...\n DD=")
    DD = autocorr(dat, bins)
    DD[DD == 0] = 1.0
    # Nd = len(dat)
    # DD = DD/(Nd*(Nd-1.0))
    print (DD)
    return DD

def RRcalc(datR, bins):
    print ("Calculating RR...\n RR=")
    RR = rautocorr(datR, bins)
    RR[RR == 0] = 1.0
    # Nr = len(datR)
    # RR = RR/(Nr*(Nr-1.0))
    print (RR)
    return RR

def DRcalc(dat, bins):
    print ("Calculating DR...\n DR=")
    DR = crosscorr(dat, bins)
    DR[DR == 0] = 1.0
    # Nd = len(dat)
    # Nr = len(datR)
    # DR = DR/(Nd*Nr)
    print (DR/2.0)
    return DR/2.0

def RDcalc(datR, bins):
    print ("Calculating RD...\n RD=")
    RD = crosscorrd(datR, bins)
    RD[RD == 0] = 1.0
    # Nd = len(dat)
    # Nr = len(datR)
    # DR = DR/(Nd*Nr)
    print (RD/2.0)
    return RD/2.0

def autocorr(dat, bins):
    counts_DD = dbt.two_point_correlation(dat, bins)
    DD = np.diff(counts_DD)
    return DD

# def autocorrp(dat, bins):
#     pool = Pool(processes=pcpus)
#     counts_DD = pool.map(dbt.two_point_correlation, (dat, bins))
#     DD = np.diff(counts_DD)
#     return DD

def rautocorr(datR, bins):
    counts_RR = rbt.two_point_correlation(datR, bins)
    RR = np.diff(counts_RR)
    return RR

def crosscorr(dat, bins):
    counts_DR = rbt.two_point_correlation(dat, bins)
    DR = np.diff(counts_DR)
    return 2.0*DR

def crosscorrd(datR, bins):
    counts_RD = dbt.two_point_correlation(datR, bins)
    RD = np.diff(counts_RD)
    return 2.0*RD

def poserr(xi, DD):
    print ("Calculating Poisson error")
    return (1.0+xi)/np.sqrt(DD)
# alternatively
# rbt=BallTree(dat,metric='pyfunc',func=metric)
# counts_RD=rbt.two_point_correlation(dat,bins)

def DDwcalc(dat, bins, metric, weights):
    print ("Calculating DD with weights (parallelized)...\n DD=")
    # DD = autocorrw(dat, bins, metric, weights)
    # Nd = len(dat)
    DD = multi_autocp(dat, bins, metric, weights, Nd, pcpus)
    DD[DD == 0] = 1.0
    # DD = DD/(Nd*(Nd-1.0)) # factor of 2 cancels with 1/2 that needs to be done to remove double counting of pairs
    print (DD)
    return DD

def RRwcalc(datR, bins, metric, rweights):
    print ("Calculating RR with weights (parallelized)...\n RR=")
    # RR = autocorrw(datR, bins, metric, weights)
    # Nr = len(datR)
    RR = multi_autocpr(datR, bins, metric, rweights, Nr, pcpus)
    RR[RR == 0] = 1.0
    # RR = RR/(Nr*(Nr-1.0))
    print (RR)
    return RR

def DRwcalc(dat, datR, bins, metric, rweights):
    print ("Calculating DR with weights (parallelized)...\n DR=")
    # DR = crosscorrw(dat, datR, bins, metric, rweights)
    # Nd = len(dat)
    # Nr = len(datR)
    # DR = multi_crosscp(dat, datR, bins, metric, rweights, Nd, pcpus)
    DR = multi_crosscpdr(dat, datR, bins, metric, rweights, Nd, pcpus)
    DR[DR == 0] = 1.0
    # DR = DR/(Nd*Nr)
    print (DR/2.0)
    return DR/2.0

def RDwcalc(dat, datR, bins, metric, weights):
    print ("Calculating RD with weights...\n RD=")
    # DR = crosscorrwrd(dat, datR, bins, metric, weights)
    # Nd = len(dat)
    # Nr = len(datR)
    DR = multi_crosscp(dat, datR, bins, metric, weights, Nr, pcpus)
    DR[DR == 0] = 1.0
    # DR = DR/(Nd*Nr)
    print (DR/2.0)
    return DR/2.0

def autocorrw(dat, bins, metric, weights):
    # dbt = BallTree(dat, metric='pyfunc', func=metric)
    DD = np.zeros(len(bins)-1)
    binmax = max(bins)
    for i in tqdm(range(len(dat))):
        ind = dbt.query_radius(dat[i].reshape(1, -1), binmax)
        # wts=np.array([])
        for j in ind:
            # print ("i j")
            # print (i, j)
            # print ("ind[ind>i]")
            # print (ind[ind>i])
            dist0 = dist.cdist([dat[i], ], dat[j[j>i]], metric)[0]
            DD += np.histogram(dist0, bins=bins, weights=weights[j[j>i]])[0]
            # print (dist0,weights[j])
    print(DD)
    return DD

def crosscorrw(dat, datR, bins, metric, rweights):
    # rbt = BallTree(datR, metric='pyfunc', func=metric)
    DR = np.zeros(len(bins)-1)
    binmax = max(bins)
    for i in tqdm(range(len(dat))):
        ind = rbt.query_radius(dat[i].reshape(1, -1), binmax)
        # wts=np.array([])
        for j in ind:
            dist0 = dist.cdist([dat[i], ], datR[j], metric)[0]
            DR += np.histogram(dist0, bins=bins, weights=rweights[j])[0]
            # print (dist0,weights[j])
    return DR

def crosscorrwrd(dat, datR, bins, metric, weights):
    # dbt = BallTree(dat, metric='pyfunc', func=metric)
    RD = np.zeros(len(bins)-1)
    # p=multiprocessing.Pool(processes=multiprocessing.cpu_count())
    # RD=p.map(rdcalc, range(len(datR)))
    binmax = max(bins)
    for i in tqdm(range(len(datR))):
    # def rdcalc():
        ind = dbt.query_radius(datR[i].reshape(1, -1), binmax)
        #  wts=np.array([])
        for j in ind:
            dist0 = dist.cdist([datR[i], ], dat[j], metric)[0]
            RD += np.histogram(dist0, bins=bins, weights=weights[j])[0]
                # print (dist0,weights[j])
            # return RD
    print(RD)
    return RD

def autocorrwp(dat, bins, metric, weights, rNd, multi=False, queue=0):
    # dbt = BallTree(dat, metric='pyfunc', func=metric)
    DD = np.zeros(len(bins)-1)
    binmax = max(bins)
    for i in tqdm(rNd):
        ind = dbt.query_radius(dat[i].reshape(1, -1), binmax)
        # wts=np.array([])
        for j in ind:
            # print ("i j")
            # print (i, j)
            # print ("ind[ind>i]")
            # print (ind)
            # print (ind[ind>i])
            dist0 = dist.cdist([dat[i], ], dat[j[j>i]], metric)[0]
            DD += np.histogram(dist0, bins=bins, weights=weights[j[j>i]])[0]
            # print (dist0,weights[j])
    if multi:
        queue.put(DD)
    else:
        return DD
    # print (DD)
    return DD

def crosscorrwrdp(dat, datR, bins, metric, weights, rNr, multi=False, queue=0):
    # dbt = BallTree(dat, metric='pyfunc', func=metric)
    RD = np.zeros(len(bins)-1)
    binmax = max(bins)
    # p=multiprocessing.Pool(processes=multiprocessing.cpu_count())
    # RD=p.map(rdcalc, range(len(datR)))
    for i in tqdm(rNr):
    # def rdcalc():
        ind = dbt.query_radius(datR[i].reshape(1, -1), binmax)
        #  wts=np.array([])
        for j in ind:
            dist0 = dist.cdist([datR[i], ], dat[j], metric)[0]
            RD += np.histogram(dist0, bins=bins, weights=weights[j])[0]
    if multi:
        queue.put(RD)
    else:
        return RD
    # print(RD)
    return RD

def crosscorrwdrp(dat, datR, bins, metric, rweights, rNd, multi=False, queue=0):
    # dbt = BallTree(dat, metric='pyfunc', func=metric)
    DR = np.zeros(len(bins)-1)
    binmax = max(bins)
    # p=multiprocessing.Pool(processes=multiprocessing.cpu_count())
    # RD=p.map(rdcalc, range(len(datR)))
    for i in tqdm(rNd):
    # def rdcalc():
        ind = rbt.query_radius(dat[i].reshape(1, -1), binmax)
        #  wts=np.array([])
        for j in ind:
            dist0 = dist.cdist([dat[i], ], datR[j], metric)[0]
            DR += np.histogram(dist0, bins=bins, weights=rweights[j])[0]
    if multi:
        queue.put(DR)
    else:
        return DR
    # print(RD)
    return DR

def autocorrwpr(datR, bins, metric, rweights, rNr, multi=False, queue=0):
    # dbt = BallTree(dat, metric='pyfunc', func=metric)
    RR = np.zeros(len(bins)-1)
    binmax = max(bins)
    for i in tqdm(rNr):
        ind = rbt.query_radius(datR[i].reshape(1, -1), binmax)
        # print (ind)
        # wts=np.array([])
        for j in ind:
            # print("i")
            # print (i)
            # print ("j")
            # print (j)
            # print ("j[j>i]")
            # print (j[j>i])
            dist0 = dist.cdist([datR[i], ], datR[j[j>i]], metric)[0]
            RR += np.histogram(dist0, bins=bins, weights=rweights[j[j>i]])[0]
            # print (dist0,weights[j])
    if multi:
        queue.put(RR)
    else:
        return RR
    # print (RR)
    return RR

def multi_autocp(dat, bins, metric, weights, Nd, CORES=pcpus):

    DD = np.zeros(len(bins)-1)
    queues = [RetryQueue() for i in range(CORES)]
    args = [(dat, bins, metric, weights, range(int(Nd*i/CORES), int(Nd*(i+1)/CORES)), True, queues[i]) for i in range(CORES)]
    jobs = [Process(target=autocorrwp, args=(a)) for a in args]
    for j in jobs: j.start()
    for q in queues: DD += q.get()
    for j in jobs: j.join()

    return DD

def multi_autocpr(datR, bins, metric, rweights, Nr, CORES=pcpus):

    RR = np.zeros(len(bins)-1)
    queues = [RetryQueue() for i in range(CORES)]
    args = [(datR, bins, metric, rweights, range(int(Nr*i/CORES), int(Nr*(i+1)/CORES)), True, queues[i]) for i in range(CORES)]
    jobs = [Process(target=autocorrwpr, args=(a)) for a in args]
    for j in jobs: j.start()
    for q in queues: RR += q.get()
    for j in jobs: j.join()

    return RR

def multi_crosscp(dat, datR, bins, metric, weights, Nr, CORES=pcpus):

    DR = np.zeros(len(bins)-1)
    queues = [RetryQueue() for i in range(CORES)]
    args = [(dat, datR, bins, metric, weights, range(int(Nr*i/CORES), int(Nr*(i+1)/CORES)), True, queues[i]) for i in range(CORES)]
    jobs = [Process(target=crosscorrwrdp, args=(a)) for a in args]
    for j in jobs: j.start()
    for q in queues: DR += q.get()
    for j in jobs: j.join()

    return DR

def multi_crosscpdr(dat, datR, bins, metric, rweights, Nd, CORES=pcpus):

    DR = np.zeros(len(bins)-1)
    queues = [RetryQueue() for i in range(CORES)]
    args = [(dat, datR, bins, metric, rweights, range(int(Nd*i/CORES), int(Nd*(i+1)/CORES)), True, queues[i]) for i in range(CORES)]
    jobs = [Process(target=crosscorrwdrp, args=(a)) for a in args]
    for j in jobs: j.start()
    for q in queues: DR += q.get()
    for j in jobs: j.join()

    return DR

def retry_on_eintr(function, *args, **kw):
    while True:
        try:
            return function(*args, **kw)
        except IOError, e:
            if e.errno == KeyboardInterrupt:
                break
            elif e.errno == errno.EINTR:
                continue
            else:
                raise "multiprocessing pooling queue error", e

class RetryQueue(Queue):

    def get(self, block=True, timeout=None):
        return retry_on_eintr(Queue.get, self, block, timeout)


<gh_stars>1-10
"""
Process Whole slide images and their respective annotations into .mat file as per the paper & related code from  "HoVer-Net: Simultaneous Segmentation and Classification of Nuclei in Multi-Tissue Histology Images"

----------------------------------------------------------------------------------------------------
Each ground truth file is stored as a .mat file, with the keys:
'inst_map'
'type_map'
'inst_type'
'inst_centroid'

'inst_map' is a 1000x1000 array containing a unique integer for each individual nucleus. i.e the map ranges from 0 to N, where 0 is the background and N is the number of nuclei.

'type_map' is a 1000x1000 array where each pixel value denotes the class of that pixel. The map ranges from 0 to 7, where 7 is the total number of classes in CoNSeP.

"""

import os
import openslide
from xml.dom import minidom
import numpy as np
import openslide
from openslide import open_slide
from glob import glob
import cv2
import matplotlib.pyplot as plt
import scipy.io as sio
from PIL import Image
import scipy
import scipy.ndimage
from shapely.geometry import Polygon
from skimage import draw
import xml.etree.ElementTree as ET
import argparse
from misc.viz_utils import visualize_instances

def main(input_loc, output_loc):

    # Read svs files from the desired path
    count = 0
    if(input_loc):
        data_path = input_loc
    else:
        data_path = '/usr/local/opt/work/hover_net_modified/MoNuSAC' #Path to read data from
    if(output_loc):
        destination_path = output_loc # Path to save n-ary masks corresponding to xml files
    else:
        destination_path = '/usr/local/opt/work/hover_net_modified/'

    os.chdir(destination_path)

    # Create MoNuSAC folder
    try:
        os.mkdir(destination_path+'/MoNuSAC_processed')
    except OSError:
        print ("Creation of the mask directory %s failed" % destination_path)

    # Create sub-folders in the same pattern as CoNSeP
    # -- Images
    # -- Labels
    # -- Overlay
    try:
        os.mkdir(destination_path+'/MoNuSAC_processed/Images')
    except OSError:
        print ("Creation of the mask directory %s failed" % (destination_path + "\Images"))
    try:
        os.mkdir(destination_path+'/MoNuSAC_processed/Labels')
    except OSError:
        print ("Creation of the mask directory %s failed" % (destination_path + "\Labels"))
    try:
        os.mkdir(destination_path+'/MoNuSAC_processed/Overlay')
    except OSError:
        print ("Creation of the mask directory %s failed" % (destination_path + "\Overlay"))

    os.chdir(destination_path+'/MoNuSAC_processed')#Create folder named as MoNuSAC_masks
    patients = [x[0] for x in os.walk(data_path)]#Total patients in the data_path
    print('No. of Patients: ',len(patients))

    # Define Integer encoding for MoNuSAC classes (0:Background)
    nuclei_type_dict = {
            'Epithelial': 1, # ! Please ensure the matching ID is unique
            'Lymphocyte': 2,
            'Macrophage': 3,
            'Neutrophil': 4,
        }

    for patient_loc in patients:
        patient_name = patient_loc[len(data_path)+1:]#Patient name
        """## To make patient's name directory in the destination folder
        try:
            os.mkdir(patient_name)
        except OSError:
            print ("\n Creation of the patient's directory %s failed" % patient_name)
        """
        ## Read sub-images of each patient in the data path
        sub_images = glob(patient_loc+'/*.svs')
        for sub_image_loc in sub_images:

            gt = 0
            sub_image_name = sub_image_loc[len(data_path)+len(patient_name)+1:-4]
            print('File Name being processed:', sub_image_name)

            ## To make sub_image directory under the patient's folder
            """
            sub_image = './'+patient_name+'/'+sub_image_name #Destination path
            try:
                os.mkdir(sub_image)
            except OSError:
                print ("\n Creation of the patient's directory %s failed" % sub_image)
            """
            image_name = sub_image_loc
            img = openslide.OpenSlide(image_name)

            # If svs image needs to save in png
            cv2.imwrite(destination_path+'/MoNuSAC_processed/Images'+sub_image_name+'.png', np.array(img.read_region((0,0),0,img.level_dimensions[0])))
            #og_img = cv2.imread(destination_path+'/MoNuSAC_processed/Images'+sub_image_name+'.png')
            #og_img = cv2.cvtColor(og_img, cv2.COLOR_BGR2RGB)

            # Read xml file
            xml_file_name  = image_name[:-4]
            xml_file_name = xml_file_name+'.xml'
            print('XML:',xml_file_name)
            tree = ET.parse(xml_file_name)
            root = tree.getroot()

            n_ary_mask = np.transpose(np.zeros((img.read_region((0,0),0,img.level_dimensions[0]).size)))
            type_map = np.transpose(np.zeros((img.read_region((0,0),0,img.level_dimensions[0]).size)))

            #Generate n-ary mask for each cell-type
            for k in range(len(root)):
                label = [x.attrib['Name'] for x in root[k][0]]
                label = label[0]

                for child in root[k]:
                    for x in child:
                        r = x.tag
                        if r == 'Attribute':
                            count = count+1
                            print(count)
                            label = x.attrib['Name']
                            class_value = nuclei_type_dict[label]
                            print(label,':',class_value)

                            '''# Create directory for each label
                            sub_path = sub_image+'/'+label

                            try:
                                os.mkdir(sub_path)
                            except OSError:
                                print ("Creation of the directory %s failed" % label)
                            else:
                                print ("Successfully created the directory %s " % label)
                            '''

                        if r == 'Region':
                            regions = []
                            vertices = x[1]
                            coords = np.zeros((len(vertices), 2))
                            for i, vertex in enumerate(vertices):
                                coords[i][0] = vertex.attrib['X']
                                coords[i][1] = vertex.attrib['Y']
                            regions.append(coords)
                            poly = Polygon(regions[0])

                            vertex_row_coords = regions[0][:,0]
                            vertex_col_coords = regions[0][:,1]
                            fill_row_coords, fill_col_coords = draw.polygon(vertex_col_coords, vertex_row_coords, n_ary_mask.shape)
                            gt = gt+1 #Keep track of giving unique valu to each instance in an image
                            n_ary_mask[fill_row_coords, fill_col_coords] = gt
                            type_map[fill_row_coords, fill_col_coords] = class_value
                            # Stack togethor the inst_map & type_map

                            #nary_path = destination_path+'/MoNuSAC_processed/Labels'+sub_image_name+'_nary.tif'

                            #overlay_path = destination_path+'/MoNuSAC_processed/Overlay'+sub_image_name+'.png'

                            sio.savemat(destination_path+'/MoNuSAC_processed/Labels'+sub_image_name+".mat", {'inst_map':n_ary_mask,'class_map':type_map})
                            #overlay_image = visualize_instances(type_map, canvas=None)
                            #cv2.imwrite(overlay_path, overlay_image)
                            #cv2.imwrite(class_path, type_map)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--input',help='MoNuSAC data location')
    parser.add_argument('--output',help='MoNuSAC output location')
    args = parser.parse_args()

    main(args.input,args.output)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Nov 21 16:58:35 2021

@author: pablobottero
"""

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import openpyxl
from pandas.api.types import CategoricalDtype
import scipy.stats as stats
import seaborn as sns

os.chdir('/Users/pablobottero/github/master/python/data_analysis/tasks/second_one/')

#Abrir fichero
rank = pd.read_csv('billionaires_index.csv', sep=',', decimal=',')
print(rank.shape)
rank

#Obtener las personas de España de la lista
rank[rank["Country"] == "Spain"]

#Crear un fichero con los 50 primeros para trabajar con ellos
rank.head(50).to_csv("50b.csv")
bsl = pd.read_csv('50b.csv', sep=',', decimal=',') # Billionaire Sorted List

#Agrupar valores por industria y país
ig = bsl['Industry'].value_counts()
cg = bsl['Country'].value_counts()

ig.to_csv("industries.csv")
cg.to_csv("countries.csv")
bsl_ig = pd.read_csv('industries.csv', sep=',', decimal=',')
bsl_cg = pd.read_csv('countries.csv', sep=',', decimal=',')

#QC_done

# Imprimir un gráfico con el listado de  billonarios por industria
bsl.Industry.describe()
plt.hist (bsl.Industry)
i_table = pd.crosstab(index=bsl["Industry"], columns="count")
n=i_table.sum()
i_table2 = (i_table/n)*100
print(i_table2)
plt.bar((i_table.index), i_table['count'])
plt.xlabel('Industries')
plt.ylabel('%')
plt.title('Figure 2. Percentaje of billionaires \n by industry.')
plt.savefig('grafica.svg')
plt.show

#QC_done

# Imprimir un gráfico con el listado de  billonarios por pais
bsl_cg.Country:Number.describe()
plt.hist (bsl_cg.Country)
c_table = pd.crosstab(index=bsl["Country"], columns="count")
n=c_table.sum()
c_table2 = (c_table/n)*100
print(c_table2)
plt.bar(c_table2.index, c_table2['count'])
plt.xlabel('Industries')
plt.ylabel('%')
plt.title('Figure 1. Percentaje of billionaires \n by industry.')
plt.show

print(c_table.sort())
-----------------------
#Hacer graficos también por países
#Analizar los porcentajes por país y por industria
#Sacar el pocentaje de españoles en la lista completa
#Obtener los porcentajes de ricos por país en comparación con la lista (500)


<reponame>reemagit/flowcentrality
import pandas as pd
#from MyPythonLibs3 import bio
import networkx as nx
from tqdm.autonotebook import tqdm
from concurrent.futures import ProcessPoolExecutor, as_completed
import numpy as np
import os
from os.path import join, exists
from scipy.stats import mannwhitneyu
import matplotlib.pyplot as plt

def create_gidmap():
    from funcs.gidmap import GeneIDMap
    gmap = GeneIDMap()
    gmap.load_table('biomart', '../data/gene_info_genenames.txt', col2idtype_map={'Approved symbol':'symbol','NCBI Gene ID':'entrez','Ensembl gene ID':'ensembl'})
    #gmap.load_dict('ppi', ms.GENE_ID_CACHE_DATA)
    return gmap

gm = create_gidmap()

def read_gene_list(filepath):
    return pd.read_csv(filepath, header=None)[0].tolist()
def write_gene_list(filepath, genelist):
    pd.Series(genelist).to_csv(filepath, index=False, header=False)
def write_genesets_list(filepath, genesets):
    np.savetxt(filepath, genesets, fmt='%d', delimiter='\t')
def read_genesets_list(filepath):
    return np.loadtxt(filepath, delimiter='\t')

def read_network(filepath):
    return nx.read_edgelist(filepath, nodetype=int, delimiter='\t')

def read_flows(filepath):
    return pd.read_csv(filepath,sep='\t',index_col=0)

def read_paths(filepath):
    with open(filepath, 'r') as f:
        return [list(map(int,line.rstrip().split(','))) for line in f.readlines()]
def write_paths(filepath, paths):
    with open(filepath, 'w') as f:
        for path in paths:
            f.write(','.join(map(str,path)) + '\n')

def read_text(filepath):
    with open(filepath, 'r') as f:
        return [line.rstrip() for line in f.readlines()]

def write_text(filepath, lines):
    with open(filepath, 'w') as f:
        for line in lines:
            f.write(line + '\n')

def read_expr(filepath):
    return pd.read_csv(filepath, sep='\t',index_col='ENTREZ_GENE_ID')

def write_expr(filepath, expr):
    expr.to_csv(filepath, sep='\t')

def read_nonrectangular(filepath):
    with open(filepath, 'r') as f:
        return [list(map(float,line.rstrip().split(','))) for line in f.readlines()]
def write_nonrectangular(filepath, lines):
    with open(filepath, 'w') as f:
        for line in lines:
            f.write(','.join(map(str,line)) + '\n')

def read_config(filepath):
    import yaml
    with open(filepath, 'r') as stream:
        try:
            return yaml.safe_load(stream)
        except yaml.YAMLError as exc:
            print(exc)

def create_dir_if_not_exist(dirpath):
    if not exists(dirpath):
        os.makedirs(dirpath)

def parallel_process(function, array, n_jobs=None, use_kwargs=False, front_num=0, leave=True, mute=False):
    """
        A parallel version of the map function with a progress bar.

        Args:
            array (array-like): An array to iterate over.
            function (function): A python function to apply to the elements of array
            n_jobs (int, default=16): The number of cores to use
            use_kwargs (boolean, default=False): Whether to consider the elements of array as dictionaries of
                keyword arguments to function
            front_num (int, default=3): The number of iterations to run serially before kicking off the parallel job.
                Useful for catching bugs
        Returns:
            [function(array[0]), function(array[1]), ...]
    """
    bar_func = tqdm
    if mute:
        def identity(x, *args, **kwargs):
            return x
        bar_func = identity

    front = []
    #We run the first few iterations serially to catch bugs
    if front_num > 0:
        front = [function(**a) if use_kwargs else function(a) for a in array[:front_num]]
    #If we set n_jobs to 1, just run a list comprehension. This is useful for benchmarking and debugging.
    if n_jobs==1:
        return front + [function(**a) if use_kwargs else function(a) for a in tqdm(array[front_num:])]
    #Assemble the workers
    with ProcessPoolExecutor(max_workers=n_jobs) as pool:
        #Pass the elements of array into function
        if use_kwargs:
            futures = [pool.submit(function, **a) for a in array[front_num:]]
        else:
            futures = [pool.submit(function, a) for a in array[front_num:]]
        kwargs = {
            'total': len(futures),
            'unit': 'it',
            'unit_scale': True,
            'leave': leave
        }
        #Print out the progress as tasks complete

        for f in bar_func(as_completed(futures), **kwargs):
            pass
    out = []
    #Get the results from the futures.
    for i, future in bar_func(enumerate(futures)):
        try:
            out.append(future.result())
        except Exception as e:
            out.append(e)
    return front + out

def pack_variables(list_var, *other_vars):
    return [[elem]+list(other_vars) for elem in list_var]

def fisher_overlap(n_overlap, n1, n2, n_background):
    import fisher
    return fisher.pvalue(n_background - (n1 + n2 - n_overlap), n2 - n_overlap, n1 - n_overlap, n_overlap).right_tail

def fisher_overlap_set(set1, set2, background, check_consistency=True):
    if check_consistency:
        assert len(set1) == len(set(set1))
        assert len(set2) == len(set(set2))
        assert len(background) == len(set(background))
        assert set(set1).issubset(set(background))
        assert set(set2).issubset(set(background))
    n_overlap = len(set(set1).intersection(set(set2)))
    n1 = len(set1)
    n2 = len(set2)
    n_background = len(background)
    return fisher_overlap(n_overlap=n_overlap, n1=n1, n2=n2, n_background=n_background)

def non_nan(arr):
    arr = np.asarray(arr)
    return arr[~np.isnan(arr)]
def mwpval(obs,rdm):
    return mannwhitneyu(non_nan(obs),non_nan(rdm),alternative='greater').pvalue

def trim_text(string, maxchar=10):
    if len(string) > maxchar:
        return string[:maxchar-3] + '...'
    return string

def plot_significances(data_list, pos_list = None, delta_ratio = 30, bounds=None, no_fliers=False):
    if no_fliers:
        def whiskmax(data):
            iqr = np.percentile(data, 75) - np.percentile(data, 25)
            return data[data <= np.percentile(data, 75) + 1.5 * iqr].max()
        def whiskmin(data):
            iqr = np.percentile(data, 75) - np.percentile(data, 25)
            return data[data >= np.percentile(data, 25) - 1.5 * iqr].min()
        min_ = min([whiskmin(non_nan(data)) for pair in data_list for data in pair])
        max_ = max([whiskmax(non_nan(data)) for pair in data_list for data in pair])
    else:
        min_ = min([min(non_nan(data)) for pair in data_list for data in pair])
        max_ = max([max(non_nan(data)) for pair in data_list for data in pair])
    delta = (max_ - min_) / delta_ratio
    offset = 0
    if bounds is None:
        bounds = [1e-10, 1e-4, 0.05]
    bounds = [-np.inf] + bounds + [np.inf]
    sigtexts = ['***','**','*','n.s.']
    for pair,pos in zip(data_list, pos_list):
        pval = mwpval(non_nan(pair[1]), non_nan(pair[0]))
        sigtext = sigtexts[np.digitize(pval, bounds)-1]
        x1,x2 = pos
        y = max_+offset+delta
        plt.plot([x1, x1, x2, x2], [y, y+delta, y+delta, y], lw=1., c='black')
        plt.text((x1 + x2) * .5, y+delta, sigtext, ha='center', va='bottom', color='black')
        offset += delta * 4
    if plt.gca().get_ylim()[1] < max_ + offset + delta:
        plt.gca().set_ylim([None, max_ + offset + 3*delta])

def plot_significances_paths(data_list, pos=None, delta_ratio = 30, no_fliers=False):
    if pos is None:
        pos = [1,2,3]
    plot_significances([[data_list[1],data_list[2]], [data_list[0],data_list[2]]], [[pos[1],pos[2]],[pos[0],pos[2]]], delta_ratio = delta_ratio, no_fliers=no_fliers)


from scipy.fftpack import fft
import TicClassBgremover
import cv2
import numpy as np
from PIL import Image

br=TicClassBgremover.bgremover()
device='cuda'

print('start...')
size=256
model = br.load_model(device=device)
fimg, bimg = br.load_img(file_img='./dataset/demo/3.jpg',size=size)
print('mask ......')
PILbg = br.square_pad(bimg, 0)
PILbg = PILbg.resize((size, size), Image.ANTIALIAS)
cap = cv2.VideoCapture(0)

while 1:
    ret, frame = cap.read() #读取
    # img = cv2.cvtColor(np.asarray(frame), cv2.COLOR_RGB2BGR)
    ff=Image.fromarray(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB))
    PILimage = br.square_pad(ff, 0)
    PILimage = PILimage.resize((size, size), Image.ANTIALIAS)

    y=br.ai(PILimage, model, device=device)
    pp,p=br.ai_mask(y)

    fPil,blendPil=br.ai_blend(PILimage, PILbg,pp)
    imgp = cv2.cvtColor(np.asarray(blendPil), cv2.COLOR_RGB2BGR)

    cv2.imshow("capture", imgp) #显示
    if cv2.waitKey(10) & 0xff == ord('q'): #按q退出
        break


import numpy as np
from scipy.stats import rankdata
from sklearn.base import BaseEstimator, TransformerMixin

from julia import Julia
Julia(compiled_modules=False)
from julia import Relief as Relief_jl
from skrelief.relieff import Relieff

class TuRF(BaseEstimator, TransformerMixin):
    """sklearn compatible implementation of the TuRF algorithm.

    Reference:
        <NAME> and <NAME>. Tuning ReliefF for genome-wide
        genetic analysis. In <NAME>, <NAME>, and <NAME>.
        Rajapakse, editors, Evolutionary Computation,Machine Learning and
        Data Mining in Bioinformatics, pages 166–175. Springer, 2007.

    Args:
        n_features_to_select (int): number of features to select from dataset.
        num_it (int): number of iterations.
        rba (object): feature weighting algorithm wrapped by the VLSRelief algorithm. If equal
        to None, the default ReliefF RBA implemented in Julia is used.

    Attributes:
        n_features_to_select (int): number of features to select from dataset.
        num_it (int): number of iterations.
        _rba (object): feature weighting algorithm wrapped by the TuRF algorithm.
    """

    def __init__(self, n_features_to_select=10, num_it=10, rba=None):
        self.n_features_to_select = n_features_to_select
        self.num_it = num_it
        self._rba = rba

    def fit(self, data, target):
        """
        Rank features using TuRF feature selection algorithm

        Args:
            data (numpy.ndarray): matrix of data samples
            target (numpy.ndarray): vector of target values of samples

        Returns:
            (object): reference to self
        """

        # Compute feature weights and rank.
        if self._rba is not None:
            self.weights = Relief_jl.turf(data, target, self.num_it, self.rba_wrap)
        else:
            self.weights = Relief_jl.turf(data, target, self.num_it)
        self.rank = rankdata(-self.weights, method='ordinal')

        # Return reference to self.
        return self

    def transform(self, data):
        """
        Perform feature selection using computed feature ranks.

        Args:
            data (numpy.ndarray): matrix of data samples on which to perform feature selection.

        Returns:
            (numpy.ndarray): result of performing feature selection.
        """

        # select n_features_to_select best features and return selected features.
        msk = self.rank <= self.n_features_to_select  # Compute mask.
        return data[:, msk]  # Perform feature selection.


<filename>aurora/pipelines/time_series_helpers.py
import scipy.signal as ssig

from aurora.time_series.windowing_scheme import WindowingScheme

def validate_sample_rate(run_ts, config):
    if run_ts.sample_rate != config.sample_rate:
        print(
            f"sample rate in run time series {run_ts.sample_rate} and "
            f"processing config {config.sample_rate} do not match"
        )
        raise Exception
    return

def apply_prewhitening(config, run_xrts_input):
    if config["prewhitening_type"] == "first difference":
        run_xrts = run_xrts_input.diff("time")
    else:
        run_xrts = run_xrts_input
    return run_xrts

def apply_recoloring(config, stft_obj):
    if config["prewhitening_type"] == "first difference":
        from aurora.time_series.frequency_domain_helpers import get_fft_harmonics
        from numpy import pi

        freqs = get_fft_harmonics(config.num_samples_window, config.sample_rate)
        prewhitening_correction = 1.0j * 2 * pi * freqs  # jw
        stft_obj /= prewhitening_correction
    return stft_obj

def run_ts_to_stft_scipy(config, run_xrts_orig):
    """
    Parameters
    ----------
    config
    run_xrts

    Returns
    -------

    """
    import xarray as xr

    run_xrts = apply_prewhitening(config, run_xrts_orig)

    windowing_scheme = WindowingScheme(
        taper_family=config.taper_family,
        num_samples_window=config.num_samples_window,
        num_samples_overlap=config.num_samples_overlap,
        taper_additional_args=config.taper_additional_args,
        sample_rate=config.sample_rate,
    )
    # stft_obj = run_xrts.copy(deep=True)
    stft_obj = xr.Dataset()
    for channel_id in run_xrts.data_vars:
        ff, tt, specgm = ssig.spectrogram(
            run_xrts[channel_id].data,
            fs=config.sample_rate,
            window=windowing_scheme.taper,
            nperseg=config.num_samples_window,
            noverlap=config.num_samples_overlap,
            detrend="linear",
            scaling="density",
            mode="complex",
        )

        # drop Nyquist>
        ff = ff[:-1]
        specgm = specgm[:-1, :]

        import numpy as np

        specgm *= np.sqrt(2)

        # make time_axis
        tt = tt - tt[0]
        tt *= config.sample_rate
        time_axis = run_xrts.time.data[tt.astype(int)]

        xrd = xr.DataArray(
            specgm.T,
            dims=["time", "frequency"],
            coords={"frequency": ff, "time": time_axis},
        )
        stft_obj.update({channel_id: xrd})

    stft_obj = apply_recoloring(config, stft_obj)

    return stft_obj

def run_ts_to_stft(config, run_xrts_orig):
    """

    Parameters
    ----------
    config : ShortTimeFourierTransformConfig object
    run_ts ; xarray.core.dataset.Dataset, normally extracted from mth5.RunTS

    Returns
    -------

    """
    from aurora.time_series.windowed_time_series import WindowedTimeSeries

    windowing_scheme = WindowingScheme(
        taper_family=config.taper_family,
        num_samples_window=config.num_samples_window,
        num_samples_overlap=config.num_samples_overlap,
        taper_additional_args=config.taper_additional_args,
        sample_rate=config.sample_rate,
    )

    run_xrts = apply_prewhitening(config, run_xrts_orig)

    windowed_obj = windowing_scheme.apply_sliding_window(
        run_xrts, dt=1.0 / config.sample_rate
    )
    windowed_obj = WindowedTimeSeries.detrend(data=windowed_obj, detrend_type="linear")

    tapered_obj = windowed_obj * windowing_scheme.taper
    # stft_obj = WindowedTimeSeries.apply_stft(data=tapered_obj,
    #                                          sample_rate=windowing_scheme.sample_rate,
    #                                          detrend_type="linear",
    # scale_factor=windowing_scheme.linear_spectral_density_calibration_factor)

    stft_obj = windowing_scheme.apply_fft(
        tapered_obj, detrend_type=config.extra_pre_fft_detrend_type
    )
    stft_obj = apply_recoloring(config, stft_obj)

    return stft_obj

def run_ts_to_calibrated_stft(run_ts, run_obj, config, units="MT"):
    """
    Parameters
    ----------
    run_ts
    run_obj
    config
    units

    Returns
    -------

    """
    stft_obj = run_ts_to_stft(config, run_ts.dataset)
    stft_obj = calibrate_stft_obj(stft_obj, run_obj, units=units)

    return stft_obj

def calibrate_stft_obj(stft_obj, run_obj, units="MT", channel_scale_factors=None):
    """

    Parameters
    ----------
    stft_obj
    run_obj
    units
    scale_factors : dict
        keyed by channel, supports a single scalar to apply to that channels data
        Useful for debugging.  Should not be used in production and should throw a
        warning if it is not None

    Returns
    -------

    """
    for channel_id in stft_obj.keys():
        mth5_channel = run_obj.get_channel(channel_id)
        channel_filter = mth5_channel.channel_response_filter
        if not channel_filter.filters_list:
            print("WARNING UNEXPECTED CHANNEL WITH NO FILTERS")
            if channel_id == "hy":
                channel_filter = run_obj.get_channel("hx").channel_response_filter
        calibration_response = channel_filter.complex_response(stft_obj.frequency.data)
        if channel_scale_factors:
            try:
                channel_scale_factor = channel_scale_factors[channel_id]
            except KeyError:
                channel_scale_factor = 1.0
            calibration_response /= channel_scale_factor
        if units == "SI":
            print("Warning: SI Units are not robustly supported issue #36")

        stft_obj[channel_id].data /= calibration_response
    return stft_obj


import tensorflow as tf
import numpy as np
import random
import cv2
import os
import matplotlib.pyplot as plt
import pickle
from scipy import misc
import math

import config as cfg
from resnet import *
from input_fn import Data
#np.random.seed(1)
#tf.set_random_seed(1)
#random.seed(1)

batch_size = cfg.BATCH_SIZE
keep_prob = 0.8
regular_weight = 0.001

def get_sess():
    tf_config = tf.ConfigProto()
    tf_config.gpu_options.allow_growth = True
    return tf.Session(config=tf_config)

def get_img2(pic0):
    input_pics = []
    y_data = []
    y_rank = []
    for p in pic0:
        pic_name, ppl = p
        im = cv2.imread(os.path.join(cfg.FILE_PATH, pic_name))
        im = cv2.cvtColor(im,cv2.COLOR_BGR2RGB)
        im = cv2.resize(im, (cfg.IMG_SIZE, cfg.IMG_SIZE))
        input_pics.append(im)
        y_data.append(ppl)
        y_rank.append(math.floor(ppl+2.5))

    input_pics = np.reshape(input_pics, (-1, cfg.IMG_SIZE, cfg.IMG_SIZE, 3))
    y_data = np.reshape(y_data, (-1, 1))
    y_rank = np.reshape(y_rank, (-1))
    return input_pics, y_data, y_rank

data = Data()

# Preprocessing
input_shape = (batch_size, cfg.IMG_SIZE, cfg.IMG_SIZE, 3)
imgs0 = tf.placeholder(tf.uint8, shape=[batch_size, cfg.IMG_SIZE, cfg.IMG_SIZE, 3])
batch_imgs = []
for i in range(imgs0.shape[0]):
    image = imgs0[i]
    height = cfg.IMG_SIZE
    width = cfg.IMG_SIZE
    with tf.name_scope('eval_image', [image, height, width]):
        if image.dtype != tf.float32:
            image = tf.image.convert_image_dtype(image, dtype=tf.float32)
        if height and width:
            image = tf.div(image, 255.0)
            image = tf.expand_dims(image,0)
            if i == 0:
                batch_imgs = image
            else:
                batch_imgs = tf.concat([batch_imgs, image], 0)
        images = batch_imgs

# Initialize
y_value = tf.placeholder(dtype=tf.float32, shape=[batch_size,1])
is_train = tf.placeholder(tf.bool, name="is_train")

arg_scope = resnet_arg_scope(weight_decay=regular_weight,
                     batch_norm_decay=0.9,
                     batch_norm_epsilon=1e-5,
                     batch_norm_scale=True,
                     activation_fn=tf.nn.relu,
                     use_batch_norm=True)

# ResNet
with slim.arg_scope(arg_scope):
    net, end_points = resnet_v2_50(images, is_training=is_train)

# Final Conv
with slim.arg_scope([slim.conv2d, slim.max_pool2d, slim.avg_pool2d]):
    with tf.variable_scope('Logits_out'):
        # (12 x 12) x 2048
        net = slim.avg_pool2d(net, kernel_size=[4, 4], stride=4, padding='VALID', scope='AvgPool_7x7')
        # 3 x 3 x 256
        net = slim.conv2d(net, 1, [3, 3], activation_fn=None, padding='VALID', weights_regularizer=slim.l2_regularizer(regular_weight), scope='Last_conv')
        # 1 x 1 x 1
        net = tf.squeeze(net, [1, 2], name='SpatialSqueeze')

# Variable To Train and Restore
checkpoint_exclude_scopes = "Logits_out"
exclusions = []
if checkpoint_exclude_scopes:
    exclusions = [scope.strip() for scope in checkpoint_exclude_scopes.split(',')]
print (exclusions)
variables_to_restore = []
variables_to_train = []
for var in slim.get_model_variables():
    excluded = False
    print(var.op.name)
    for exclusion in exclusions:
        if var.op.name.startswith(exclusion):
            excluded = True
            variables_to_train.append(var)
            print ("ok")
            print (var.op.name)
            break
    if not excluded:
        variables_to_restore.append(var)

# file = './output/LOG_res_e4_pre.txt'
epoch = 1800
with get_sess() as sess:
    tf.logging.info("Creating the model...")
    sess.run(tf.global_variables_initializer())
    saver2 = tf.train.Saver(tf.global_variables(), max_to_keep=5)
    saver2.restore(sess, './savenet/loss_128_ave1_pre_e4/loss_128_ave1_pre_e4-%s'%epoch)

    ############## test set in study cities #################
    city_names = ['Beijing', 'Guangzhou', 'Shanghai', 'Harbin', 'Kunming', 'Wuhan', 'Lasa', 'Lanzhou']
    for city_name in city_names:
        print(city_name)
        test_file = open('./output/loss_128_ave1_pre_e4/test_%s_%s.txt'%(epoch, city_name),'w')
        train_city_data = data.generate_train_city_data(city_name)
        l = len(train_city_data)
        for k in range(l):
            test_input_pics, test_y_data, test_name = data.obtain_train_city_data(k)
            # test acc
            pred = sess.run(net, feed_dict={imgs0:test_input_pics, y_value:test_y_data, is_train:False})

            test_file.write("####### step:%d >> pic_name: %s >> real_value: %s >> pred: %s\n"%(k,test_name[0], test_y_data[0][0], pred[0][0]))

    ############## additional test cities ##################
    # Shenyang, Jinan, Tianjin, Shijiazhuang, Hefei, Changsha, Nanchang, Shenzhen, Luoyang, Dalian
    city_names = [['Shenyang',12315,12364,4170,4194]] # [[city, min_lon, max_lon, min_lat, max_lat]]
    for city_name in city_names:
        print(city_name)
        test_file = open('./output/loss_128_ave1_pre_e4/test_%s_%s.txt'%(epoch, city_name[0]),'w')
        test_city_key = data.generate_test_city_data(file='./data/%s.txt'%city_name[0], lon1=city_name[1], lon2=city_name[2], lat1=city_name[3], lat2=city_name[4])
        l = len(test_city_key)
        for k in range(l):
            test_input_pics, test_y_data, test_name = data.obtain_test_city_data(k)

            # test acc
            pred = sess.run(net, feed_dict={imgs0:test_input_pics, y_value:test_y_data, is_train:False})
            test_file.write("####### step:%d >> pic_name: %s >> real_value: %s >> pred: %s\n"%(k,test_name[0], test_y_data[0][0], pred[0][0]))


from scipy.signal import butter, lfilter
import numpy as np
from librosa.core.time_frequency import mel_frequencies
import warnings

def butter_bandpass(lowcut, highcut, fs, order=3):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype="band")
    return b, a

def butter_bandpass_filter(data, lowcut, highcut, fs, order=3):
    if highcut > int(fs / 2):
        warnings.warn("Highcut is too high for bandpass filter. Setting to nyquist")
        highcut = int(fs / 2)
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = lfilter(b, a, data)
    return y

def window_rms(a, window_size):
    a2 = np.power(a, 2)
    window = np.ones(window_size) / float(window_size)
    return np.sqrt(np.convolve(a2, window, "valid"))

def RMS(data, rate, rms_stride, rms_window, rms_padding, noise_thresh):
    """
    Take data, run and RMS filter over it
    """

    # we compute root mean squared over a window, where we stride by rms_stride seconds for speed
    rms_data = window_rms(
        data.astype("float32")[:: int(rms_stride * rate)],
        int(rate * rms_window * rms_stride),
    )
    rms_data = rms_data / np.max(rms_data)

    # convolve a block filter over RMS, then threshold it, so to call everything with RMS > noise_threshold noise
    block_filter = np.ones(int(rms_padding * rms_stride * rate))  # create our filter

    # pad the data to be filtered
    rms_threshed = np.concatenate(
        (
            np.zeros(int(len(block_filter) / 2)),
            np.array(rms_data > noise_thresh),
            np.zeros(int(len(block_filter) / 2)),
        )
    )
    # convolve on our filter
    sound_threshed = np.array(np.convolve(rms_threshed, block_filter, "valid") > 0)[
        : len(rms_data)
    ]

    return rms_data, sound_threshed

import sys
import os

def prepare_mel_matrix(hparams, rate, return_numpy=True, GPU_backend=False):
    """ Create mel filter
    """
    # import tensorflow if needed
    if "tf" not in sys.modules:
        if not GPU_backend:
            os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  # see issue #152
            os.environ["CUDA_VISIBLE_DEVICES"] = ""
        import tensorflow as tf
        tf.enable_eager_execution()
        assert tf.executing_eagerly()

    # create a filter to convolve with the spectrogram
    mel_matrix = tf.signal.linear_to_mel_weight_matrix(
        num_mel_bins=hparams.num_mel_bins,
        num_spectrogram_bins=int(hparams.n_fft / 2) + 1,
        sample_rate=rate,
        lower_edge_hertz=hparams.mel_lower_edge_hertz,
        upper_edge_hertz=hparams.mel_upper_edge_hertz,
        dtype=tf.dtypes.float32,
        name=None,
    )

    # gets the center frequencies of mel bands
    mel_f = mel_frequencies(
        n_mels=hparams.num_mel_bins + 2,
        fmin=hparams.mel_lower_edge_hertz,
        fmax=hparams.mel_upper_edge_hertz,
    )

    # Slaney-style mel is scaled to be approx constant energy per channel (from librosa)
    enorm = tf.dtypes.cast(
        tf.expand_dims(
            tf.constant(
                2.0
                / (mel_f[2 : hparams.num_mel_bins + 2] - mel_f[: hparams.num_mel_bins])
            ),
            0,
        ),
        tf.float32,
    )

    mel_matrix = tf.multiply(mel_matrix, enorm)
    mel_matrix = tf.divide(mel_matrix, tf.reduce_sum(mel_matrix, axis=0))
    if return_numpy:
        return mel_matrix.numpy()
    else:
        return mel_matrix


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np

from scipy.fftpack import dct
from collections import deque

class MFCCKaldi:
    '''
    TODO port Kaldi mfcc to Python. Use similar parameters as
    in suggested in __init__ function
    '''

    def __init__(self, sourcerate=16000, framesize=512,
                 usehamming=True, preemcoef=0.97,
                 numchans=26, ceplifter=22, numceps=12,
                 enormalise=True, zmeansource=True, usepower=True, usec0=True,
                 usecmn=False, usedelta=True, useacc=True, n_last_frames=0,
                 lofreq=125, hifreq=3800, mel_banks_only=False):
        self.sourcerate = sourcerate
        self.framesize = framesize
        self.usehamming = usehamming
        self.preemcoef = preemcoef
        self.numchans = numchans
        self.ceplifter = ceplifter
        self.enormalise = enormalise
        self.zmeansource = zmeansource
        self.usepower = usepower
        self.usec0 = usec0
        self.usecmn = usecmn
        self.usedelta = usedelta
        self.useacc = useacc
        self.numceps = numceps
        self.lofreq = lofreq
        self.hifreq = hifreq
        self.mel_banks_only = mel_banks_only

    def param(self, frame):
        """Compute the MFCC coefficients in a way similar to the HTK."""

class MFCCFrontEnd:
    """This is an a CLOSE approximation of MFCC coefficients computed by the HTK.

    The frame size should be a number of power of 2.

    TODO: CMN is not implemented. It should normalise only teh cepstrum, not the delta or acc coefficients.

    It was not tested to give exactly the same results the HTK. As a result,
    it should not be used in conjunction with models trained on speech
    parametrised with the HTK.

    Over all it appears that this implementation of MFCC is worse than the one from the HTK.
    On the VAD task, the HTK features score 90.8% and the this features scores only 88.7%.
    """

    def __init__(self, sourcerate=16000, framesize=512,
                 usehamming=True, preemcoef=0.97,
                 numchans=26, ceplifter=22, numceps=12,
                 enormalise=True, zmeansource=True, usepower=True, usec0=True, usecmn=False,
                 usedelta=True, useacc=True, n_last_frames = 0,
                 lofreq=125, hifreq=3800, mel_banks_only = False):
        self.sourcerate = sourcerate
        self.framesize = framesize
        self.usehamming = usehamming
        self.preemcoef = preemcoef
        self.numchans = numchans
        self.ceplifter = ceplifter
        self.enormalise = enormalise
        self.zmeansource = zmeansource
        self.usepower = usepower
        self.usec0 = usec0
        self.usecmn = usecmn
        self.usedelta = usedelta
        self.useacc = useacc
        self.numceps = numceps
        self.lofreq = lofreq
        self.hifreq = hifreq
        self.mel_banks_only = mel_banks_only

        self.prior = 0.0

        self.n_last_frames = n_last_frames
        self.mfcc_queue = deque(maxlen=4 + n_last_frames)
        self.mfcc_delta_queue = deque(maxlen=4 + n_last_frames)

        self.init_hamming()
        self.init_mel_filter_bank()
        self.init_cep_liftering_weights()

    def freq_to_mel(self, freq):
        return 1127 * np.log(1.0 + freq / 700.0)

    def mel_to_freq(self, mel):
        return 700 * (np.exp(mel / 1127) - 1.0)

    def init_hamming(self):
        self.hamming = np.hamming(self.framesize)

    def init_mel_filter_bank(self):
        """Initialise the triangular mel freq filters."""

        minMel = self.freq_to_mel(self.lofreq)
        maxMel = self.freq_to_mel(self.hifreq)

#    print "MM", minMel, "MM", maxMel

        # Create a matrix for triangular filters, one row per filter
        filterMatrix = np.zeros((self.numchans, self.framesize / 2 + 1))

        melRange = np.array(xrange(self.numchans + 2))
#    print "MR", melRange

        melCenterFilters = melRange * (maxMel - minMel) / (
            self.numchans + 1) + minMel
#    print "MCF", melCenterFilters

        dfreq = self.sourcerate / self.framesize
        # each array index represent the center of each triangular filter
        centerIndex = np.array(
            np.round(self.mel_to_freq(melCenterFilters) / dfreq), int)
#    print "CI", centerIndex

        for i in xrange(self.numchans):
            start, centre, end = centerIndex[i:i + 3]
            k1 = np.float32(centre - start)
            k2 = np.float32(end - centre)
            up = (np.array(xrange(start, centre)) - start) / k1
            down = (end - np.array(xrange(centre, end))) / k2

            filterMatrix[i][start:centre] = up
            filterMatrix[i][centre:end] = down

        self.mel_filter_bank = filterMatrix.transpose()
#    print "SMFB", self.mel_filter_bank.shape

    def init_cep_liftering_weights(self):
        cep_lift_weights = np.zeros((self.numceps, ))
        a = np.pi / self.ceplifter
        b = self.ceplifter / 2.0
        for i in range(self.numceps):
            cep_lift_weights[i] = 1.0 + b * np.sin(i * a)

        self.cep_lift_weights = cep_lift_weights

    def preemphasis(self, frame):
        out_frame = np.zeros_like(frame)
        out_frame[0] = frame[0] - self.preemcoef * self.prior
        for i in range(1, len(frame)):
            out_frame[i] = frame[i] - self.preemcoef * frame[i - 1]

        self.prior = frame[-1]

        return out_frame

    def param(self, frame):
        """Compute the MFCC coefficients in a way similar to the HTK."""
        # zero mean
        if self.zmeansource:
            frame = frame - np.mean(frame)
        # preemphasis
        frame = self.preemphasis(frame)
        # apply hamming window
        if self.usehamming:
            frame = self.hamming * frame

        complex_spectrum = np.fft.rfft(frame)
#    print "LCS", len(complex_spectrum)
        power_spectrum = complex_spectrum.real * complex_spectrum.real + \
            complex_spectrum.imag * complex_spectrum.imag
        # compute only power spectrum if required
        if not self.usepower:
            power_spectrum = np.sqrt(power_spectrum)

#    print "SPS",power_spectrum.shape
        mel_spectrum = np.dot(power_spectrum, self.mel_filter_bank)
        # apply mel floor
        for i in range(len(mel_spectrum)):
            if mel_spectrum[i] < 1.0:
                mel_spectrum[i] = 1.0
        mel_spectrum = np.log(mel_spectrum)

        if self.mel_banks_only:
            mfcc = mel_spectrum
            self.mfcc_queue.append(mel_spectrum)
        else:
            cepstrum = dct(mel_spectrum, type=2, norm='ortho')
            c0 = cepstrum[0]
            htk_cepstrum = cepstrum[1:self.numceps + 1]
            # cepstral liftering
            cep_lift_mfcc = self.cep_lift_weights * htk_cepstrum

            if self.usec0:
                mfcc = np.append(cep_lift_mfcc, c0)
            else:
                mfcc = cep_lift_mfcc

            # compute delta and acceleration coefficients if requested
            self.mfcc_queue.append(mfcc)

#        print len(self.mfcc_queue)

            if self.usedelta:
#      print "LMQ", len(self.mfcc_queue)
                if len(self.mfcc_queue) >= 2:
                    delta = np.zeros_like(mfcc)
                    for i in range(1, len(self.mfcc_queue)):
                        delta += self.mfcc_queue[i] - self.mfcc_queue[i - 1]
                    delta /= len(self.mfcc_queue) - 1

                    self.mfcc_delta_queue.append(delta)
                else:
                    delta = np.zeros_like(mfcc)

            if self.useacc:
                if len(self.mfcc_delta_queue) >= 2:
                    acc = np.zeros_like(mfcc)
                    for i in range(1, len(self.mfcc_delta_queue)):
                        acc += self.mfcc_delta_queue[i] - \
                            self.mfcc_delta_queue[i - 1]
                    acc /= len(self.mfcc_delta_queue) - 1
                else:
                    acc = np.zeros_like(mfcc)

            if self.usedelta:
                mfcc = np.append(mfcc, delta)
            if self.useacc:
                mfcc = np.append(mfcc, acc)

        for i in range(self.n_last_frames):
            if len(self.mfcc_queue) > i + 1 :
                mfcc = np.append(mfcc, self.mfcc_queue[-1-i-1])
            else:
                mfcc = np.append(mfcc, np.zeros_like(self.mfcc_queue[-1]))

        return mfcc.astype(np.float32)


'''
Calculate single-category modularity correlations with a given performance
metric. You can enter your own files or try our defaults! Note that
we recommend using csv files, as this script reads data using csv.
'''

# imports
import csv
import argparse
from scipy.stats import spearmanr as sr

# process command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument("--modularity_file",
    help = "name of file with modularity scores, columns = categories,"
            + "1st row = column headers, no row headers",
    default = "data/3_2.csv")
parser.add_argument("--metrics_file",
    help = "name of file with task performance metrics, 1 column, no headers,"
            + "assuming length aligns with modularity_file",
    default = "movies_accuracy.csv")
parser.add_argument("--out_file",
    help = "name of file to write correlation to, if desired",
    default = "out.txt")
args = parser.parse_args()

# scores: stores modularity scores - outer key = k, inner key = level
scores = {}

# reading in modularity data
with open(args.modularity_file, newline = "") as csvfile:
    reader = csv.reader(csvfile)
    rows = [row for row in reader]
    for i in range(len(rows[0])):
        scores[rows[0][i]] = [float(score[i]) for score in rows[1:]]

metric = []

# reading in performance data, assuming 1-column list w/ no header
with open(args.metrics_file, "r") as csvfile:
    reader = csv.reader(csvfile)
    rows = [row for row in reader]
    for i in range(len(rows)):
        metric.append(float(rows[i][0]))

# writes results to a file, separates columns by commas
with open(args.out_file, "w") as o:
    o.write("category,correlation\n") # column headers
    for cat in scores.keys():
        o.write(cat + "," + str(sr(scores[cat], metric).correlation) + "\n")


<reponame>guochunhe/Vitis-AI
# Copyright 2019 Xilinx Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
import os
from skimage import io
import cv2
from scipy.spatial.distance import cdist
import argparse

import torch.nn.functional as F
from network.resnet18 import Resnet18
from network.resnet_small import Resnetsmall
from ipdb import set_trace

def load_reid_model(net_name, model_path):
    '''
    function to load my trained reid model
    '''
    assert net_name == 'facereid_small' or net_name == 'facereid_large'
    if net_name == 'facereid_small':
        net = Resnetsmall()
    else:
        net = Resnet18()

    checkpoint = torch.load(model_path)
    pretrain_dict = checkpoint['state_dict']
    model_dict = net.state_dict()
    for i in pretrain_dict:
        if 'classifier' in i or 'fc' in i:  continue
        net.state_dict()[i].copy_(pretrain_dict[i])
    net = net.cuda()
    net.eval()
    return net

def im_preprocess(image):
    image = np.asarray(image, np.float32)
    image /=255.0
    image -= np.array([0.485, 0.456, 0.406]).reshape(1, 1, -1)
    image /= np.array([0.229, 0.224, 0.225]).reshape(1, 1, -1)
    image = np.transpose(image, [2,0,1]) # data format: hwc->chw
    return image

def parse_args():
    parser = argparse.ArgumentParser(description='FaceReid demo')
    parser.add_argument('--data_dir', default = '../data/test_imgs/', help='directory to test data')
    parser.add_argument('--network', choices=['facereid_small', 'facereid_large'],
                        help='set network, choose from [facereid_small, facereid_large] ')
    parser.add_argument('--model_path', required=True, help='path to model')
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    img_dir = args.data_dir
    query = ['1.jpg']
    gallery = ['2.jpg',
                '3.jpg',
                '4.jpg',
                '5.jpg',
                '6.jpg',
                '7.jpg']
    img_list = query + gallery

    resize_wh = 80 if args.network == 'facereid_small' else 96

    imgs = [io.imread(os.path.join(img_dir, imgname)) for imgname in img_list]
    imgs = np.asarray([im_preprocess(cv2.resize(p, (resize_wh, resize_wh))) for p in imgs], dtype=np.float32)
    reid_model = load_reid_model(args.network, args.model_path)
    print('[INFO] Load model: {}, model path: {}'.format(args.network, args.model_path))

    with torch.no_grad():
        im_var = Variable(torch.from_numpy(imgs))
        im_var = im_var.cuda()
        feats = F.normalize(reid_model(im_var)).cpu()

    q_feats = feats[:len(query)]
    g_feats = feats[len(query):]
    distmat = cdist(q_feats, g_feats)
    print('[INFO] Query-gallery-distance matrix: \n  ', distmat)

    for i in range(len(query)):
        g_idx = np.where(distmat[i]==np.min(distmat[i]))[0][0]
        print('[INFO] Link query image {} with gallery image {} as same id'.format(query[i], gallery[g_idx]))


import logging
from fractions import Fraction
from contextlib import contextmanager
import pexpect
from crl.interactivesessions.shells.terminalclient import TerminalComm

__copyright__ = 'Copyright (C) 2019, Nokia'

LOGGER = logging.getLogger(__name__)

class LostStrComm(object):

    def __init__(self, probability_of_lost, modifier=None):
        self._count = 0
        self._probability_of_lost = probability_of_lost
        self._modifier = modifier or self._default_modifier
        self._strcomm = None
        self._write_empty = False
        self._in_lost_context = False

    @staticmethod
    def _default_modifier(*_):
        return b''

    @contextmanager
    def in_lost(self):
        self._in_lost_context = True
        try:
            yield None
        finally:
            self._in_lost_context = False

    def set_strcomm(self, strcomm):
        self._strcomm = strcomm

    @property
    def comm(self):
        return self._strcomm.comm

    def write_str(self, s):
        self._strcomm.comm.clear()
        self._strcomm.write_str(s)
        written = self._strcomm.comm.written
        new_written = self._modified_s(written)
        if new_written or self._write_empty:
            LOGGER.debug('===== writing: %s', repr(new_written))
            self._strcomm.comm.write_direct(new_written)

    def read_str(self):
        ret = self._strcomm.read_str()
        modified_ret = self._modified_s(ret)
        LOGGER.debug('==== Received: %s, is lost: %d', ret, not modified_ret)
        LOGGER.debug('==== Modified return %s', modified_ret)
        if not modified_ret:
            raise pexpect.TIMEOUT('==== Message is lost: {!r}'.format(ret))

        return ret

    def _modified_s(self, s):
        if self._is_msg_lost():
            new_s = self._modifier(s)
            LOGGER.debug('==== Msg modified, count=%s, s=%s, new_s=%s',
                         self._count, s, repr(new_s))
            s = new_s
        else:
            LOGGER.debug('==== Msg success, count=%s, %s', self._count, s)
        return s

    def _is_msg_lost(self):
        p = Fraction(self._probability_of_lost)
        ret = self._count % p.denominator < p.numerator and self._in_lost_context
        self._count += 1
        if ret:
            LOGGER.debug('==== message LOST')
        else:
            LOGGER.debug('==== message SUCCESS')
        return ret

class CustomTerminalComm(TerminalComm):
    def __init__(self, *args, **kwargs):
        super(CustomTerminalComm, self).__init__(*args, **kwargs)
        self._written = b''

    @property
    def written(self):
        return self._written

    def clear(self):
        self._written = b''

    def _write(self, s):
        self._written += s

    def write_direct(self, s):
        super(CustomTerminalComm, self)._write(s)