import io

import os

from typing import List, Union

import dask.dataframe as dd

import pandas as pd

import validators

from openomics.io.files import get_pkg_data_filename

BCR_PATIENT_BARCODE_COL = "bcr_patient_barcode"

HISTOLOGIC_SUBTYPE_COL = "histologic_subtype"

PATHOLOGIC_STAGE_COL = "pathologic_stage"

TUMOR_NORMAL_COL = 'tumor_normal'

PREDICTED_SUBTYPE_COL = 'predicted_subtype'

TUMOR = "Tumor"

NORMAL = "Normal"

__all__ = ['ClinicalData']

class ClinicalData:

    """This class manages the clinical data tables to handle the patient's

    phenotype data, as well as the treatment, and sample data associated to each

    patient.

    """

    pathologic_stage_map = {'Stage IA': 'Stage I', 'Stage IB': 'Stage I',

                            'Stage IIA': 'Stage II', 'Stage IIB': 'Stage II',

                            'Stage IIIA': 'Stage III', 'Stage IIIB': 'Stage III'}

    def __init__(self,

                 file: Union[str, io.StringIO, pd.DataFrame, dd.DataFrame],

                 patient_index: str,

                 columns: List[str] = None):

        """

        Args:

            file (str, io.StringIO, pd.DataFrame): either a path to the

                patients clinical data file, or a DataFrame.

            patient_index (str): the patient's ID column name

            columns (List[str]): default None. Specifies the columns to import,

                if None, then import all columns.

        """

        # self.cohort_name = cohort_name

        self.patient_column = patient_index

        if columns and patient_index not in columns:

            columns.append(patient_index)

        if isinstance(file, (pd.DataFrame, dd.DataFrame)):

            self.patient = file

        elif isinstance(file, io.StringIO):

            file.seek(0)  # Needed since the file was previous read to extract columns information

            self.patient = pd.read_table(file,

                                         skiprows=[1, 2],

                                         na_values=["[Not Available]", "[Unknown]", "[Not Applicable]",

                                                    "[Discrepancy]"],

                                         usecols=columns

                                         )

        elif isinstance(file, str) and validators.url(file):

            dataurl, filename = os.path.split(file)

            file = get_pkg_data_filename(dataurl + "/", filename)

            self.patient = pd.read_table(file)

        elif isinstance(file, str) and os.path.isfile(file):

            self.patient = pd.read_table(file,

                                         skiprows=[1, 2],

                                         na_values=["[Not Available]", "[Unknown]", "[Not Applicable]",

                                                    "[Discrepancy]"],

                                         usecols=columns

                                         )

        else:

            raise FileNotFoundError("{}".format(file))

        self.patient_barcodes = self.patient[patient_index].tolist()

        self.patient.set_index(patient_index, inplace=True)

        # Rename columns

        self.patient.rename({"ajcc_pathologic_tumor_stage": PATHOLOGIC_STAGE_COL,

                             "histological_type": HISTOLOGIC_SUBTYPE_COL,

                             "histologic_diagnosis.1": HISTOLOGIC_SUBTYPE_COL}, axis=1, inplace=True)

        self.patient.replace({PATHOLOGIC_STAGE_COL: ClinicalData.pathologic_stage_map}, inplace=True)

    @classmethod

    def name(self):

        """Returns the name of the class, i.e. 'ClinicalData'"""

        return self.__class__.__name__

    def build_clinical_samples(self, all_samples, index="bcr_patient_barcode"):

        """Build table with samples clinical data from patients :param

        all_samples:

        Args:

            all_samples:

            index:

        """

        self.samples = pd.DataFrame(index=all_samples)

        self.samples.index.name = index

        self.samples.index = self.samples.index.str[:-4]  # Cut sample barcode for TCGA

        # Merge patients clinical data with patient barcode as index

        # target = pd.merge(target, self.patient,

        #                      how="left", left_on="patient_barcode", right_on="patient_barcode")

        self.samples = self.samples.join(self.patient, on=index, how="left", rsuffix="_")

        # self.samples.dropna(axis=0, subset=["bcr_patient_barcode"], inplace=True) # Remove samples without clinical data

        self.samples = self.samples[self.samples[PATHOLOGIC_STAGE_COL] != "[Discrepancy]"]

        self.samples.loc[self.samples.index.str.contains(

            "-11"), TUMOR_NORMAL_COL] = NORMAL  # Change stage label of normal samples to "Normal"

        self.samples.loc[self.samples.index.str.contains(

            "-01"), TUMOR_NORMAL_COL] = TUMOR  # Change stage label of normal samples to "Normal"

    def add_drug_response_data(self, file_path="nationwidechildrens.org_clinical_drug.txt",

                               patient_column="bcr_patient_barcode",

                               columns=None,

                               drug_name_col=None, response_column=None):

        """

        Args:

            file_path:

            patient_column:

            columns:

            drug_name_col:

            response_column:

        """

        if columns is None:

            columns = ['bcr_patient_barcode', 'pharmaceutical_therapy_drug_name',

                       'pharmaceutical_therapy_type', 'treatment_best_response']

        if not os.path.exists(file_path):

            raise FileNotFoundError(file_path)

        self.drug_name_col = drug_name_col

        self.response_column = response_column

        self.drugs = pd.read_table(file_path,

                                   sep="\t",

                                   skiprows=[1, 2],

                                   na_values=["[Not Available]", "[Unknown]", "[Not Applicable]"],

                                   usecols=columns

                                   )

        self.drugs.set_index(patient_column, inplace=True)

    def add_biospecimen_data(self, file_path="genome.wustl.edu_biospecimen_sample.txt",

                             patient_col_name="bcr_patient_barcode",

                             columns=['bcr_sample_barcode', 'sample_type']):

        """

        Args:

            file_path:

            patient_col_name:

            columns:

        """

        if not os.path.exists(file_path):

            raise FileNotFoundError(file_path)

        self.biospecimen = pd.read_table(file_path, sep="\t", skiprows=[1, ],

                                         na_values=["[Not Available]", "[Unknown]", "[Not Applicable]"],

                                         usecols=columns

                                         )

        self.sample_barcodes = self.biospecimen[patient_col_name].tolist()

        self.biospecimen.set_index(patient_col_name, inplace=True)

    def get_patient_barcodes(self):

        return self.patient_barcodes

    def get_sample_barcodes(self):

        return self.sample_barcodes