import os, sys
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import connected_components
from sklearn.preprocessing import OneHotEncoder
import pandas as pd
import numpy as np
import pickle
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.sampler import WeightedRandomSampler
import torch
from .params import data_path, main_data_path, cohort_data_path, kinship_path, withdraw_path, fam_path
from .utils import get_fields, get_row_last_values, remove_kinships, save_dict, load_dict, print_sys


class ukbb_cohort:
    
    def __init__(self, main_data_path, cohort_data_path, withdraw_path, keep_relatives = False):
        self.keep_relatives = keep_relatives
        self.cohort_data_path = cohort_data_path
        self.main_data_path = main_data_path
        
        if keep_relatives:
            cohort_path = os.path.join(cohort_data_path, 'cohort_with_relatives.pkl')
        else:
            cohort_path = os.path.join(cohort_data_path, 'cohort_no_relatives.pkl')
        
        if not os.path.exists(cohort_path):
            print_sys('construct from scratch...')
            '''
            exclusions:
            the original uk biobank nature paper supplementary S 3.4:
            22006: Genetic ethnic grouping -> to retain only white british ancestry

            https://www.frontiersin.org/articles/10.3389/fgene.2022.866042/full
            22018: genetic relatedness exclusions
            22019: sex chromosome aneuploidy 
            31 <-> 22001: mismatch between self-reported sex and genetically determined sex 
            22010: recommended genomic analysis exclusions, signs of insufficient data quality

            (optional) further remove relatives based on KING relative scores, choose the first one in relative group
            remove the list of eids who do not want to be in the study anymore            
            '''
            
            all_field_ids = [22006, 22018, 22019, 22001, 22010, 31]
            df_main = get_fields(all_field_ids, main_data_path)
            cur_size = len(df_main)
            print_sys('Total sample size: ' + str(cur_size))
            df_main = df_main[df_main['22006-0.0'] == 1]
            print_sys('Keeping only white british ancestry (ID: 22006), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
            cur_size = len(df_main)

            df_main = df_main[df_main['22018-0.0'].isnull()]
            print_sys('Removing genetics related samples (ID: 22018), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
            cur_size = len(df_main)

            df_main = df_main[df_main['22019-0.0'].isnull()]
            print_sys('Removing sex chromosome aneuploidy (ID: 22019), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
            cur_size = len(df_main)

            df_main = df_main[df_main['31-0.0'] == df_main['22001-0.0']]
            print_sys('Removing samples with mismatched self-reported sex and genetic determined sex (ID: 31 <-> 22001), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
            cur_size = len(df_main)

            df_main = df_main[df_main['22010-0.0'].isnull()]
            print_sys('Removing samples with genomic data quality (ID: 22010), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
            cur_size = len(df_main)
            
            save_dict(os.path.join(cohort_data_path, 'cohort_with_relatives.pkl'), df_main.eid.values)
            
            kinship_mask = remove_kinships(df_main.eid)
            df_main = df_main[kinship_mask]
            save_dict(os.path.join(cohort_data_path, 'cohort_no_relatives.pkl'), df_main.eid.values)
        else:
            print_sys('Found local copy...')
            
        self.cohort = load_dict(cohort_path)
        print_sys('There are ' + str(len(self.cohort)) + ' samples!')
        
        if keep_relatives:
            self.no_rel_eid = load_dict(os.path.join(cohort_data_path, 'cohort_no_relatives.pkl'))
            
        if os.path.exists(withdraw_path):
            ## todo: when there is a withdraw file, implement this...
            pass
        
    def get_covariates(self, to_plink = False, plink_num_pca = 15, return_full = False, plink_filter = False):
        '''
        covariates:

        31: sex
        21003: Age when attended assessment centre
        22009: pca
        54: assessment center
        batch from params.fam_path file
        '''
        covar_path = os.path.join(self.cohort_data_path, 'covariates_all.pkl')
        if os.path.exists(covar_path):
            print_sys('Found local copy...')
            self.covar = load_dict(covar_path)
        else:
            print_sys('construct co-variates from scratch...')
            df_covar = get_fields([31, 54, 21003, 22009], self.main_data_path)
            column_name_map = {'22009-0.' + str(i): 'pca ' + str(i) for i in range(1, 41)}
            column_name_map['31-0.0'] = 'sex'
            column_name_map['21003-0.0'] = 'age'
            column_name_map['54-0.0'] = 'assessment_center'
            self.covar = df_covar.rename(columns = column_name_map)
            
            enc = OneHotEncoder(handle_unknown='ignore')
            enc.fit(self.covar['assessment_center'].unique().reshape(-1,1))
            center_array = enc.transform(self.covar['assessment_center'].values.reshape(-1,1)).toarray()
            center_one_hot = pd.DataFrame(center_array).astype('int').rename(columns = dict(zip(range(22), ['center_' + str(i) for i in range(22)])))
            
            self.covar = self.covar.drop(['21003-1.0', '21003-2.0', '21003-3.0', 'assessment_center', '54-1.0', '54-2.0', '54-3.0'], axis = 1)
            self.covar = self.covar.join(center_one_hot)
            
            df_fam = pd.read_csv(fam_path)
            enc = OneHotEncoder(handle_unknown='ignore')
            enc.fit(df_fam.trait.unique().reshape(-1,1))
            batch_one_hot = enc.transform(df_fam['trait'].values.reshape(-1,1)).toarray()
            batch_num = batch_one_hot.shape[1]
            id2batch = dict(zip(df_fam.fid.values, batch_one_hot.astype(int)))
            df_batch = pd.DataFrame(np.stack(self.covar['eid'].apply(lambda x: id2batch[x] if x in id2batch else np.zeros(batch_one_hot.shape[1]).astype(int)).values)).rename(columns = dict(zip(range(batch_num), ['batch_' + str(i) for i in range(batch_num)])))
            self.covar = self.covar.join(df_batch)
            
            save_dict(covar_path, self.covar)
            print_sys('Done! Saving...')
            
        if not to_plink:
            if return_full:
                return self.covar.reset_index(drop = True)
            else:
                return self.covar[self.covar.eid.isin(self.cohort)].reset_index(drop = True)
        else:
            plink_path = os.path.join(self.cohort_data_path, 'covar_pca' + str(plink_num_pca) + '_all_real_value')
            if plink_filter:
                plink_path += '_null_removed'
            plink_path += '.txt'
            
            if not os.path.exists(plink_path):
                pca_columns = [i for i in self.covar.columns.values if (i[:3]=='pca') and int(i.split()[-1]) <= plink_num_pca]
                #center_one_hot_columns = ['center_' + str(i) for i in range(22)]
                #batch_columns = ['batch_' + str(i) for i in range(batch_num)]
                #self.covar[['eid', 'eid', 'age', 'sex'] + pca_columns + center_one_hot_columns + batch_columns].to_csv(plink_path, header=None, index=None, sep=' ')
                center = np.argmax(self.covar.loc[:, self.covar.columns.str.contains('center')].values, axis = 1)
                batch = np.argmax(self.covar.loc[:, self.covar.columns.str.contains('batch')].values, axis = 1)
                self.covar = self.covar.iloc[:, :43]
                self.covar['assessment_center'] = center
                self.covar['batch'] = batch
                if plink_filter:
                    self.covar = self.covar[self.covar.eid.isin(self.cohort)].reset_index(drop = True)
                self.covar[['eid', 'eid', 'age', 'sex', 'assessment_center', 'batch'] + pca_columns].to_csv(plink_path, header=None, index=None, sep=' ')
            self.covar_plink = pd.read_csv(plink_path, header = None, sep = ' ')
            return self.covar_plink
    
    def get_external_traits(self, trait_name, to_plink = False, to_str = True, random_seed = 42, sep_cohort = False, randomize = False, use_sample_size = False, sample_size = -1, randomize_seed = 42):
        '''
        example:
        standing heights: 50
        '''
        if trait_name in ['body_BALDING1', 'cancer_BREAST', 'disease_ALLERGY_ECZEMA_DIAGNOSED', 'disease_HYPOTHYROIDISM_SELF_REP', 'other_MORNINGPERSON', 'pigment_SUNBURN']:
            trait_type = 'binary'
        else:
            trait_type = 'continuous'

        pheno_path = os.path.join(self.cohort_data_path, str(trait_name) + '_pheno.pkl')
        if os.path.exists(pheno_path):
            print_sys('Found local copy...')
            self.pheno = load_dict(pheno_path)
        else:
            print_sys('construct phenotype from scratch...')
            
            self.pheno = pd.read_csv(os.path.join(data_path, 'full_gwas', trait_name+'_'+trait_type+'.csv'))
            self.pheno['eid'] = self.pheno.eid.astype('int')
            self.pheno = self.pheno[self.pheno['pheno'].notnull()]
            if trait_type == 'binary':
                self.pheno['pheno'] += 1
                self.pheno['pheno'] = self.pheno['pheno'].astype(int)
            save_dict(pheno_path, self.pheno)
            print_sys('Done! Saving...')
            
            
        
        # filtering to cohorts incl. with/without relatives            
        self.pheno = self.pheno[self.pheno.eid.isin(self.cohort)].reset_index(drop = True)
            
        if to_str:
            self.pheno['eid'] = self.pheno['eid'].astype('str')
        if not to_plink:
            return self.pheno
        else:
            plink_path = os.path.join(self.cohort_data_path, str(trait_name) + '_plink')
            if self.keep_relatives:
                plink_path = plink_path + '_with_relatives'
            else:
                plink_path = plink_path + '_no_relatives'
                
            if use_sample_size:
                plink_path = plink_path + '_' + str(sample_size) + '_' + str(random_seed)
            
            if sep_cohort:
                plink_path += '_sep_cohort'
             
            if randomize:
                plink_path += '_randomize' + str(randomize_seed)
            
            plink_path = plink_path + '.txt'
            
            if randomize:
                self.pheno['pheno'] = self.pheno['pheno'].sample(frac = 1, random_state = randomize_seed).values
            
            if use_sample_size:
                from sklearn.model_selection import train_test_split
                print('random_seed:', random_seed)
                pheno_shuffle = self.pheno.sample(frac = 1, random_state = random_seed)
                all_ids, y = pheno_shuffle.eid.values, pheno_shuffle['pheno'].values
                train_val_ids, test_ids, y_train_val, y_test = all_ids[:sample_size], all_ids[sample_size:], y[:sample_size], y[sample_size:]
                if sep_cohort:
                    self.pheno = self.pheno[self.pheno.eid.isin(test_ids)]
                else:
                    self.pheno = self.pheno[self.pheno.eid.isin(train_val_ids)]
                        
            if not os.path.exists(plink_path):
                print_sys('Saving...')
                self.pheno[['eid', 'eid', self.pheno.columns.values[-1]]].to_csv(plink_path, header=None, index=None, sep=' ')
            else:
                print_sys('Already existed! Loading...')
            
            self.pheno_plink = pd.read_csv(plink_path, header = None, sep = ' ')
            return self.pheno_plink
    
    
    
    
    def get_phenotype(self, field_id, aggregate = 'last_value', to_plink = False, to_str = True, normalize = 'None', frac = 1, random_seed = 42, fastgwa_match = False, icd10 = False, icd10_level = 2, sep_cohort = False, randomize = False, use_sample_size = False, sample_size = -1, randomize_seed = 42):
        '''
        example:
        standing heights: 50
        '''
        pheno_path = os.path.join(self.cohort_data_path, str(field_id) + '_pheno.pkl')
        if os.path.exists(pheno_path):
            print_sys('Found local copy...')
            self.pheno = load_dict(pheno_path)
        else:
            print_sys('construct phenotype from scratch...')
            if icd10:
                ## field_id is icd10 level
                icd10_df = self.get_icd10(to_plink = True, level = icd10_level, get_all = True)
                self.pheno = icd10_df[['FID', field_id]].rename(columns = {'FID': 'eid'})
                self.pheno['eid'] = self.pheno.eid.astype('int')
            else:
                ## from raw data field id
                self.pheno = get_fields([field_id], self.main_data_path)
            save_dict(pheno_path, self.pheno)
            print_sys('Done! Saving...')
        
        if len(self.pheno.columns.values) > 2:
            print_sys('There are multiple index for this phenotype... aggregate...')
            if aggregate == 'last_value':
                print_sys('Getting the latest measure...')
                tmp = pd.DataFrame()
                tmp['eid'] = self.pheno.loc[:, 'eid']
                tmp[str(field_id)] = get_row_last_values(self.pheno.iloc[:, 1:])
                self.pheno = tmp
                print_sys('There are ' + str(len(self.pheno[self.pheno[str(field_id)].isnull()])) + ' samples with NaN values. Removing them ...')
                self.pheno = self.pheno[self.pheno[str(field_id)].notnull()]
                
        if fastgwa_match:
            # get the number of without relatives:
            if not self.keep_relatives:
                raise ValueError('If you turned fastgwa_match = True, then keep_relatives = True!')
            self.rel_ratio = len(self.pheno[self.pheno.eid.isin(self.no_rel_eid)])/len(self.pheno[self.pheno.eid.isin(self.cohort)])
        
        # filtering to cohorts incl. with/without relatives            
        self.pheno = self.pheno[self.pheno.eid.isin(self.cohort)].reset_index(drop = True)

        if normalize != 'None':
            y = self.pheno[str(field_id)].values
            if normalize == 'log':
                y = np.log(y)
            elif normalize == 'std':
                y = (y - np.mean(y))/np.std(y)
            elif normalize == 'quantile_normalization':
                from sklearn.preprocessing import quantile_transform
                y = quantile_transform(y.reshape(-1,1), output_distribution = 'normal', random_state = 42).reshape(-1)
            self.pheno[str(field_id)] = y 
            
        if to_str:
            self.pheno['eid'] = self.pheno['eid'].astype('str')
        if not to_plink:
            return self.pheno
        else:
            plink_path = os.path.join(self.cohort_data_path, str(field_id) + '_plink')
            if self.keep_relatives:
                plink_path = plink_path + '_with_relatives'
            else:
                plink_path = plink_path + '_no_relatives'
                
            if normalize != 'None':
                plink_path = plink_path + '_' + str(normalize)
            if use_sample_size:
                plink_path = plink_path + '_' + str(sample_size) + '_' + str(random_seed)
            else:
                if frac != 1:
                    plink_path = plink_path + '_' + str(frac) + '_' + str(random_seed)
                
            if fastgwa_match:
                plink_path += '_match'
                
            if sep_cohort:
                plink_path += '_sep_cohort'
             
            if randomize:
                plink_path += '_randomize' + str(randomize_seed)
                
            
            plink_path = plink_path + '.txt'
            
            if randomize:
                self.pheno[str(field_id)] = self.pheno[str(field_id)].sample(frac = 1, random_state = randomize_seed).values
            
            
            if use_sample_size:
                from sklearn.model_selection import train_test_split
                
                if icd10:
                    df_cases = self.pheno[self.pheno[str(field_id)] == 2]
                    df_cases_shuffle = df_cases.sample(frac = 1, random_state = random_seed)
                    all_ids, y = df_cases_shuffle.eid.values, df_cases_shuffle[str(field_id)].values
                    train_val_ids, test_ids, y_train_val, y_test = all_ids[:sample_size], all_ids[sample_size:], y[:sample_size], y[sample_size:]
                    train_val_ids = np.concatenate((train_val_ids, self.pheno[self.pheno[str(field_id)] == 1].eid.values))
                    self.pheno = self.pheno[self.pheno.eid.isin(train_val_ids)]
                    if sep_cohort:
                        raise NotImplementedError
                else:
                    print('random_seed', random_seed)
                    pheno_shuffle = self.pheno.sample(frac = 1, random_state = random_seed)
                    all_ids, y = pheno_shuffle.eid.values, pheno_shuffle[str(field_id)].values
                    train_val_ids, test_ids, y_train_val, y_test = all_ids[:sample_size], all_ids[sample_size:], y[:sample_size], y[sample_size:]
                    if fastgwa_match:
                        raise ValueError('Not used anymore...')
                    if sep_cohort:
                        self.pheno = self.pheno[self.pheno.eid.isin(test_ids)]
                    else:
                        self.pheno = self.pheno[self.pheno.eid.isin(train_val_ids)]
            else:
                if frac!=1:
                    from sklearn.model_selection import train_test_split
                    all_ids, y = self.pheno.eid.values, self.pheno[str(field_id)].values                
                    train_val_ids, test_ids, y_train_val, y_test = train_test_split(all_ids, y, test_size=frac, random_state=random_seed)
                    if fastgwa_match:
                        train_val_ids, test_ids, y_train_val, y_test = train_test_split(train_val_ids, y_train_val, test_size=1-self.rel_ratio, random_state=42)
                    if sep_cohort:
                        self.pheno = self.pheno[self.pheno.eid.isin(test_ids)]
                    else:
                        self.pheno = self.pheno[self.pheno.eid.isin(train_val_ids)]
                        
            if not os.path.exists(plink_path):
                self.pheno[['eid', 'eid', self.pheno.columns.values[-1]]].to_csv(plink_path, header=None, index=None, sep=' ')
            else:
                print_sys('Already existed! Loading...')
            
            self.pheno_plink = pd.read_csv(plink_path, header = None, sep = ' ')
            return self.pheno_plink
            
        
    def get_icd10(self, to_plink = False, level = 2, get_all = False):
        '''
        icd10: 41270
        '''
        pheno_path = os.path.join(self.cohort_data_path, 'icd10.pkl')
        level_str = 'level' + str(level)
        if os.path.exists(pheno_path):
            print_sys('Found local copy...')
            self.icd10 = load_dict(pheno_path)
        else:
            print_sys('construct from scratch...')
            icd10_raw_concat = get_fields([41270], self.main_data_path)
            icd10_columns = icd10_raw_concat.columns.values[1:]
            icd10_tuple = icd10_raw_concat.apply(lambda x: (x.eid, x[icd10_columns][x[icd10_columns].notnull()].values), axis = 1)
            icd10 = pd.DataFrame(list(icd10_tuple.values)).rename(columns = {0: 'eid', 1: 'level3'})
            icd10['level2'] = icd10['level3'].apply(lambda x: np.unique([i[:3] for i in x]))
            save_dict(pheno_path, icd10)
            print_sys('Done! Saving...')
            self.icd10 = icd10
        if get_all:
            self.pheno = self.icd10.reset_index(drop = True)
        else:
            self.pheno = self.icd10[self.icd10.eid.isin(self.cohort)].reset_index(drop = True)
        if not to_plink:
            return self.pheno
        else:            
            if self.keep_relatives or get_all:
                plink_path = os.path.join(self.cohort_data_path, 'icd10_plink_with_relatives_' + level_str + '.txt')
            else:
                plink_path = os.path.join(self.cohort_data_path, 'icd10_plink_no_relatives_' + level_str + '.txt')
            
            if os.path.exists(plink_path):
                print_sys("Found local copy...")
                self.icd10_plink = pd.read_csv(plink_path, sep=' ')
            else:
                print_sys('transforming to plink files... takes around 1 min...')
                unique_icd10 = np.unique([item for sublist in self.pheno[level_str].values for item in sublist])
                icd10_2_idx = dict(zip(unique_icd10, range(len(unique_icd10))))
                idx_2_icd10 = dict(zip(range(len(unique_icd10)), unique_icd10))

                self.pheno[level_str + '_idx'] = self.pheno[level_str].apply(lambda x: [icd10_2_idx[i] for i in x])

                tmp = np.zeros((len(self.pheno), len(unique_icd10)), dtype=np.int8)
                for idx, i in enumerate(self.pheno[level_str + '_idx'].values):
                    tmp[idx, i] = 1

                icd10_plink = pd.DataFrame(tmp).rename(columns = idx_2_icd10)
                icd102sample_size = dict(icd10_plink.sum(axis = 0))
                icd_100 = [i for i,j in icd102sample_size.items() if j > 100]
                icd10_plink = icd10_plink + 1
                icd10_plink['IID'] = self.pheno.eid.values
                icd10_plink['FID'] = self.pheno.eid.values
                icd10_plink = icd10_plink.loc[:, ['FID', 'IID'] + icd_100]
                print_sys('Only using ICD10 codes with at least 100 cases...')
                print_sys('There are ' + str(len(icd_100)) + ' ICD10 codes with at least 100 cases.')
                icd10_plink.to_csv(plink_path, index=None, sep=' ')
                self.icd10_plink = icd10_plink
                
            return self.icd10_plink