import sys

import pandas as pd

def read_(file):

    # read file

    if file.endswith('.csv'):

        df = pd.read_csv(file, index_col=0)

    elif file.endswith('.csv.gz'):

        df = pd.read_csv(file, compression='gzip', index_col=0)

    else:

        print('\n[Error]: The program cannot infer the format of {} . Currently, only the csv format is supported, please ensure that the file name suffix is .csv or .csv.gz.'.format(file))

        sys.exit(0)

    return df

def read_omics(args):

    omics = []

    for file in args.omic_file:

        df = read_(file)

        df = df.fillna(0)  # fill nan with 0

        omics.append(df)

    return omics

def read_label(args):

    file = args.label_file

    df = read_(file)

    df = df.rename(

        columns={df.columns.values[0]: 'label'})

    return df

def read_clin(args):

    file = args.clin_file

    df = None

    if not file is None:

        df = read_(file)

        # fill na

        df = df.fillna(0)

    return df

def process(df_omics, df_label, df_clin):

    # extract patient id

    patients = [df_tmp.index.to_list() for df_tmp in df_omics]

    patients.append(df_label.index.to_list())

    if not df_clin is None:

        patients.append(df_clin.index.to_list())

    # get shared patients between different data

    patients_shared = patients[0]

    for i in range(1, len(patients)):

        patients_shared = list(set(patients_shared).intersection(patients[i]))

    # extract shared patients' data

    for i in range(len(df_omics)):

        df_omics[i] = df_omics[i].loc[patients_shared, :].sort_index()

    df_label = df_label.loc[patients_shared, :].sort_index()

    if not df_clin is None:

        df_clin = df_clin.loc[patients_shared, :].sort_index()

    return df_omics, df_label, df_clin

# api

def read_dataset(args):

    # 1. read raw dataset

    # (1) read omics dataset

    df_omics = read_omics(args)

    # (2) read label

    df_label = read_label(args)

    # (3) read clinical feature

    df_clin = read_clin(args)

    # 2. process

    df_omics, df_label, df_clin = process(df_omics, df_label, df_clin)

    # 3. return clean dataset

    return df_omics, df_label, df_clin