import pandas as pd

import numpy as np

from scipy import stats

from scipy.stats import norm

from sklearn.preprocessing import MinMaxScaler

import warnings

import subprocess

import re

def abundance_filter(d, abundance):

    return d.loc[:, d.mean() >= abundance]

def isDigit(x):

    try:

        float(x)

        return True

    except ValueError:

        return False

def missing_filter(d, missing_ratio):

    if d.applymap(np.isreal).all().sum() == d.shape[1]:

        d = d.loc[:, d.applymap(np.isnan).sum() <= d.shape[0] * missing_ratio]

        d = d.fillna(0)

    else:

        d = d.loc[:, d.applymap(lambda x:isDigit(x)).sum() >= d.shape[0] * missing_ratio]

        d_array = d.values

        d_array[~d.applymap(lambda x:isDigit(x))] = 0

        d.loc[:, :] = d_array

        d = d.astype(float)

    #d = d.loc[:, (d == 0).sum() <= d.shape[0] * missing_ratio]

    return d

def log2_scale(d):

    return np.log2(d+1)

def ln_scale(d):

    return np.log(d+1)

def log10_scale(d):

    return np.log10(d+1)

def normalize_scale(d):

    d = d + 1

    d = d.apply(lambda x: stats.boxcox(x)[0])

    d.loc[:,:] = MinMaxScaler().fit_transform(d.values)

    return d

def ppoints(n, a=None):

    try:

        n = np.float64(len(n))

    except TypeError:

        n = np.float64(n)

    if a is None:

        a = 3.0/8 if(n <= 10) else 1.0/2

    return (np.arange(n) + 1 - a)/(n + 1 - 2*a)

def qqnorm(y):

    ina = np.isnan(y)

    if ina.sum() > 0:

        yN = y

        y = y[~ina]

    n = y.shape[0]

    if n == 0:

        print('y is empty or has only NAs')

        return np.array([])

    x = np.around(norm.ppf(ppoints(n)[np.argsort(np.argsort(y))]), decimals=15)

    if ina.sum() > 0:

        y = x

        x = yN

        x[~ina] = y

    return x

def trait_correct(pc, y):

    pc1 = pd.concat([pd.DataFrame(np.ones((y.shape[0], 1)), index=pc.index), pc], axis=1)

    vhat = np.dot(np.linalg.pinv(np.dot(pc1.T, pc1)), np.dot(pc1.T, y))

    if len(vhat.shape) == 1:

        y_corr = y - np.dot(pc, vhat[1:])

    else:

        y_corr = y - np.dot(pc, vhat[1:, :])

    return y_corr

def pc_calc(bed, pc_num):

    try:

        from rpy2.robjects.packages import importr

        from rpy2.robjects import pandas2ri

        from rpy2.rinterface_lib.embedded import RRuntimeError

        import rpy2.robjects as robjects

        pandas2ri.activate()

        warnings.filterwarnings("ignore")

        base = importr('base')

        utils = importr('utils')

        if not base.require('bigsnpr', quietly=True)[0]:

            utils_path = subprocess.check_output('locate modas/utils', shell=True, text=True, encoding='utf-8')

            # utils_path = '/'.join(re.search('\n(.*site-packages.*)\n', utils_path).group(1).split('/')[:-1])

            utils_path = re.search('\n(.*site-packages.*)\n', utils_path).group(1)

            if not utils_path.endswith('utils'):

                utils_path = '/'.join(utils_path.split('/')[:-1])

            utils.install_packages(utils_path + '/Matrix_1.6-5.tar.gz', repos=robjects.rinterface.NULL, type='source',

                                   quiet=True)

            utils.install_packages('bigsnpr', dependence=True, repos='https://cloud.r-project.org', quiet=True)

        robjects.r['options'](warn=-1)

        bigsnpr = importr('bigsnpr')

        bigstatsr = importr('bigstatsr')

        base.sink('/dev/null')

        g = bigsnpr.snp_readBed(bed, backingfile=base.tempfile())

        g = bigsnpr.snp_attach(g)

        # svd = bigsnpr.snp_autoSVD(g.rx2[0], infos_chr=g[2]['chromosome'], ncores=1,

        #                           infos_pos=g[2]['physical.pos'], thr_r2=np.nan, k=pc_num)

        svd = bigsnpr.snp_autoSVD(g.rx2('genotypes'), infos_chr=g.rx2('map').rx2('chromosome'),

                                  infos_pos=g.rx2('map').rx2('physical.pos'), thr_r2=np.nan, k=pc_num)

        base.sink()

    except RRuntimeError:

        return None

    else:

        pc = bigstatsr.predict_big_SVD(svd)

        # pc = base.data_frame(pc, row_names=g[1]['sample.ID'])

        #pc = base.cbind(pc.rownames, pc)

        pc = pd.DataFrame(pc, index=g.rx2('fam').rx2('sample.ID'))

        pc.columns = ['PC' + str(i) for i in range(1, pc_num+1)]

        return pc