from collections import OrderedDict

import pandas as pd
from loguru import logger

from .constants import MEASUREMENT_DF_LABEL, DESIGN_DF_LABEL, PADJ_COL_PREFIX, PVALUE_COL_PREFIX, FC_COL_PREFIX
from .info import get_info


class SingleOmicsData():
    def __init__(self, data_type, measurement_df, design_df, feature_annot_df=None, significant_df=None):
        self.data_type = data_type

        self.original_measurement_df = measurement_df.copy()
        self.original_design_df = design_df.copy()
        self.original_feature_annot_df = feature_annot_df
        self.original_significant_df = significant_df

        # extract columns containing p-values and fold changes into its own dataframe
        if significant_df is not None:
            self.significant_df = significant_df
        else:  # try to extract from the data
            keep_df, drop_df = self._get_significant_df(measurement_df)
            if not drop_df.empty:
                self.significant_df = drop_df
                measurement_df = keep_df

        # clean data
        cleaned_measurement_df, cleaned_design_df = self._clean_data(measurement_df, design_df)
        msg = 'cleaned_measurement_df = %s, cleaned_design_df = %s' % (
            cleaned_measurement_df.shape,
            cleaned_design_df.shape
        )
        assert cleaned_measurement_df.shape[1] == cleaned_design_df.shape[0], msg

        self.initial_measurement_df = cleaned_measurement_df
        self.design_df = cleaned_design_df
        self.feature_annot_df = feature_annot_df

        # An ordered dict of dataframes as it's processed through the pipeline.
        # The ordering is important since data is processed sequentially
        # The last entry is the current result to use
        self.processed_dfs = OrderedDict()

    @property
    def data_df(self):
        od = self.processed_dfs
        if len(od) > 0:
            return od[next(reversed(od))]
        else:
            return self.initial_measurement_df

    def get_initial_measurement_df(self):
        return self.initial_measurement_df.copy()

    def _get_significant_df(self, df):
        # old data in GraphOmics was using PADJ_COL_PREFIX rather than PVALUE_COL_PREFIX
        df = df.copy()
        df.columns = df.columns.str.replace(PADJ_COL_PREFIX, PVALUE_COL_PREFIX)

        # find significant information in columns starting with FC_COL_PREFIX, PVALUE_COL_PREFIX
        drop_cols = [FC_COL_PREFIX, PVALUE_COL_PREFIX]
        drop_cols = tuple([x.lower() for x in drop_cols])
        to_drop = list(filter(lambda x: x.lower().startswith(drop_cols), df.columns))
        to_keep = [col for col in df.columns if col not in to_drop]

        keep_df = df[to_keep]
        drop_df = df[to_drop]
        return keep_df, drop_df

    def _clean_data(self, measurement_df, design_df):

        # drop duplicate rows and columns by values
        # measurement_df = self._drop_dupes_by_values(measurement_df, MEASUREMENT_DF_LABEL)
        # design_df = self._drop_dupes_by_values(design_df, DESIGN_DF_LABEL) # don't do this!

        # drop duplicate rows and columns by sample names
        measurement_df = self._drop_dupes_by_colnames(measurement_df, MEASUREMENT_DF_LABEL)
        design_df = self._drop_dupes_by_colnames(design_df.transpose(), DESIGN_DF_LABEL).transpose()

        # keep common samples having both measurements and metadata
        measurement_df, design_df = self._keep_common_samples(measurement_df, design_df)
        return measurement_df, design_df

    def _drop_dupes_by_values(self, df, label):
        # drop duplicate rows, keep no duplicates
        no_dupe_rows = df.drop_duplicates(keep='first')

        # drop duplicate columns, keep no duplicates
        no_dupe = no_dupe_rows.transpose().drop_duplicates(keep='first').transpose()

        # print message is something has been dropped
        if df.shape != no_dupe.shape:
            logger.warning('Dropped duplicate from %s by values: %d rows and %d cols' % (
                label,
                df.shape[0] - no_dupe.shape[0],
                df.shape[1] - no_dupe.shape[1]
            ))
        return no_dupe

    def _drop_dupes_by_colnames(self, df, label):
        # find columns that have the same name
        # https://stackoverflow.com/questions/14984119/python-pandas-remove-duplicate-columns
        cleaned_df = df.loc[:, ~df.columns.duplicated()]

        n_cols_initial = df.shape[1]
        n_cols_cleaned = cleaned_df.shape[1]
        diff = n_cols_initial - n_cols_cleaned
        if diff > 0:
            logger.warning('Dropped %d duplicate sample names from %s' % (diff, label))
        return cleaned_df

    def _keep_common_samples(self, measurement_df, design_df):
        # find common sample names (rows and columns) in measurement and design dfs
        cols = measurement_df.columns.values
        rows = design_df.index.values
        common = set(cols).intersection(set(rows))

        # select the common row and col names
        selected_cols = [col for col in cols if col in common]
        selected_rows = [row for row in rows if row in common]
        cleaned_measurement_df = measurement_df[selected_cols]
        cleaned_design_df = design_df.loc[selected_rows]

        diff = measurement_df.shape[1] - cleaned_measurement_df.shape[1]
        if diff > 0:
            logger.warning('Dropped %d columns from measurement dataframe due to missing metadata' % diff)

        diff = design_df.shape[0] - cleaned_design_df.shape[0]
        if diff > 0:
            logger.warning('Dropped %d columns from sample metadata due to missing measurements' % diff)

        return cleaned_measurement_df, cleaned_design_df

    def __repr__(self):
        dtype_str = self.data_type
        shape = self.data_df.shape
        return '%s data with (%d, %d) measurements' % (dtype_str, shape[0], shape[1])


class MultiOmicsData():
    def __init__(self, publication=None, url=None):
        self.views = {}
        self.publication = publication
        self.url = url

    def add_data(self, omics_data):
        # check if list
        if not isinstance(omics_data, list):
            # if not, put it in a list
            items = [omics_data]
        else:
            items = omics_data

        # push list items to dictionary
        for item in items:
            view_name = item.data_type
            self.views[view_name] = item

    def has_data(self, data_type):
        return data_type in self.views

    def get_data(self, data_type):
        try:
            return self.views[data_type]
        except KeyError:
            return None

    def get_dfs(self, data_type):
        if self.has_data(data_type):
            data = self.get_data(data_type)
            return data.data_df, data.design_df
        else:
            return None, None

    def get_info(self, entity_id, data_type):
        return get_info(entity_id, data_type)

    def to_mofa(self):
        res = pd.DataFrame()
        for v in self.views:
            data = self.views[v]
            df = data.data_df
            df['feature'] = df.index
            df = df.melt(id_vars='feature',var_name='sample')
            df['view'] = data.data_type
            df = df.join(data.design_df, on = 'sample')
            res = res.append(df)
        return res

    def __repr__(self):
        msg = 'Multi-omics data container'
        if self.publication is not None:
            msg += '\n- publication: %s' % self.publication
        if self.url is not None:
            msg += '\n- URL: %s' % self.url
        if len(self.views) > 0:
            msg += '\n- Views: %d modalities' % len(self.views)
            for v in self.views:
                msg += '\n\t - %s' % self.views[v]
        return msg