import pandas as pd
import scipy
import pylab
import numpy as np

from statsmodels.formula.api import OLS

from easydev import Progress, AttrDict

from gdsctools.stats import MultipleTesting
from gdsctools import readers
from gdsctools.boxplots import BoxPlots
from gdsctools.settings import ANOVASettings
from gdsctools.anova_results import ANOVAResults

__all__ = ['ANOVA']

# Not that Logging is not used: it is not pickable and prevent
# multicore analysis.
class ANOVA(object): #Logging):
    """ANOVA analysis of the IC50 vs Feature matrices

    This class is the core of the analysis. It can be used to
    compute

    #. One association between a drug and a feature
    #. The association**S** between a drug and a set of features
    #. All assocations between a set of deugs and a set of features.

    For instance here below, we read an IC50 matrix and compute the
    association for a given drug with a specific feature.

    Note that genomic features are not provided as input but a default
    file is provided with this package that contains 677 genomic
    features for 1001 cell lines. If your IC50 contains unknown cell lines,
    you can provide your own file.

    .. plot::
        :include-source:
        :width: 80%

        from gdsctools import IC50, ANOVA, ic50_test
        ic = IC50(ic50_test)
        an = ANOVA(ic)
        # This is to select a specific tissue
        an.set_cancer_type('breast')
        df = an.anova_one_drug_one_feature('Drug_1047_IC50',
            'TP53_mut', show=True)

    :Details about the anova analysis: In the example above, we perform a
        regression/anova test based on OLS regression. This is done for
        one feature one drug across all cell lines (tissue) in the method
        :meth:`anova_one_drug`. The regression
        takes into account the following factors: tissue, MSI and features.
        The order matters. If there is only one tissue, this factor is
        dropped. If the number of MSI values is less than a pre-defined
        parameter (see :class:`~gdsctools.settings.ANOVASettings`), it is
        dropped. The other
        methods :meth:`anova_one_drug` and :meth:`anova_all` are wrappers
        around :meth:`anova_one_drug_one_feature` to loop over all drugs, and
        loop over all drugs and all features, respectively.

    V17 :
        gdsc.volcano_FDR_interpolation = False
        gdsc.settings.pvalue_correction_method = 'qvalue'

    """
    def __init__(self, ic50, genomic_features=None,
            drug_decode=None, verbose=True, low_memory=True,
            set_media_factor=False):
        """.. rubric:: Constructor

        :param DataFrame IC50: a dataframe with the IC50. Rows should be
            the COSMIC identifiers and columns should be the Drug names
            (or identifiers)
        :param features: another dataframe with rows as in the IC50 matrix
            and columns as features.  The first 3 columns must be named
            specifically to hold tissues, MSI (see format).
        :param drug_decode: a 3 column CSV file with drug's name and targets
            see :mod:`readers` for more information.
        :param verbose: verbosity in "WARNING", "ERROR", "DEBUG", "INFO"

        The attribute :attr:`settings` contains specific settings related
        to the analysis or visulation.
        """
        self.verbose = verbose

        # We first need to read the IC50 using a dedicated reader
        self.ic50 = readers.IC50(ic50)

        # Create a dictionary version of the data
        # to be accessed per drug where NA have already been
        # removed. Each drug is a dictionary with 2 keys:
        # Y for the data and indices for the cosmicID where
        # there is an IC50 measured.
        ic50_parse = self.ic50.df.copy().unstack().dropna()
        self.ic50_dict = dict([(d, {'indices': ic50_parse.ix[d].index,
            'Y':ic50_parse.ix[d].values}) for d in self.ic50.drugIds])

        # Reads features if provided, otherwise use a default data set
        if genomic_features is None:
            # Reads default version provided with the package
            self.features = readers.GenomicFeatures()
        else:
            self.features = readers.GenomicFeatures(genomic_features)

        if self.features.found_media is False and \
                set_media_factor is True:
            if self.verbose:
                print('Populating MEDIA Factor in the Genomic Feature matrix')
            self.features.fill_media_factor()

        #: a CSV with 3 columns used in the report
        self.read_drug_decode(drug_decode)

        # create the multiple testing factory used in anova_all()
        self.multiple_testing = MultipleTesting()

        # We prune the genomic features by settings the cosmic ids of
        # the features to be those of the cosmic ids of the IC50. See
        # readers module. This affectation, prune the features dataframe
        # automatically. This fails if a cosmic identifier is not
        # found in the features' cosmic ids, so let us catch the error
        # before hand to give a
        unknowns = set(self.ic50.cosmicIds).difference(
                set(self.features.cosmicIds))

        if len(unknowns) > 0:
            print("WARNING: " +
                "%s cosmic identifiers in your IC50 " % len(unknowns) +
                "could not be found in the genomic feature matrix. "+
                "They will be dropped. Consider using a user-defined " +
                "genomic features matrix")

        self.ic50.drop_cosmic(list(unknowns))
        self.features.cosmicIds = self.ic50.cosmicIds
        #self.cosmicIds = self.ic50.cosmicIds

        #: an instance of :class:`~gdsctools.settings.ANOVASettings`
        self.settings = ANOVASettings()
        self.settings.low_memory = low_memory

        # alias to all column names to store results
        # cast to list (Python3).
        self.column_names = list(ANOVAResults().mapping.keys())

        # skip assoc_id for now
        self._odof_dict = dict([(name, None)
            for name in self.column_names])

        # a cache to store ANOVA results for each drug
        self.individual_anova = {}

        # must be called if ic50 or features are changed.
        self.init()

    def _autoset_msi_factor(self):
        if self.features.found_msi:
            # if the number of pos. (or neg.) factors is not large enough then
            # the MSI factor is not used
            msi_name = self.features.colnames.msi
            self.msi_factor = self.features.df[msi_name]
            total = len(self.msi_factor)
            positives = self.msi_factor.sum()
            negatives = total - positives

            # we must have at least 2 positives or 2 negative
            # This is therefore a < comparison here below. See in
            # _get_one_drug_one_feature_data that we use >= which
            # is consistent.
            if positives < self.settings.MSI_factor_threshold:
                self.settings.include_MSI_factor = False
            if negatives < self.settings.MSI_factor_threshold:
                self.settings.include_MSI_factor = False
        else:
            self.settings.include_MSI_factor = False

    def _autoset_tissue_factor(self):
        # select tissue based on the features
        tissue_name = self.features.colnames.tissue
        self.tissue_factor = self.features.df[tissue_name]
        if len(self.tissue_factor.unique()) == 1:
            # there is only one tissue
            tissue = self.tissue_factor.unique()[0]
            self.settings.analysis_type = tissue
            self.settings.directory = tissue
        else:
            # this is a PANCAN analysis
            self.settings.analysis_type = 'PANCAN'

    def _autoset_media_factor(self):

        if self.settings.analysis_type != 'PANCAN':
            self.settings.include_media_factor = False
        elif self.features.found_media is True:
            self.settings.include_media_factor = True
            colname = self.features.colnames.media
            self.media_factor = self.features.df[colname]
        else:
            self.settings.include_media_factor = False

    def set_cancer_type(self, ctype=None):
        """Select only a set of tissues.

        Input IC50 may be PANCAN (several cancer tissues).
        This  function can be used to select a subset of tissues.
        This function changes the :attr:`ic50` dataframe and possibly
        the feature as well if some are not relevant anymore (sum of the
        column is zero for instance).

        """
        if ctype is None:
            return

        if ctype == 'PANCAN':
            return

        if isinstance(ctype, str):
            ctype = [str(ctype)]

        for this in ctype:
            assert this in self.features.tissues, "%s not found" % ctype

        # keep only features that correspond to the tissue
        self.features.keep_tissue_in(ctype)

        self.ic50.df = self.ic50.df.ix[self.features.df.index]
        self.init()

    def read_settings(self, settings):
        """Read settings and update cancer type if set"""
        self.settings.from_json(settings)
        self.set_cancer_type(self.settings.analysis_type)

    def init(self):
        # Some preprocessing to speed up data access
        ic50_parse = self.ic50.df.copy().unstack().dropna()
        # for each drug, we store the IC50s (Y) and corresponding indices
        # of cosmic identifiers
        self.ic50_dict = dict([
            (d, {'indices': ic50_parse.ix[d].index,
             'Y': ic50_parse.ix[d].values}) for d in self.ic50.drugIds])

        # save the tissues
        self._autoset_tissue_factor()

        # and MSI (Microsatellite instability) status of the samples.
        self._autoset_msi_factor()

        # and (growth) media factor
        self._autoset_media_factor()

        # dictionaries to speed up code.
        self.features_dict = {}
        self.msi_dict = {}
        self.tissue_dict = {}
        self.media_dict = {}
        # fill the dictionaries for each drug once for all
        for drug_name in self.ic50.drugIds:
            indices = self.ic50_dict[drug_name]['indices']
            # if we were to store all drugs /features, this takes
            # 1Gb of memory for 265 drugs and 680 features. This is
            # therefore not scalable, especially for multiprocessing.
            if self.settings.low_memory is True:
                pass
            else:
                self.features_dict[drug_name] = self.features.df.ix[indices]

            # MSI, media and tissue are not large data files and can be store
            if self.features.found_msi:
                self.msi_dict[drug_name] = self.msi_factor.ix[indices]
            if self.features.found_media:
                self.media_dict[drug_name] = self.media_factor.ix[indices]

            self.tissue_dict[drug_name] = self.tissue_factor.ix[indices]

        # some preprocessing for the OLS computation.
        # We create the dummies for the tissue factor once for all
        # Note that to agree with R convention, we have to resort the column
        # to agree with R convention that is a<B==b<c instead of
        # where A<B<C<a<b<c (in R)
        self._tissue_dummies = pd.get_dummies(self.tissue_factor)
        columns = self._tissue_dummies.columns
        columns = sorted(columns, key=lambda s: s.lower())
        columns = ['C(tissue)[T.' + x + ']' for x in columns]
        self._tissue_dummies.columns = columns

        if self.settings.include_media_factor:
            self._media_dummies = pd.get_dummies(self.media_factor)
            columns = self._media_dummies.columns
            columns = ['C(media)[T.' + x + ']' for x in columns]
            self._media_dummies.columns = columns
            for col in columns:
                self._tissue_dummies[col] = self._media_dummies[col]

        N = len(self._tissue_dummies)
        self._tissue_dummies['C(msi)[T.1]'] = [1]*N
        self._tissue_dummies['feature'] = [1] * N
        self._tissue_dummies.insert(0, 'Intercept', [1] * N)

        # drop first feature in the tissues that seems to be used as a
        # reference in the regression
        tissues = [x for x in self._tissue_dummies.columns if 'tissue' in x]
        self._tissue_dummies.drop(tissues[0], axis=1, inplace=True)

        if self.settings.include_media_factor:
            media = [x for x in self._tissue_dummies.columns if 'media' in x]
            self._tissue_dummies.drop(media[0], axis=1, inplace=True)

        # reset the buffer.
        self.individual_anova = {}

        for this in ['tissue', 'media', 'msi', 'feature']:
            if this in self._get_analysis_mode():
                print(this.upper() + " FACTOR : included")
            else:
                print(this.upper() + " FACTOR : NOT included")

    def _get_cosmics(self):
        return self.ic50.cosmicIds
    def _set_cosmics(self, cosmics):
        self.ic50.cosmicIds = cosmics
        self.features.cosmicIds = cosmics
        self.init()
        self.individual_anova = {}
    cosmicIds = property(_get_cosmics, _set_cosmics,
        doc="get/set the cosmic identifiers in the IC50 and feature matrices")

    def _get_drug_names(self):
        return self.ic50.drugIds
    def _set_drug_names(self, drugs):
        self.ic50.drugIds = drugs
        self.init()
        # not need to init this again ? self.individual_anova = {}
    drugIds = property(_get_drug_names, _set_drug_names,
            doc="Get/Set drug identifers")

    def _get_feature_names(self):
        shift = self.features.shift
        return self.features.features[shift:]
    def _set_features_names(self, features):
        self.features.features = features
        self.init()
        self.individual_anova = {}
    feature_names = property(_get_feature_names, _set_features_names,
            doc="Get/Set feature names")

    def _get_analysis_mode(self):
        modes = []
        if self.settings.analysis_type == 'PANCAN':
            modes.append('tissue')

        if self.settings.include_MSI_factor is True:
            modes.append('msi')

        if self.settings.include_media_factor is True:
            modes.append('media')

        modes.append('feature')
        return modes

    def diagnostics(self):
        """Return dataframe with information about the analysis

        """
        n_drugs = len(self.ic50.drugIds)
        n_features = len(self.features.features) - self.features.shift
        n_combos = n_drugs * n_features
        feasible = 0
        pb = Progress(n_drugs, 1)
        counter = 0
        for drug in self.ic50.drugIds:
            for feature in self.features.features[self.features.shift:]:
                dd = self._get_one_drug_one_feature_data(drug, feature,
                        diagnostic_only=True)
                if dd.status is True:
                    feasible += 1
            counter += 1
            pb.animate(counter)

        results = {
                'n_drug': n_drugs,
                'n_combos': n_combos,
                'feasible_tests': feasible,
                'percentage_feasible_tests': float(feasible)/n_combos*100}
        return results

    def _get_one_drug_one_feature_data(self, drug_name, feature_name,
            diagnostic_only=False):
        """
        return: a dictionary with relevant information. There is also
            a test to see if the data can be analysis or not. This is
            stored ad a boolean value with key called *status*.
        """
        # dictionary  struture to hold results (can set values as attributes)
        dd = AttrDict()

        # select IC50 of a given drug
        # a fast way to select non-NA values from 1 column:
        # dropna is actually faster than a method using a mask.
        #dd.Y = self.ic50.df[drug_name].dropna()
        #indices = dd.Y.index
        #dd.masked_features = self.features.df[feature_name][indices]
        #dd.masked_tissue = self.tissue_factor[indices]
        #dd.masked_msi = self.msi_factor[indices]
        #dd.positive_feature = dd.masked_features.values.sum()
        #dd.negative_feature = len(dd.masked_features) - dd.positive_feature
        #dd.positive_msi = dd.masked_msi.values.sum()
        #dd.negative_msi = len(dd.masked_msi) - dd.positive_msi
        # using a mask instead of indices is 30% slower
        #mask = self.ic50.df[drug_name].isnull()==False
        #dd.masked_features = self.features.df[feature_name][mask]
        #dd.masked_tissue = self.tissue_factor[mask]
        #dd.masked_msi = self.msi_factor[mask]

        # Amother version using a dictionary instead of dataframer is actually
        # 2-3 times faster. It requires to transform the dataframe into a
        # dictionary once for all and dropping the NA as well.
        # Now, the next line takes no time
        dd.Y = self.ic50_dict[drug_name]['Y']

        # an alias to the indices
        indices = self.ic50_dict[drug_name]['indices']

        # select only relevant tissues/msi/features
        if self.settings.low_memory is True:
            # This line takes 50% of the time
            dd.masked_features = self.features.df.loc[indices, feature_name]
        else:
            dd.masked_features = self.features_dict[drug_name][feature_name]

        dd.masked_tissue = self.tissue_dict[drug_name]
        if self.features.found_msi:
            dd.masked_msi = self.msi_dict[drug_name]
            dd.positive_msi = dd.masked_msi.values.sum()
            dd.negative_msi = len(dd.masked_msi) - dd.positive_msi

        if self.features.found_media:
            dd.masked_media = self.media_dict[drug_name]

        # compute length of pos/neg features and MSI
        dd.positive_feature = dd.masked_features.values.sum()
        dd.negative_feature = len(dd.masked_features) - dd.positive_feature

        # Some validity tests to run the analysis or not
        feature_threshold = self.settings.feature_factor_threshold
        msi_threshold = self.settings.MSI_factor_threshold

        A = self.settings.include_MSI_factor and\
            dd.positive_feature >= feature_threshold and\
            dd.negative_feature >= feature_threshold and\
            dd.negative_msi >= msi_threshold and \
            dd.positive_msi >= msi_threshold

        B = (not self.settings.include_MSI_factor) and\
            dd.positive_feature >= feature_threshold and\
            dd.negative_feature >= feature_threshold
        # We could of course use the mean() and std() functions from pandas or
        # numpy. We could also use the glass and cohens functions from the
        # stats module but the following code is much faster because it
        # factorises the computations of mean and variance
        dd.positives = dd.Y[dd.masked_features.values == 1]
        dd.negatives = dd.Y[dd.masked_features.values == 0]
        dd.Npos = len(dd.positives)
        dd.Nneg = len(dd.negatives)

        # additional information
        dd.feature_name = feature_name
        dd.drug_name = drug_name

        # FIXME is False does not give the same results as == False
        # in the test test_anova.py !!
        if (A == False) and (B == False):
            dd.status = False
            return dd
        else:
            dd.status = True

        if diagnostic_only is True:
            return dd

        # compute mean and std of pos and neg sets;using mean() takes 15us and
        # using the already computed sum and N takes 5us
        pos_sum = dd.positives.sum()
        neg_sum = dd.negatives.sum()
        dd.pos_IC50_mean = pos_sum / dd.Npos
        dd.neg_IC50_mean = neg_sum / dd.Nneg
        dd.delta_mean_IC50 = dd.pos_IC50_mean - dd.neg_IC50_mean

        # note the ddof to agree with R convention.
        dd.pos_IC50_std = dd.positives.std(ddof=1)
        dd.neg_IC50_std = dd.negatives.std(ddof=1)

        dd.pos_IC50_std = np.sqrt(( (dd.positives**2).sum() -
            pos_sum**2/dd.Npos)/(dd.Npos-1.))
        dd.neg_IC50_std = np.sqrt(( (dd.negatives**2).sum() -
            neg_sum**2/dd.Nneg)/(dd.Nneg-1.))

        # Compute Cohens and Glass effect size. Since underlying code
        # has lots in common, we do not use the modules but add
        # the code here below
        md = np.abs(dd.pos_IC50_mean - dd.neg_IC50_mean)
        dd.pos_glass = md / dd.pos_IC50_std
        dd.neg_glass = md / dd.neg_IC50_std

        csd = (dd.Npos - 1.) * dd.pos_IC50_std**2 + \
                (dd.Nneg - 1.) * dd.neg_IC50_std**2
        csd /= dd.Npos + dd.Nneg - 2.  # make sure this is float
        dd.effectsize_ic50 = md / np.sqrt(csd)

        # Note that equal_var is a user parameter and affects
        # results. The ANOVA_results.txt obtained from SFTP
        # have different values meaning that the equal.var param
        # was set to False. Note that pvalue is stored at index 1
        dd.ttest = self._get_ttest(dd.negatives, dd.positives)
        return dd

    def _get_ttest(self, sample1, sample2):
        # this computes the ttest.
        import scipy
        return scipy.stats.ttest_ind(sample1, sample2,
                equal_var=self.settings.equal_var_ttest)[1]

    def read_drug_decode(self, filename=None):
        """Read file with the DRUG information

        .. seealso:: :class:`gdsctools.readers.DrugDecode`
        """
        # Read the DRUG decoder file into a DrugDecode/Reader instance
        self.drug_decode = readers.DrugDecode(filename)

    def anova_one_drug_one_feature(self, drug_id,
            feature_name, show=False,
            production=False, directory='.'):
        """Compute ANOVA and various tests on one drug and one feature

        :param drug_id: a valid drug identifier
        :param feature_name: a valid feature name
        :param bool show: show some plots
        :param str directory: where to save the figure.
        :param bool production: if False, returns a dataframe otherwise
            a dictionary. This is to speed up analysis when scanning
            the drug across all features.

        .. note:: **for developer** this is the core of tha analysis
            and should be kept as fast as possible. 95% of the time is spent
            here.

        .. note:: **for developer** Data used in this function comes from
            _get_one_drug_one_feature_data method, which should also be kept
            as fast as possible.
        """
        if drug_id not in self.drugIds:
            raise ValueError('Unknown drug name %s. Use e.g., %s'
                    % (drug_id, self.drugIds[0]))

        if feature_name not in self.feature_names:
            # we start index at 3 to skip tissue/name/msi
            raise ValueError('Unknown feature name %s. Use e.g. one of %s'
                    % (feature_name, self.feature_names[0:3]))

        # This extract the relevant data and some simple metrics
        # This is now pretty fast accounting for 45 seconds
        # for 265 drugs and 988 features
        odof = self._get_one_drug_one_feature_data(drug_id, feature_name)
        drug_name = self.drug_decode.get_name(drug_id)
        drug_target = self.drug_decode.get_target(drug_id)

        # if the status is False, it means the number of data points
        # in a category (e.g., positive feature) is too low.
        # If so, nothing to do, we return an 'empty' dictionary
        if odof.status is False:
            results = self._odof_dict.copy()
            results['FEATURE'] = feature_name
            results['DRUG_ID'] = drug_id
            results['DRUG_NAME'] = drug_name
            results['DRUG_TARGET'] = drug_target
            results['N_FEATURE_pos'] = odof.Npos
            results['N_FEATURE_neg'] = odof.Nneg
            if production is True:
                # return a dict
                return results
            else:
                # or a dataframe; note that index is not relevant here but
                # required.
                df = pd.DataFrame(results, index=[1])
                return df

        # with the data extract, we can now compute the regression.

        # In R or statsmodels, the regression code is simple since
        # it is based on the formula notation (Y~C(msi)+feature)
        # This is also possible in statsmodels library,  however,
        # this relies on patsy, which is very slow as compared to the
        # statsmodels without formula.
        #### self._mydata = pd.DataFrame({'Y':self.Y,
        ####    'tissue':self.masked_tissue,
        ####       'msi': self.masked_msi, 'feature':self.masked_features})
        #### self.data_lm = ols('Y ~ C(tissue) + C(msi) + feature',
        ####  data=self._mydata, missing='none').fit() #Specify C is category

        # IMPORTANT: the order of the factors in the formula
        # is important. It does not change the total sum of square errors
        # but may change individual effects of the categorical
        # components.

        # Instead of using ols function, we use the OLS one so we cannot
        # use formula. Instead, we need to create manually the input
        # data. In the case of categorical data (tissue), we need to
        # create the dummy variable, which is done in the constructor
        # once for all (slow otherwise).
        if self.settings.analysis_type == 'PANCAN':
            # IMPORTANT: tissues are sorted alphabetically in R aov
            # function. Same in statsmodels but capitalised names
            # are sorted differently. In R, a<b<B<c but in Python,
            # A<B<C<a<b<c. So, 'aero' tissue is before 'Bladder' in R,
            # not in python. Since in a linear regression
            # models, the order of the factor matters and the first
            # factor is used as a reference, we decided to use same
            # convention as in R.
            # see http://statsmodels.sourceforge.net/devel/contrasts.html
            # for a good explanation

            #self._mydata = pd.DataFrame({'Y': odof.Y.copy(),
            #    'tissue':odof.masked_tissue,
            #    'msi':  odof.masked_msi, 'feature': odof.masked_features})
            #self.data_lm2 = ols('Y ~ C(tissue) + C(msi) + feature',
            #    data=self._mydata).fit() #Specify C for Categorical

            # from statsmodels.stats.anova import anova_lm
            # import statsmodels.formula.api as smf
            # df  = pd.DataFrame({'Y': odof.Y.copy(),
            #   'tissue':odof.masked_tissue,'media'
            #    odof.masked_media, 'msi':  odof.masked_msi,
            #   'feature': odof.masked_features})
            # lm = smf.ols('Y~C(tissue)+C(media)+C(msi)+feature',
            #    data=df).fit()
            #  anova_lm(lm)
            # The code above gives same answer as the code in gdsctools
            # but is slower

            # We could use pd.get_dummies but pretty slow
            # instead we create the full matrix in init() method.
            # One issue is that some columns end up with sum == 0
            # and needs to be dropped.
            df = self._tissue_dummies.ix[odof.masked_tissue.index]
            todrop = df.columns[df.values.sum(axis=0) == 0]
            if len(todrop) > 0: # use if since drop() is slow
                df = df.drop(todrop, axis=1)

            # Here we set other variables with dataframe columns' names as
            # expected by OLS.
            if self.settings.include_media_factor == False:
                todrop = [x for x in df.columns if
                        x.startswith('C(media)')]
                df = df.drop(todrop, axis=1)

            df['C(msi)[T.1]'] = odof.masked_msi.values
            df['feature'] = odof.masked_features.values

            self.Y = odof.Y
            self.EV = df.values
            # The regression and anova summary are done here
            #
            if self.settings.regression_method == 'ElasticNet':
                self.data_lm = OLS(odof.Y, df.values).fit_regularized(
                        alpha=self.settings.regression_alpha,
                        L1_wt=self.settings.regression_L1_wt)
            elif self.settings.regression_method == 'OLS':
                self.data_lm = OLS(odof.Y, df.values).fit()
            elif self.settings.regression_method == 'Ridge':
                self.data_lm = OLS(odof.Y, df.values).fit_regularized(
                        alpha=self.settings.regression_alpha,
                        L1_wt=0)
            elif self.settings.regression_method == 'Lasso':
                self.data_lm = OLS(odof.Y, df.values).fit_regularized(
                        alpha=self.settings.regression_alpha,
                        L1_wt=1)

            # example of computing null model ?
            # Example of computing pvalues ourself
            """self.samples1 = []
            self.samples2 = []
            self.samples3 = []
            Y = odof.Y.copy()
            pb = Progress(10000,20)
            for i in range(0,10000):
                #pylab.shuffle(Y)
                #data_lm = OLS(Y, df.values).fit()
                data_lm = OLS(Y+0.3*pylab.randn(len(Y)), df.values).fit()
                anova_pvalues = self._get_anova_summary(data_lm,
                    output='dict')
                self.samples1.append(anova_pvalues['msi'])
                self.samples2.append(anova_pvalues['feature'])
                self.samples3.append(anova_pvalues['tissue'])
                pb.animate(i)
            """

        elif self.settings.include_MSI_factor is True:
            #self._mydata = pd.DataFrame({'Y': odof.Y,
            #    'msi':  odof.masked_msi, 'feature': odof.masked_features})
            #self.data_lm = ols('Y ~ C(msi) + feature',
            #    data=self._mydata).fit() #Specify C for Categorical
            df = pd.DataFrame()
            df['C(msi)[T.1]'] = odof.masked_msi.values
            df['feature'] = odof.masked_features.values
            df.insert(0, 'Intercept', [1] * (odof.Npos + odof.Nneg))
            self.data_lm = OLS(odof.Y, df.values).fit()
        else:
            df = pd.DataFrame()
            df['feature'] = odof.masked_features.values
            df.insert(0, 'Intercept', [1] * (odof.Npos + odof.Nneg))
            self.data_lm = OLS(odof.Y, df.values).fit()
            #self._mydata = pd.DataFrame({'Y': odof.Y,
            #    'feature': odof.masked_features})
            #self.data_lm = ols('Y ~ feature',
            #    data=self._mydata).fit() #Specify C for Categorical

        self.anova_pvalues = self._get_anova_summary(self.data_lm,
                 output='dict')

        # Store the pvalues. Note that some may be missing so we use try
        # except, which is faster than if/else
        try:
            tissue_PVAL = self.anova_pvalues['tissue']
        except:
            tissue_PVAL = None

        try:
            MSI_PVAL = self.anova_pvalues['msi']
        except:
            MSI_PVAL = None

        try:
            FEATURE_PVAL = self.anova_pvalues['feature']
        except:
            FEATURE_PVAL = None

        try:
            MEDIA_PVAL = self.anova_pvalues['media']
        except:
            MEDIA_PVAL = None

        if show is True:
            boxplot = BoxPlots(odof, savefig=self.settings.savefig,
                    directory=directory)
            boxplot.boxplot_association(fignum=1)

            # a boxplot to show cell lines effects. This requires
            # the settings.analyse_type to be PANCAN
            if self.settings.analysis_type == 'PANCAN':
                boxplot.boxplot_pancan(fignum=2, mode='tissue')
            if self.settings.include_MSI_factor:
                boxplot.boxplot_pancan(fignum=3, mode='msi')

        results = {'FEATURE': feature_name,
                'DRUG_ID': drug_id,
                'DRUG_NAME': drug_name,
                'DRUG_TARGET': drug_target,
                'N_FEATURE_pos': odof.Npos,
                'N_FEATURE_neg': odof.Nneg,
                'FEATURE_pos_logIC50_MEAN': odof.pos_IC50_mean,
                'FEATURE_neg_logIC50_MEAN': odof.neg_IC50_mean,
                'FEATURE_delta_MEAN_IC50': odof.delta_mean_IC50,
                'FEATURE_pos_IC50_sd': odof.pos_IC50_std,
                'FEATURE_neg_IC50_sd': odof.neg_IC50_std,
                'FEATURE_IC50_effect_size': odof.effectsize_ic50,
                'FEATURE_pos_Glass_delta': odof.pos_glass,
                'FEATURE_neg_Glass_delta': odof.neg_glass,
                'ANOVA_FEATURE_pval': FEATURE_PVAL,
                'ANOVA_TISSUE_pval': tissue_PVAL,
                'ANOVA_MSI_pval': MSI_PVAL,
                'ANOVA_MEDIA_pval': MEDIA_PVAL,
                'FEATURE_IC50_T_pval': odof.ttest # pvalues is in index 1
                }

        # 12% of the time here
        if production is True:
            return results
        else:
            df = pd.DataFrame(results, index=[1])
            return df

    def optimise_elastic_net(self, drug_name, feature_name, N=20, Nalpha=20):
        lwts = pylab.linspace(0, 1, N)
        alphas = pylab.linspace(0, 5, Nalpha)

        mses = np.zeros((N, Nalpha))

        pb = Progress(N)
        for i, lwt in enumerate(lwts):
            for j, alpha in enumerate(alphas):
                self.settings.regression_method = 'ElasticNet'
                self.settings.regression_alpha = alpha
                self.settings.regression_L1_wt = lwt
                odof = self.anova_one_drug_one_feature(drug_name,
                        feature_name)
                anova = self._get_anova_summary(self.data_lm,
                        output='dataframe')
                mses[i,j] = self.data_lm.bic
            pb.animate(i+1)
        return mses

    def optimise_ridge(self, drug_name, feature_name, alphas=None):
        return self._opt_ridge_lasso(drug_name, feature_name,
                'Ridge', alphas=alphas)

    def optimise_lasso(self, drug_name, feature_name, alphas=None):
        return self._opt_ridge_lasso(drug_name, feature_name,
                'Lasso', alphas=alphas)

    def _opt_ridge_lasso(self, drug_name, feature_name, method, alphas=None):

        if alphas is None:
            alphas = pylab.linspace(0,1, 20)

        mses = []
        params = []
        method_buf = self.settings.regression_method
        alpha_buf = self.settings.elastic_net.alpha

        pb = Progress(len(alphas))
        for j, alpha in enumerate(alphas):
            self.settings.regression_method = method
            self.settings.elastic_net.alpha = alpha
            odof = self.anova_one_drug_one_feature(drug_name,
                    feature_name)
            anova = self._get_anova_summary(self.data_lm,
                    output='dataframe')
            #mses.append(anova.ix['Residuals']['Sum Sq'])
            mses.append(anova.ix['tissue']['F value'])
            #mses.append(anova['Sum Sq'].sum())
            pb.animate(j+1)
            params.append(self.data_lm.params)
        self.settings.regression_method = method_buf
        self.settings.elastic_net.alpha = alpha_buf
        return alphas, mses, params

    # no need to optimise anymore
    def _get_anova_summary(self, data_lm, output='dict'):
        # could use this with statsmodels but somehow anova_lm with typ I
        # does not work, which is the one used in R version, so we implement
        # the anova here
        q, r = np.linalg.qr(data_lm.model.data.exog)
        effects = np.dot(q.T, data_lm.model.data.endog)

        # In the regression, the first tissue is dropped hence -1
        # The degree of freedom for tissues is N - 1
        # self.features.tissues contains all tissues even those that
        # were dropped due to lack of pos or neg features. So, we must use
        modes = self._get_analysis_mode()
        Ncolumns = data_lm.model.data.exog.shape[1]
        Ntissue = Ncolumns
        Ntissue -= 1 # remove intercept
        Ntissue -= 1 # remove feature, which are always taken into account
        if 'msi' in modes:
            Ntissue -= 1
        if 'media' in modes:
            Nmedia = len(self._media_dummies.columns)-1
        else:
            Nmedia = 0
        Ntissue -= Nmedia

        # create the W matrix using tissue and MSI if requested
        # default is that the 3 features are used
        if 'tissue' in modes and 'msi' in modes and 'media' in modes:
            dof = [Ntissue, Nmedia, 1, 1]
            indices = ['tissue', 'media', 'msi', 'feature', 'Residuals']
            # 4 stands for intercept + tissue + msi +feature
            arr = np.zeros((5, Ncolumns))
            arr[1, slice(1, Ntissue+1)] = 1
            arr[2, slice(Ntissue+1, Ntissue+Nmedia+1)] = 1
            arr[3, Ntissue + Nmedia + 1] = 1
            arr[4, Ntissue + Nmedia + 2] = 1
            self.arr = arr
        elif 'tissue' in modes and 'msi' in modes:
            dof = [Ntissue, 1, 1]
            indices = ['tissue', 'msi', 'feature', 'Residuals']
            # 4 stands for intercept + tissue + msi +feature
            arr = np.zeros((4, Ncolumns))
            arr[1, slice(1, Ntissue)] = 1
            arr[2, Ntissue + 1] = 1
            arr[3, Ntissue + 2] = 1
            self.arr = arr
        elif 'tissue' not in modes and 'msi' in modes:
            dof = [1, 1]
            indices = ['msi', 'feature', 'Residuals']
            # 3 stands for intercept + msi +feature
            arr = np.zeros((3, Ncolumns))
            arr[1, 1] = 1
            arr[2, 2] = 1
        elif 'tissue' not in modes and 'msi' not in modes:
            dof = [1]
            indices = ['feature', 'Residuals']
            # 3 stands for intercept + msi +feature
            arr = np.zeros((2, Ncolumns))
            arr[1, 1] = 1
        arr[0, 0] = 1                   # intercept

        sum_sq = np.dot(arr, effects**2)[1:] # drop the intercep
        mean_sq = sum_sq / np.array(dof)
        Fvalues = mean_sq / (data_lm.ssr / data_lm.df_resid)
        F_pvalues = scipy.stats.f.sf(Fvalues, dof, data_lm.df_resid)
        sum_sq = np.append(sum_sq, data_lm.ssr)
        mean_sq = np.append(mean_sq, data_lm.mse_resid)
        F_pvalues = np.append(F_pvalues, None)
        Fvalues = np.append(Fvalues, None)
        dof.append(data_lm.model.df_resid)
        #indices.append('Residuals')
        # dataframe is slow, return just the dict of pvalues by default
        if output == 'dataframe':
            anova = pd.DataFrame({'Sum Sq': sum_sq, 'Mean Sq': mean_sq,
                'Df': dof, 'F value': Fvalues, 'PR(>F)': F_pvalues},
                index=indices,
                columns=['Df', 'Sum Sq', 'Mean Sq', 'F value', 'PR(>F)']
                )
            return anova
        elif self.settings.analysis_type == 'PANCAN':
            if self.settings.include_media_factor:
                dd = {'tissue': F_pvalues[0],
                      'media': F_pvalues[1],
                      'msi':F_pvalues[2],
                      'feature':F_pvalues[3]}
            else:
                dd = {'tissue': F_pvalues[0],
                      'msi':F_pvalues[1],
                      'feature':F_pvalues[2]}
            return dd
        elif self.settings.include_MSI_factor is True:
            return {'msi': F_pvalues[0], 'feature':F_pvalues[1]}
        else:
            return {'feature': F_pvalues[0]}

        #return anova

    def _draft(self):
        # using sklearn
        #ols = linear_model.LinearRegression()
        #f = ols.fit(an.dff, an.Y)
        #sse = sum(np.square((f.predict(an.dff).T - an.Y))) /
        #           float(an.dff.shape[0] - an.dff.shape[1])
        # ssr = sum(np.square((f.predict(an.dff).T - an.Y.mean())))
        pass

    def _test(self):
        # for drug1047 and featuer ABCB1_mut
        print("""
        Analysis of Variance Table

        Response: Y
        Df  Sum Sq Mean Sq F value  Pr(>F)
        TISSUEpattern  26  352.35 13.5517  9.2685 < 2e-16 ***
        MSIpattern      1    5.31  5.3094  3.6313 0.05705 .
        FEATpattern     1    3.19  3.1861  2.1791 0.14028
        Residuals     817 1194.55  1.4621
        """)

    #98% of time in  method anova_one_drug_one_feature
    def anova_one_drug(self, drug_id, animate=True, output='object'):
        """Computes ANOVA for a given drug across all features

        :param str drug_id: a valid drug identifier.
        :param animate: shows the progress bar
        :return: a dataframe

        Calls :meth:`anova_one_drug_one_feature` for each feature.

        """
        # some features can be dropped ??

        # drop first and second columns that are made of strings
        # works under python2 but not python 3. Assume that the 2 first
        #columns are the sample name and tissue feature
        # Then, we keep only cases with at least 3 features.
        # MSI could be used but is not like in original R code.
        features = self.features.df.copy()
        # need to skip the FACTOR to keep only features
        shift = self.features.shift

        features = features[features.columns[shift:]]
        mask = features.sum(axis=0) >= 3

        # TODO: MSI, tissues, name must always be kept
        #
        selected_features = features[features.columns[mask]]

        # scan all features for a given drug
        assert drug_id in self.ic50.df.columns
        N = len(selected_features.columns)
        pb = Progress(N, 10)
        res = {}
        # note that we start at idnex 4 to drop sample name, tissue and MSI
        for i, feature in enumerate(selected_features.columns):
            # production True, means we do not want to create a DataFrame
            # for each call to the anova_one_drug_one_feature function
            # Instead, we require dictionaries
            this = self.anova_one_drug_one_feature(drug_id, feature,
                    production=True)
            if this['ANOVA_FEATURE_pval'] is not None:
                res[feature] = this
            if animate is True:
                pb.animate(i+1)

        # if production is False:
        # df = pid.concat(res, ignore_index=True)
        df = pd.DataFrame.from_records(res)
        df = df.T

        df = ANOVAResults().astype(df)
        if len(df) == 0:
            return df

        # append DRUG_NAME/DRUG_TARGET columns
        df = self.drug_decode.drug_annotations(df)

        # TODO: drop rows where ANOVA_FEATURE_PVAL is None
        if output != 'object':
            df = self.add_pvalues_correction(df)
            return df
        else:
            df = self.add_pvalues_correction(df)

            res = ANOVAResults(df)
            res.settings = ANOVASettings(**self.settings)
            return res

    def anova_all(self, animate=True, drugs=None):
        """Run all ANOVA tests for all drugs and all features.

        :param drugs: you may select a subset of drugs
        :param animate: shows the progress bar
        :return: an :class:`~gdsctools.anova_results.ANOVAResults`
            instance with the dataframe
            stored in an attribute called **df**

        Loops over all drugs calling :meth:`anova_one_drug` for each
        drug and concatenating all results together. Note that once all
        data are gathered, an extra column containing the FDR corrections
        is added to the dataframe using :meth:`add_pvalues_correction`
        method. An extra column  named "ASSOC_ID" is also added with
        a unique identifer sorted by ascending FDR.

        .. note:: A thorough comparison with version v17 give the same FDR
            results (difference ~1e-6); Note however that the qvalue results
            differ by about 0.3% due to different smoothing in R and Python.
        """
        # drop DRUG where number of IC50 (non-null) is below 5
        # axis=0 is default but we emphasize that sum is over
        # column (i.e. drug
        vv = (self.ic50.df.isnull() == False).sum(axis=0)
        # FIXME: should be in one_drug_one_feature ??
        drug_names = vv.index[vv >= self.settings.minimum_nonna_ic50]

        # if user provided a list of drugs, use them:
        if drugs is not None:
            # todo: check valifity of the drug names
            drug_names = drugs[:]

        pb = Progress(len(drug_names), 1)
        drug_names = list(drug_names)
        pylab.shuffle(drug_names)
        if animate is True:
            pb.animate(0)

        for i, drug_name in enumerate(drug_names):
            if drug_name in self.individual_anova.keys():
                pass
            else:
                res = self.anova_one_drug(drug_name, animate=False,
                                          output='dataframe')
                self.individual_anova[drug_name] = res
            if animate is True:
                pb.animate(i+1)
        print("\n")
        if len(self.individual_anova) == 0:
            return ANOVAResults()

        df = pd.concat(self.individual_anova, ignore_index=True)

        if len(df) == 0:
            return df
        # sort all data by ANOVA p-values
        try:
            df.sort_values('ANOVA_FEATURE_pval', inplace=True)
        except:
            df.sort('ANOVA_FEATURE_pval', inplace=True)

        # all ANOVA have been computed individually for each drug and each
        # feature. Now, we need to compute the multiple testing corrections
        if self.settings.pvalue_correction_level == 'global':
            df = self.add_pvalues_correction(df)

        # insert a unique identifier as first column
        df.insert(0, 'ASSOC_ID', range(1, len(df) + 1))

        self.df = df
        # order the column names as defined in the __init__ method
        df = df[self.column_names]
        df.reset_index(inplace=True, drop=True)

        results = ANOVAResults()
        results.df = df
        results.settings = ANOVASettings(**self.settings)
        return results

    def add_pvalues_correction(self, df, colname='ANOVA_FEATURE_pval'):
        """Add the corrected pvalues column in a dataframe based on pvalues

        The default method (FDR correction) is stored in
        :attr:`settings.pvalue_correction_method` and can be changed to other
        methods (e.g., *qvalue*)

        .. seealso:: :meth:`anova_all`,
            :class:`~gdsctools.stats.MultipleTesting`
        """
        if len(df) == 0:
            return

        # extract pvalues
        data = df[colname].values

        # set the method and compute new pvalues
        self.multiple_testing.method = self.settings.pvalue_correction_method
        new_pvalues = self.multiple_testing.get_corrected_pvalues(data)
        new_pvalues *= 100
        # insert new columns.
        colname = 'ANOVA_FEATURE_FDR'

        try:
            df.insert(len(df.columns), colname, new_pvalues)
        except:
            # replaces it otherwise
            df[colname] = new_pvalues
        return df

    def reset_buffer(self):
        self.individual_anova = {}

    def __str__(self):
        txt = self.ic50.__str__()
        txt += "\n" + self.features.__str__()
        return txt

    def __repr__(self):
        txt = self.__str__()
        return txt


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import datetime as dt
from datetime import datetime

#%config InlineBackend.figure_format = 'retina'
#%matplotlib inline
#%%
path='C:/Users/Mr.Goldss/Desktop/health care digital ticket EDA/'
filename='7_23_2020 Facility Itemized Tracker (1)'
df=pd.read_excel(path+filename+'.xlsx',sheet_name='New ALL')

df.shape

#%%                     fix basic formatting: preprocessing
'''data preprocessing'''

def formatting_preprocessing(df):
    want_col=['Date', 'Facility Name',
              'Facility Type', 'Item', 'Product Source', 'Type', 'UOM',
              'Qty', 'Total']
    df=df[want_col]
    df.columns=['Delivered Date', 'Facility Name',
       'Facility Type', 'Item Description', 'Product Source', 'Type', 'Units',
       'QTY Per UOM', 'Total Amount']
    df=df[df['Delivered Date'].notna()]
    df=df[df['Item Description'].notna()]
    df=df[(df['Units']>0) & (df['Units'].notna())]
    df['Product Source']=df['Product Source'].fillna('SNS')
    df_clean=df
    return df_clean
df_clean=formatting_preprocessing(df)

def del_wrong_ems(df_clean):

    del_cond1=(df_clean['Delivered Date']>='2020-04-01') & (df_clean['Delivered Date']<='2020-04-30')
    del_cond2=(df_clean['Facility Type'] =='EMS') & (df_clean['Facility Name'] !='Phoenix Fire Resource Management')
    to_be_drop_EMS=df_clean.loc[del_cond1 & del_cond2]
    Ems_drop_list=to_be_drop_EMS.index
    df_clean_copy=df_clean
    df_clean_EMS=df_clean_copy.drop(Ems_drop_list)
    return df_clean_EMS
df_clean_EMS=del_wrong_ems(df_clean)

def deep_clean_str_manipulate(df_clean_EMS):

    df_clean_EMS['QTY Per UOM']=df_clean_EMS['QTY Per UOM'].fillna(1)
    df_clean_EMS['Type']=df_clean_EMS['Type'].str.replace('Surgical Masks','Surgical Mask')
    df_clean_EMS['Type']=df_clean_EMS['Type'].str.lower()
    df_clean_EMS.groupby('Type')['Total Amount'].sum()

    df_clean_EMS['Facility Type']=df_clean_EMS['Facility Type'].str.lower()
    df_clean_EMS['Facility Type']=df_clean_EMS['Facility Type'].str.strip()
    df_clean_EMS['Facility Type']=df_clean_EMS['Facility Type'].str.capitalize()

    to_replace={'Lctf':'Ltcf','Ltfc':'Ltcf'}
    for key in to_replace.keys():
        df_clean_EMS['Facility Type']=df_clean_EMS['Facility Type'].str.replace(key,to_replace[key])

    df_clean_EMS['Facility Type']=df_clean_EMS['Facility Type'].str.replace('Acute care','Acute')

    df_clean_EMS['Item Description']=df_clean_EMS['Item Description'].str.lower()

    s=[]
    m=[]
    l=[]
    other=[]

    for i in df_clean_EMS['Item Description']:

        if 'small' in i:
            s.append(i)
        elif 'medium' in i:
            m.append(i)
        elif 'large' in i:
            l.append(i)
        else:
            other.append(i)
    df_clean_EMS['size']=df_clean_EMS['Item Description'].apply(
        lambda x:x.replace(x,'small') if x in s
        else(x.replace(x,'medium') if x in m else( x.replace(x,'large') if x in l else x.replace(x,'universal')) ))

    df_clean_EMS['week_of_year']=df_clean_EMS['Delivered Date'].dt.strftime('%U')
    df_clean_EMS['Month']=df_clean_EMS['Delivered Date'].dt.month
    df_clean_EMS['Total Amount']=df_clean_EMS['Total Amount'].astype(int)
    return df_clean_EMS
deep_clean_str_manipulate(df_clean_EMS)
#%%
'''countplot'''
fig, ax1, = plt.subplots()
fig.set_size_inches(25, 12)
sns.countplot(x='Facility Type',data=df_clean_EMS,ax=ax1)
plt.xticks(fontsize=18,rotation=90)
plt.xlabel('Facility', fontsize=15)
plt.ylabel('count', fontsize=20)
plt.title('Beltmann Inventory Distribution', fontsize=30)

plt.tight_layout()
plt.show()
#%%
#,barplot with plt.text to label all columns'''
'''
for index, value in enumerate(df_clean_EMS['Total Amount']):
    plt.text(index,value+400, str(value),fontsize=12)
plt.tight_layout()
'''

# %%
'''sns.lineplot---->date as x axis'''
#we only care about 5 product type
type_needed=['n95','surgical mask','gowns','gloves']
fig, ax1, = plt.subplots()
fig.set_size_inches(25, 12)
sns.lineplot(x='Month',y='Total Amount',hue='Type',\
    data=df_clean_EMS[df_clean_EMS['Type'].isin(type_needed)],ax=ax1,\
        ci=False, markers=True, style='Type')
plt.xticks(fontsize=18,rotation=90)
plt.xlabel('Month', fontsize=20)
plt.ylabel('Amount', fontsize=25)
plt.title('day', fontsize=30)

plt.tight_layout()
plt.show()
#%%
type_needed=['n95','surgical mask','gowns','gloves']
fig, ax1, = plt.subplots()
fig.set_size_inches(25, 12)
sns.relplot(x='Month',y='Total Amount',hue='Type',\
    data=df_clean_EMS[df_clean_EMS['Type'].isin(type_needed)],ax=ax1,\
        ci=False, markers=True, style='Type',kind='line')
plt.xticks(fontsize=18,rotation=90)
plt.xlabel('Month', fontsize=20)
plt.ylabel('Amount', fontsize=25)
plt.title('day', fontsize=30)

plt.tight_layout()
plt.show()

# %%
'''sns.barplot()'''
#one thing need to realize is that the bar is not the sum but the mean

type_needed=['n95','surgical mask','gowns','gloves']
fig, ax1, = plt.subplots()
fig.set_size_inches(25, 12)
sns.barplot(x='Month',y='Total Amount',hue='Type',\
    data=df_clean_EMS[df_clean_EMS['Type'].isin(type_needed)],ax=ax1,\
        ci=False)
plt.yticks(fontsize=18)
plt.xticks(fontsize=20)
plt.xlabel('Month of the year', fontsize=20)
plt.ylabel('mean Amount', fontsize=25)
plt.title('day', fontsize=30)

plt.tight_layout()
plt.show()
# %%
'''switch x and y and use orient=h to transpose '''
type_needed=['n95','surgical mask','gowns','gloves']
fig, ax1, = plt.subplots()
fig.set_size_inches(25, 12)
sns.barplot(y='Month',x='Total Amount',hue='Type',\
    data=df_clean_EMS[df_clean_EMS['Type'].isin(type_needed)],ax=ax1,\
        ci=False,orient='h')
plt.yticks(fontsize=18)
plt.xticks(fontsize=20)
plt.xlabel('Month of the year', fontsize=20)
plt.ylabel('mean Amount', fontsize=25)
plt.title('day', fontsize=30)

plt.tight_layout()
plt.show()

# %%
'''how can i change the color?'''
#go to seaborn official web to find color option
# search for html color code--> find the (hex) code like this :#E315C7
type_needed=['n95','surgical mask','gowns','gloves']
fig, ax1, = plt.subplots()
fig.set_size_inches(25, 12)
sns.barplot(y='Month',x='Total Amount',hue='Type',\
    data=df_clean_EMS[df_clean_EMS['Type'].isin(type_needed)],ax=ax1,\
        ci=False,orient='h',color='#E315C7')
plt.yticks(fontsize=18)
plt.xticks(fontsize=20)
plt.xlabel('Month of the year', fontsize=20)
plt.ylabel('mean Amount', fontsize=25)
plt.title('day', fontsize=30)

plt.tight_layout()
plt.show()

# %%
'''hist is a good way to analysis the distribution of the data'''

fig, ax1, = plt.subplots()
fig.set_size_inches(25, 12)
sns.distplot(df_clean_EMS[(df_clean_EMS['Type']=='gloves')\
    &(df_clean_EMS['Facility Type']=='Ltcf')\
        &(df_clean_EMS['Total Amount']<=500)]['Total Amount'],ax=ax1)
plt.yticks(fontsize=18)
plt.xticks(fontsize=20)
plt.xlabel('give amount', fontsize=20)
plt.ylabel('count', fontsize=25)
plt.title('give out distribution', fontsize=30)

mean=df_clean_EMS[(df_clean_EMS['Type']=='gloves')\
    &(df_clean_EMS['Facility Type']=='Ltcf')\
        &(df_clean_EMS['Total Amount']<=500)]['Total Amount'].mean()
plt.axvline(mean, color='red')

plt.tight_layout()
plt.show()

# %%
def draw_type_facility_type_amount_dist(Type,FacilityType,upper_level=50000):

    fig, ax1, = plt.subplots()
    fig.set_size_inches(25, 12)
    sns.distplot(df_clean_EMS[(df_clean_EMS['Type']==Type)\
        &(df_clean_EMS['Facility Type']==FacilityType)\
            &(df_clean_EMS['Total Amount']<=upper_level)]['Total Amount'],ax=ax1)
    plt.yticks(fontsize=18)
    plt.xticks(fontsize=20)
    plt.xlabel('give amount', fontsize=20)
    plt.ylabel('count', fontsize=25)
    plt.title('give out distribution', fontsize=30)

    mean=df_clean_EMS[(df_clean_EMS['Type']==Type)\
        &(df_clean_EMS['Facility Type']==FacilityType)\
            &(df_clean_EMS['Total Amount']<=upper_level)]['Total Amount'].mean()
    plt.axvline(mean, 0,1,color='red')

    plt.tight_layout()
    plt.show()

# %%
draw_type_facility_type_amount_dist('n95','Acute',2000)

# %%
'''box plot is another good way to show distribution'''
fig, ax1, = plt.subplots()
fig.set_size_inches(25, 12)
sns.boxplot(df_clean_EMS[(df_clean_EMS['Type']=='gloves')\
    &(df_clean_EMS['Facility Type']=='Ltcf')\
        &(df_clean_EMS['Total Amount']<=5000)]['Total Amount'],ax=ax1)
plt.yticks(fontsize=18)
plt.xticks(fontsize=20)
plt.xlabel('give amount', fontsize=20)
plt.ylabel('count', fontsize=25)
plt.title('give out distribution', fontsize=30)

mean=df_clean_EMS[(df_clean_EMS['Type']=='gloves')\
    &(df_clean_EMS['Facility Type']=='Ltcf')\
        &(df_clean_EMS['Total Amount']<=500)]['Total Amount'].mean()
plt.axvline(mean, color='red')

plt.tight_layout()
plt.show()

# %%
fig, ax1, = plt.subplots()
#fig.set_size_inches(25, 12)
sns.set(rc={'figure.figsize':(12,10)})
sns.boxplot(y='Total Amount',
            x='Month',data=df_clean_EMS,ax=ax1)
plt.tight_layout()
plt.show()

# %%
'''sns.swarmplot(), which can show every data point with a distribution,
 is a good helper for sns.boxplot, '''

fig, ax1, = plt.subplots()
#fig.set_size_inches(25, 12)
sns.set(rc={'figure.figsize':(12,10)})
sns.boxplot(x='Total Amount',
            y='Month',data=df_clean_EMS[df_clean_EMS['Total Amount']<=5000],ax=ax1,orient='h',color='#E55C94')

sns.swarmplot(x='Total Amount',
            y='Month',data=df_clean_EMS[df_clean_EMS['Total Amount']<=5000],ax=ax1,orient='h',color='green')
plt.tight_layout()
plt.show()

# %%
'''same as sns.violinplot'''
fig, ax1, = plt.subplots()
#fig.set_size_inches(25, 12)
#sns.set(rc={'figure.figsize':(12,10)})
sns.boxplot(x='Total Amount',
            y='Month',data=df_clean_EMS[df_clean_EMS['Total Amount']<=5000],ax=ax1,orient='h',color='#E55C94')

sns.violinplot(x='Total Amount',
            y='Month',data=df_clean_EMS[df_clean_EMS['Total Amount']<=5000],ax=ax1,orient='h',color='green')
plt.tight_layout()
plt.show()

# %%
'''let's try sns.scatterplot()'''
fig, ax1, = plt.subplots()
#fig.set_size_inches(25, 12)
#sns.set(rc={'figure.figsize':(12,10)})
sns.scatterplot(y='Total Amount',
            x='week_of_year',data=df_clean_EMS[df_clean_EMS['Total Amount']<=5000],\
                ax=ax1,color='#E55C94')
plt.tight_layout()
plt.show()

# %%
'''what is lmplot()==> basically: scattern plot with regression line'''
#this is not a good example, basically you need to find 2 continue variable

sns.lmplot(y='Total Amount',
            x='Month',data=df_clean_EMS[df_clean_EMS['Total Amount']<=5000],height=10)

plt.tight_layout()
plt.show()

health care digital ticket EDA/python models/seaborn viz center.py


import base64
from collections import OrderedDict
import datetime

from google.cloud._helpers import UTC
from google.cloud._helpers import _date_from_iso8601_date
from google.cloud._helpers import _datetime_from_microseconds
from google.cloud._helpers import _datetime_to_rfc3339
from google.cloud._helpers import _RFC3339_NO_FRACTION
from google.cloud._helpers import _time_from_iso8601_time_naive
from google.cloud._helpers import _to_bytes

def _not_null(value, field):
    """Check whether 'value' should be coerced to 'field' type."""
    return value is not None or field.mode != 'NULLABLE'

def _int_from_json(value, field):
    """Coerce 'value' to an int, if set or not nullable."""
    if _not_null(value, field):
        return int(value)

def _float_from_json(value, field):
    """Coerce 'value' to a float, if set or not nullable."""
    if _not_null(value, field):
        return float(value)

def _bool_from_json(value, field):
    """Coerce 'value' to a bool, if set or not nullable."""
    if _not_null(value, field):
        return value.lower() in ['t', 'true', '1']

def _string_from_json(value, _):
    """NOOP string -> string coercion"""
    return value

def _bytes_from_json(value, field):
    """Base64-decode value"""
    if _not_null(value, field):
        return base64.decodestring(_to_bytes(value))

def _timestamp_from_json(value, field):
    """Coerce 'value' to a datetime, if set or not nullable."""
    if _not_null(value, field):
        # value will be a float in seconds, to microsecond precision, in UTC.
        return _datetime_from_microseconds(1e6 * float(value))

def _datetime_from_json(value, field):
    """Coerce 'value' to a datetime, if set or not nullable."""
    if _not_null(value, field):
        # value will be a string, in YYYY-MM-DDTHH:MM:SS form.
        return datetime.datetime.strptime(value, _RFC3339_NO_FRACTION)

def _date_from_json(value, field):
    """Coerce 'value' to a datetime date, if set or not nullable"""
    if _not_null(value, field):
        # value will be a string, in YYYY-MM-DD form.
        return _date_from_iso8601_date(value)

def _time_from_json(value, field):
    """Coerce 'value' to a datetime date, if set or not nullable"""
    if _not_null(value, field):
        # value will be a string, in HH:MM:SS form.
        return _time_from_iso8601_time_naive(value)

def _record_from_json(value, field):
    """Coerce 'value' to a mapping, if set or not nullable."""
    if _not_null(value, field):
        record = {}
        record_iter = zip(field.fields, value['f'])
        for subfield, cell in record_iter:
            converter = _CELLDATA_FROM_JSON[subfield.field_type]
            if subfield.mode == 'REPEATED':
                value = [converter(item['v'], subfield) for item in cell['v']]
            else:
                value = converter(cell['v'], subfield)
            record[subfield.name] = value
        return record

_CELLDATA_FROM_JSON = {
    'INTEGER': _int_from_json,
    'INT64': _int_from_json,
    'FLOAT': _float_from_json,
    'FLOAT64': _float_from_json,
    'BOOLEAN': _bool_from_json,
    'BOOL': _bool_from_json,
    'STRING': _string_from_json,
    'BYTES': _bytes_from_json,
    'TIMESTAMP': _timestamp_from_json,
    'DATETIME': _datetime_from_json,
    'DATE': _date_from_json,
    'TIME': _time_from_json,
    'RECORD': _record_from_json,
}

def _int_to_json(value):
    """Coerce 'value' to an JSON-compatible representation."""
    if isinstance(value, int):
        value = str(value)
    return value

def _float_to_json(value):
    """Coerce 'value' to an JSON-compatible representation."""
    return value

def _bool_to_json(value):
    """Coerce 'value' to an JSON-compatible representation."""
    if isinstance(value, bool):
        value = 'true' if value else 'false'
    return value

def _bytes_to_json(value):
    """Coerce 'value' to an JSON-compatible representation."""
    if isinstance(value, bytes):
        value = base64.encodestring(value)
    return value

def _timestamp_to_json(value):
    """Coerce 'value' to an JSON-compatible representation."""
    if isinstance(value, datetime.datetime):
        if value.tzinfo not in (None, UTC):
            # Convert to UTC and remove the time zone info.
            value = value.replace(tzinfo=None) - value.utcoffset()
        value = '%s %s+00:00' % (
            value.date().isoformat(), value.time().isoformat())
    return value

def _datetime_to_json(value):
    """Coerce 'value' to an JSON-compatible representation."""
    if isinstance(value, datetime.datetime):
        value = _datetime_to_rfc3339(value)
    return value

def _date_to_json(value):
    """Coerce 'value' to an JSON-compatible representation."""
    if isinstance(value, datetime.date):
        value = value.isoformat()
    return value

def _time_to_json(value):
    """Coerce 'value' to an JSON-compatible representation."""
    if isinstance(value, datetime.time):
        value = value.isoformat()
    return value

_SCALAR_VALUE_TO_JSON = {
    'INTEGER': _int_to_json,
    'INT64': _int_to_json,
    'FLOAT': _float_to_json,
    'FLOAT64': _float_to_json,
    'BOOLEAN': _bool_to_json,
    'BOOL': _bool_to_json,
    'BYTES': _bytes_to_json,
    'TIMESTAMP': _timestamp_to_json,
    'DATETIME': _datetime_to_json,
    'DATE': _date_to_json,
    'TIME': _time_to_json,
}

def _row_from_json(row, schema):
    """Convert JSON row data to row with appropriate types.

    :type row: dict
    :param row: A JSON response row to be converted.

    :type schema: tuple
    :param schema: A tuple of
                   :class:`~google.cloud.bigquery.schema.SchemaField`.

    :rtype: tuple
    :returns: A tuple of data converted to native types.
    """
    row_data = []
    for field, cell in zip(schema, row['f']):
        converter = _CELLDATA_FROM_JSON[field.field_type]
        if field.mode == 'REPEATED':
            row_data.append([converter(item['v'], field)
                             for item in cell['v']])
        else:
            row_data.append(converter(cell['v'], field))

    return tuple(row_data)

def _rows_from_json(rows, schema):
    """Convert JSON row data to rows with appropriate types."""
    return [_row_from_json(row, schema) for row in rows]

class _ConfigurationProperty(object):
    """Base property implementation.

    Values will be stored on a `_configuration` helper attribute of the
    property's job instance.

    :type name: str
    :param name:  name of the property
    """

    def __init__(self, name):
        self.name = name
        self._backing_name = '_%s' % (self.name,)

    def __get__(self, instance, owner):
        """Descriptor protocal:  accesstor"""
        if instance is None:
            return self
        return getattr(instance._configuration, self._backing_name)

    def _validate(self, value):
        """Subclasses override to impose validation policy."""
        pass

    def __set__(self, instance, value):
        """Descriptor protocal:  mutator"""
        self._validate(value)
        setattr(instance._configuration, self._backing_name, value)

    def __delete__(self, instance):
        """Descriptor protocal:  deleter"""
        delattr(instance._configuration, self._backing_name)

class _TypedProperty(_ConfigurationProperty):
    """Property implementation:  validates based on value type.

    :type name: str
    :param name:  name of the property

    :type property_type: type or sequence of types
    :param property_type: type to be validated
    """
    def __init__(self, name, property_type):
        super(_TypedProperty, self).__init__(name)
        self.property_type = property_type

    def _validate(self, value):
        """Ensure that 'value' is of the appropriate type.

        :raises: ValueError on a type mismatch.
        """
        if not isinstance(value, self.property_type):
            raise ValueError('Required type: %s' % (self.property_type,))

class _EnumProperty(_ConfigurationProperty):
    """Pseudo-enumeration class.

    Subclasses must define ``ALLOWED`` as a class-level constant:  it must
    be a sequence of strings.

    :type name: str
    :param name:  name of the property.
    """
    def _validate(self, value):
        """Check that ``value`` is one of the allowed values.

        :raises: ValueError if value is not allowed.
        """
        if value not in self.ALLOWED:
            raise ValueError('Pass one of: %s' ', '.join(self.ALLOWED))

class UDFResource(object):
    """Describe a single user-defined function (UDF) resource.

    :type udf_type: str
    :param udf_type: the type of the resource ('inlineCode' or 'resourceUri')

    :type value: str
    :param value: the inline code or resource URI.

    See
    https://cloud.google.com/bigquery/user-defined-functions#api
    """
    def __init__(self, udf_type, value):
        self.udf_type = udf_type
        self.value = value

    def __eq__(self, other):
        return(
            self.udf_type == other.udf_type and
            self.value == other.value)

class UDFResourcesProperty(object):
    """Custom property type, holding :class:`UDFResource` instances."""

    def __get__(self, instance, owner):
        """Descriptor protocol:  accessor"""
        if instance is None:
            return self
        return list(instance._udf_resources)

    def __set__(self, instance, value):
        """Descriptor protocol:  mutator"""
        if not all(isinstance(u, UDFResource) for u in value):
            raise ValueError("udf items must be UDFResource")
        instance._udf_resources = tuple(value)

class AbstractQueryParameter(object):
    """Base class for named / positional query parameters.
    """
    @classmethod
    def from_api_repr(cls, resource):
        """Factory: construct paramter from JSON resource.

        :type resource: dict
        :param resource: JSON mapping of parameter

        :rtype: :class:`ScalarQueryParameter`
        """
        raise NotImplementedError

    def to_api_repr(self):
        """Construct JSON API representation for the parameter.

        :rtype: dict
        """
        raise NotImplementedError

class ScalarQueryParameter(AbstractQueryParameter):
    """Named / positional query parameters for scalar values.

    :type name: str or None
    :param name: Parameter name, used via ``@foo`` syntax.  If None, the
                 paramter can only be addressed via position (``?``).

    :type type_: str
    :param type_: name of parameter type.  One of 'STRING', 'INT64',
                  'FLOAT64', 'BOOL', 'TIMESTAMP', 'DATETIME', or 'DATE'.

    :type value: str, int, float, bool, :class:`datetime.datetime`, or
                 :class:`datetime.date`.
    :param value: the scalar parameter value.
    """
    def __init__(self, name, type_, value):
        self.name = name
        self.type_ = type_
        self.value = value

    @classmethod
    def positional(cls, type_, value):
        """Factory for positional paramters.

        :type type_: str
        :param type_:
            name of paramter type.  One of 'STRING', 'INT64',
            'FLOAT64', 'BOOL', 'TIMESTAMP', 'DATETIME', or 'DATE'.

        :type value: str, int, float, bool, :class:`datetime.datetime`, or
                     :class:`datetime.date`.
        :param value: the scalar parameter value.

        :rtype: :class:`ScalarQueryParameter`
        :returns: instance without name
        """
        return cls(None, type_, value)

    @classmethod
    def from_api_repr(cls, resource):
        """Factory: construct paramter from JSON resource.

        :type resource: dict
        :param resource: JSON mapping of parameter

        :rtype: :class:`ScalarQueryParameter`
        :returns: instance
        """
        name = resource.get('name')
        type_ = resource['parameterType']['type']
        value = resource['parameterValue']['value']
        converted = _CELLDATA_FROM_JSON[type_](value, None)
        return cls(name, type_, converted)

    def to_api_repr(self):
        """Construct JSON API representation for the parameter.

        :rtype: dict
        :returns: JSON mapping
        """
        value = self.value
        converter = _SCALAR_VALUE_TO_JSON.get(self.type_)
        if converter is not None:
            value = converter(value)
        resource = {
            'parameterType': {
                'type': self.type_,
            },
            'parameterValue': {
                'value': value,
            },
        }
        if self.name is not None:
            resource['name'] = self.name
        return resource

class ArrayQueryParameter(AbstractQueryParameter):
    """Named / positional query parameters for array values.

    :type name: str or None
    :param name: Parameter name, used via ``@foo`` syntax.  If None, the
                 paramter can only be addressed via position (``?``).

    :type array_type: str
    :param array_type:
        name of type of array elements.  One of `'STRING'`, `'INT64'`,
        `'FLOAT64'`, `'BOOL'`, `'TIMESTAMP'`, or `'DATE'`.

    :type values: list of appropriate scalar type.
    :param values: the parameter array values.
    """
    def __init__(self, name, array_type, values):
        self.name = name
        self.array_type = array_type
        self.values = values

    @classmethod
    def positional(cls, array_type, values):
        """Factory for positional paramters.

        :type array_type: str
        :param array_type:
            name of type of array elements.  One of `'STRING'`, `'INT64'`,
            `'FLOAT64'`, `'BOOL'`, `'TIMESTAMP'`, or `'DATE'`.

        :type values: list of appropriate scalar type
        :param values: the parameter array values.

        :rtype: :class:`ArrayQueryParameter`
        :returns: instance without name
        """
        return cls(None, array_type, values)

    @classmethod
    def from_api_repr(cls, resource):
        """Factory: construct paramter from JSON resource.

        :type resource: dict
        :param resource: JSON mapping of parameter

        :rtype: :class:`ArrayQueryParameter`
        :returns: instance
        """
        name = resource.get('name')
        array_type = resource['parameterType']['arrayType']['type']
        values = [
            value['value']
            for value
            in resource['parameterValue']['arrayValues']]
        converted = [
            _CELLDATA_FROM_JSON[array_type](value, None) for value in values]
        return cls(name, array_type, converted)

    def to_api_repr(self):
        """Construct JSON API representation for the parameter.

        :rtype: dict
        :returns: JSON mapping
        """
        values = self.values
        converter = _SCALAR_VALUE_TO_JSON.get(self.array_type)
        if converter is not None:
            values = [converter(value) for value in values]
        resource = {
            'parameterType': {
                'type': 'ARRAY',
                'arrayType': {
                    'type': self.array_type,
                },
            },
            'parameterValue': {
                'arrayValues': [{'value': value} for value in values],
            },
        }
        if self.name is not None:
            resource['name'] = self.name
        return resource

class StructQueryParameter(AbstractQueryParameter):
    """Named / positional query parameters for struct values.

    :type name: str or None
    :param name: Parameter name, used via ``@foo`` syntax.  If None, the
                 paramter can only be addressed via position (``?``).

    :type sub_params: tuple of :class:`ScalarQueryParameter`
    :param sub_params: the sub-parameters for the struct
    """
    def __init__(self, name, *sub_params):
        self.name = name
        self.struct_types = OrderedDict(
            (sub.name, sub.type_) for sub in sub_params)
        self.struct_values = {sub.name: sub.value for sub in sub_params}

    @classmethod
    def positional(cls, *sub_params):
        """Factory for positional paramters.

        :type sub_params: tuple of :class:`ScalarQueryParameter`
        :param sub_params: the sub-parameters for the struct

        :rtype: :class:`StructQueryParameter`
        :returns: instance without name
        """
        return cls(None, *sub_params)

    @classmethod
    def from_api_repr(cls, resource):
        """Factory: construct paramter from JSON resource.

        :type resource: dict
        :param resource: JSON mapping of parameter

        :rtype: :class:`StructQueryParameter`
        :returns: instance
        """
        name = resource.get('name')
        instance = cls(name)
        types = instance.struct_types
        for item in resource['parameterType']['structTypes']:
            types[item['name']] = item['type']['type']
        struct_values = resource['parameterValue']['structValues']
        for key, value in struct_values.items():
            type_ = types[key]
            value = value['value']
            converted = _CELLDATA_FROM_JSON[type_](value, None)
            instance.struct_values[key] = converted
        return instance

    def to_api_repr(self):
        """Construct JSON API representation for the parameter.

        :rtype: dict
        :returns: JSON mapping
        """
        types = [
            {'name': key, 'type': {'type': value}}
            for key, value in self.struct_types.items()
        ]
        values = {}
        for name, value in self.struct_values.items():
            converter = _SCALAR_VALUE_TO_JSON.get(self.struct_types[name])
            if converter is not None:
                value = converter(value)
            values[name] = {'value': value}

        resource = {
            'parameterType': {
                'type': 'STRUCT',
                'structTypes': types,
            },
            'parameterValue': {
                'structValues': values,
            },
        }
        if self.name is not None:
            resource['name'] = self.name
        return resource

class QueryParametersProperty(object):
    """Custom property type, holding query parameter instances."""

    def __get__(self, instance, owner):
        """Descriptor protocol:  accessor

        :type instance: :class:`QueryParametersProperty`
        :param instance: instance owning the property (None if accessed via
                         the class).

        :type owner: type
        :param owner: the class owning the property.

        :rtype: list of instances of classes derived from
                :class:`AbstractQueryParameter`.
        :returns: the descriptor, if accessed via the class, or the instance's
                  query paramters.
        """
        if instance is None:
            return self
        return list(instance._query_parameters)

    def __set__(self, instance, value):
        """Descriptor protocol:  mutator

        :type instance: :class:`QueryParametersProperty`
        :param instance: instance owning the property (None if accessed via
                         the class).

        :type value: list of instances of classes derived from
                     :class:`AbstractQueryParameter`.
        :param value: new query parameters for the instance.
        """
        if not all(isinstance(u, AbstractQueryParameter) for u in value):
            raise ValueError(
                "query parameters must be derived from AbstractQueryParameter")
        instance._query_parameters = tuple(value)


# repo originally forked from https://github.com/Confusezius/Deep-Metric-Learning-Baselines

################# LIBRARIES ###############################
import warnings
warnings.filterwarnings("ignore")

import numpy as np, pandas as pd, copy, torch, random, os

from torch.utils.data import Dataset
from PIL import Image
from torchvision import transforms

"""============================================================================"""
################ FUNCTION TO RETURN ALL DATALOADERS NECESSARY ####################
def give_dataloaders(dataset, opt):
    """
    Args:
        dataset: string, name of dataset for which the dataloaders should be returned.
        opt:     argparse.Namespace, contains all training-specific parameters.
    Returns:
        dataloaders: dict of dataloaders for training, testing and evaluation on training.
    """
    #Dataset selection
    if opt.dataset=='vehicle_id':
        datasets = give_VehicleID_datasets(opt)
    elif opt.dataset=='Inaturalist':
        datasets = give_inaturalist_datasets(opt)
    else:
        raise Exception('No Dataset >{}< available!'.format(dataset))

    #Move datasets to dataloaders.
    dataloaders = {}

    for key, dataset in datasets.items():
        if isinstance(dataset, TrainDatasetsmoothap) and key == 'training':
            dataloaders[key] = torch.utils.data.DataLoader(dataset, batch_size=opt.bs,
                    num_workers=opt.kernels, sampler=torch.utils.data.SequentialSampler(dataset),
                    pin_memory=True, drop_last=True)
        else:
            is_val = dataset.is_validation
            if key == 'training':
                dataloaders[key] = torch.utils.data.DataLoader(dataset, batch_size=opt.bs,
                        num_workers=opt.kernels, shuffle=not is_val, pin_memory=True, drop_last=not is_val)
            else:
                dataloaders[key] = torch.utils.data.DataLoader(dataset, batch_size=opt.bs,
                        num_workers=6, shuffle=not is_val, pin_memory=True, drop_last=not is_val)
    return dataloaders

def give_inaturalist_datasets(opt):
    """
    This function generates a training, testing and evaluation dataloader for Metric Learning on the Inaturalist 2018 dataset.
    For Metric Learning, training and test sets are provided by given json files. Will define a train and test split
    So no random shuffling of classes.

    Args:
        opt: argparse.Namespace, contains all traininig-specific parameters.
    Returns:
        dict of PyTorch datasets for training, testing and evaluation.
    """
    #Load text-files containing classes and imagepaths.
    #Generate image_dicts of shape {class_idx:[list of paths to images belong to this class] ...}
    train_image_dict, val_image_dict  = {},{}
    with open(os.path.join(opt.source_path,'Inat_dataset_splits/Inaturalist_train_set1.txt')) as f:
        FileLines = f.readlines()
        FileLines = [x.strip() for x in FileLines]

        for entry in FileLines:
            info = entry.split('/')
            if '/'.join([info[-3],info[-2]]) not in train_image_dict:
                train_image_dict['/'.join([info[-3],info[-2]])] = []
            train_image_dict['/'.join([info[-3],info[-2]])].append(os.path.join(opt.source_path,entry))

    with open(os.path.join(opt.source_path,'Inat_dataset_splits/Inaturalist_test_set1.txt')) as f:
        FileLines = f.readlines()
        FileLines = [x.strip() for x in FileLines]

        for entry in FileLines:
            info = entry.split('/')
            if '/'.join([info[-3],info[-2]]) not in val_image_dict:
                val_image_dict['/'.join([info[-3],info[-2]])] = []
            val_image_dict['/'.join([info[-3],info[-2]])].append(os.path.join(opt.source_path,entry))

    new_train_dict = {}
    class_ind_ind = 0
    for cate in train_image_dict:
        new_train_dict["te/%d"%class_ind_ind] = train_image_dict[cate]
        class_ind_ind += 1
    train_image_dict = new_train_dict

    train_dataset = TrainDatasetsmoothap(train_image_dict, opt)

    val_dataset         = BaseTripletDataset(val_image_dict,   opt, is_validation=True)
    eval_dataset        = BaseTripletDataset(train_image_dict, opt, is_validation=True)

    #train_dataset.conversion       = conversion
    #val_dataset.conversion         = conversion
    #eval_dataset.conversion        = conversion

    return {'training':train_dataset, 'testing':val_dataset, 'evaluation':eval_dataset}
    # return {'training':train_dataset, 'testing':val_dataset, 'evaluation':eval_dataset, 'super_evaluation':super_train_dataset}

def give_VehicleID_datasets(opt):
    """
    This function generates a training, testing and evaluation dataloader for Metric Learning on the PKU Vehicle dataset.
    For Metric Learning, training and (multiple) test sets are provided by separate text files, train_list and test_list_<n_classes_2_test>.txt.
    So no random shuffling of classes.

    Args:
        opt: argparse.Namespace, contains all traininig-specific parameters.
    Returns:
        dict of PyTorch datasets for training, testing and evaluation.
    """
    #Load respective text-files
    train       = np.array(pd.read_table(opt.source_path+'/train_test_split/train_list.txt', header=None, delim_whitespace=True))
    small_test  = np.array(pd.read_table(opt.source_path+'/train_test_split/test_list_800.txt', header=None, delim_whitespace=True))
    medium_test = np.array(pd.read_table(opt.source_path+'/train_test_split/test_list_1600.txt', header=None, delim_whitespace=True))
    big_test    = np.array(pd.read_table(opt.source_path+'/train_test_split/test_list_2400.txt', header=None, delim_whitespace=True))

    #Generate conversions
    lab_conv_train = {x:i for i,x in enumerate(np.unique(train[:,1]))}
    train[:,1] = np.array([lab_conv_train[x] for x in train[:,1]])
    lab_conv = {x:i for i,x in enumerate(np.unique(np.concatenate([small_test[:,1], medium_test[:,1], big_test[:,1]])))}
    small_test[:,1]  = np.array([lab_conv[x] for x in small_test[:,1]])
    medium_test[:,1] = np.array([lab_conv[x] for x in medium_test[:,1]])
    big_test[:,1]    = np.array([lab_conv[x] for x in big_test[:,1]])

    #Generate Image-Dicts for training and different testings of shape {class_idx:[list of paths to images belong to this class] ...}
    train_image_dict    = {}
    for img_path, key in train:
        if not key in train_image_dict.keys():
            train_image_dict[key] = []
        train_image_dict[key].append(opt.source_path+'/image/{:07d}.jpg'.format(img_path))

    small_test_dict = {}
    for img_path, key in small_test:
        if not key in small_test_dict.keys():
            small_test_dict[key] = []
        small_test_dict[key].append(opt.source_path+'/image/{:07d}.jpg'.format(img_path))

    medium_test_dict    = {}
    for img_path, key in medium_test:
        if not key in medium_test_dict.keys():
            medium_test_dict[key] = []
        medium_test_dict[key].append(opt.source_path+'/image/{:07d}.jpg'.format(img_path))

    big_test_dict    = {}
    for img_path, key in big_test:
        if not key in big_test_dict.keys():
            big_test_dict[key] = []
        big_test_dict[key].append(opt.source_path+'/image/{:07d}.jpg'.format(img_path))

    attribute = np.array(pd.read_table(opt.source_path+'/attribute/model_attr.txt', header=None, delim_whitespace=True))

    new_dict = {}
    not_found = 0
    for thing in attribute:
        if lab_conv_train[thing[0]] not in train_image_dict:
            not_found +=1
        else:
            if thing[1] not in new_dict:
                new_dict[thing[1]] = []
            new_dict[thing[1]].append(lab_conv_train[thing[0]])

    train_dataset = TrainDatasetsmoothap(train_image_dict, opt)

    eval_dataset  = BaseTripletDataset(train_image_dict, opt,    is_validation=True)
    val_small_dataset     = BaseTripletDataset(small_test_dict, opt,  is_validation=True)
    val_medium_dataset    = BaseTripletDataset(medium_test_dict, opt, is_validation=True)
    val_big_dataset       = BaseTripletDataset(big_test_dict, opt,    is_validation=True)

    return {'training':train_dataset, 'testing_set1':val_small_dataset, 'testing_set2':val_medium_dataset, \
            'testing_set3':val_big_dataset, 'evaluation':eval_dataset}

################## BASIC PYTORCH DATASET USED FOR ALL DATASETS ##################################
class BaseTripletDataset(Dataset):
    """
    Dataset class to provide (augmented) correctly prepared training samples corresponding to standard DML literature.
    This includes normalizing to ImageNet-standards, and Random & Resized cropping of shapes 224 for ResNet50 and 227 for
    GoogLeNet during Training. During validation, only resizing to 256 or center cropping to 224/227 is performed.
    """
    def __init__(self, image_dict, opt, samples_per_class=8, is_validation=False):
        """
        Dataset Init-Function.

        Args:
            image_dict:         dict, Dictionary of shape {class_idx:[list of paths to images belong to this class] ...} providing all the training paths and classes.
            opt:                argparse.Namespace, contains all training-specific parameters.
            samples_per_class:  Number of samples to draw from one class before moving to the next when filling the batch.
            is_validation:      If is true, dataset properties for validation/testing are used instead of ones for training.
        Returns:
            Nothing!
        """
        #Define length of dataset
        self.n_files     = np.sum([len(image_dict[key]) for key in image_dict.keys()])

        self.is_validation = is_validation

        self.pars        = opt
        self.image_dict  = image_dict

        self.avail_classes    = sorted(list(self.image_dict.keys()))

        #Convert image dictionary from classname:content to class_idx:content, because the initial indices are not necessarily from 0 - <n_classes>.
        self.image_dict    = {i:self.image_dict[key] for i,key in enumerate(self.avail_classes)}
        self.avail_classes = sorted(list(self.image_dict.keys()))

        #Init. properties that are used when filling up batches.
        if not self.is_validation:
            self.samples_per_class = samples_per_class
            #Select current class to sample images from up to <samples_per_class>
            self.current_class   = np.random.randint(len(self.avail_classes))
            self.classes_visited = [self.current_class, self.current_class]
            self.n_samples_drawn = 0

        #Data augmentation/processing methods.
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
        transf_list = []
        if not self.is_validation:
            transf_list.extend([transforms.RandomResizedCrop(size=224) if opt.arch=='resnet50' else transforms.RandomResizedCrop(size=227),
                                transforms.RandomHorizontalFlip(0.5)])
        else:
            transf_list.extend([transforms.Resize(256),
                                transforms.CenterCrop(224) if opt.arch=='resnet50' else transforms.CenterCrop(227)])

        transf_list.extend([transforms.ToTensor(), normalize])
        self.transform = transforms.Compose(transf_list)

        #Convert Image-Dict to list of (image_path, image_class). Allows for easier direct sampling.
        self.image_list = [[(x,key) for x in self.image_dict[key]] for key in self.image_dict.keys()]
        self.image_list = [x for y in self.image_list for x in y]

        #Flag that denotes if dataset is called for the first time.
        self.is_init = True

    def ensure_3dim(self, img):
        """
        Function that ensures that the input img is three-dimensional.

        Args:
            img: PIL.Image, image which is to be checked for three-dimensionality (i.e. if some images are black-and-white in an otherwise coloured dataset).
        Returns:
            Checked PIL.Image img.
        """
        if len(img.size)==2:
            img = img.convert('RGB')
        return img

    def __getitem__(self, idx):
        """
        Args:
            idx: Sample idx for training sample
        Returns:
            tuple of form (sample_class, torch.Tensor() of input image)
        """
        if self.pars.loss == 'smoothap' or self.pars.loss == 'smoothap_element':
            if self.is_init:
                #self.current_class = self.avail_classes[idx%len(self.avail_classes)]
                self.is_init = False

            if not self.is_validation:
                if self.samples_per_class==1:
                    return self.image_list[idx][-1], self.transform(self.ensure_3dim(Image.open(self.image_list[idx][0])))

                if self.n_samples_drawn==self.samples_per_class:
                    #Once enough samples per class have been drawn, we choose another class to draw samples from.
                    #Note that we ensure with self.classes_visited that no class is chosen if it had been chosen
                    #previously or one before that.
                    counter = copy.deepcopy(self.avail_classes)
                    for prev_class in self.classes_visited:
                        if prev_class in counter: counter.remove(prev_class)

                    self.current_class   = counter[idx%len(counter)]
                    #self.classes_visited = self.classes_visited[1:]+[self.current_class]
                    # EDIT -> there can be no class repeats
                    self.classes_visited = self.classes_visited+[self.current_class]
                    self.n_samples_drawn = 0

                class_sample_idx = idx%len(self.image_dict[self.current_class])
                self.n_samples_drawn += 1

                out_img = self.transform(self.ensure_3dim(Image.open(self.image_dict[self.current_class][class_sample_idx])))
                return self.current_class,out_img
            else:
                return self.image_list[idx][-1], self.transform(self.ensure_3dim(Image.open(self.image_list[idx][0])))
        else:
            if self.is_init:
                self.current_class = self.avail_classes[idx%len(self.avail_classes)]
                self.is_init = False
            if not self.is_validation:
                if self.samples_per_class==1:
                    return self.image_list[idx][-1], self.transform(self.ensure_3dim(Image.open(self.image_list[idx][0])))

                if self.n_samples_drawn==self.samples_per_class:
                    #Once enough samples per class have been drawn, we choose another class to draw samples from.
                    #Note that we ensure with self.classes_visited that no class is chosen if it had been chosen
                    #previously or one before that.
                    counter = copy.deepcopy(self.avail_classes)
                    for prev_class in self.classes_visited:
                        if prev_class in counter: counter.remove(prev_class)

                    self.current_class   = counter[idx%len(counter)]
                    self.classes_visited = self.classes_visited[1:]+[self.current_class]
                    self.n_samples_drawn = 0

                class_sample_idx = idx%len(self.image_dict[self.current_class])
                self.n_samples_drawn += 1

                out_img = self.transform(self.ensure_3dim(Image.open(self.image_dict[self.current_class][class_sample_idx])))
                return self.current_class,out_img
            else:
                return self.image_list[idx][-1], self.transform(self.ensure_3dim(Image.open(self.image_list[idx][0])))

    def __len__(self):
        return self.n_files

flatten = lambda l: [item for sublist in l for item in sublist]

######################## dataset for SmoothAP regular training ##################################

class TrainDatasetsmoothap(Dataset):
    """
    This dataset class allows mini-batch formation pre-epoch, for greater speed

    """
    def __init__(self, image_dict, opt):
        """
        Args:
            image_dict: two-level dict, `super_dict[super_class_id][class_id]` gives the list of
                        image paths having the same super-label and class label
        """
        self.image_dict = image_dict
        self.dataset_name = opt.dataset
        self.batch_size = opt.bs
        self.samples_per_class = opt.samples_per_class
        for sub in self.image_dict:
            newsub = []
            for instance in self.image_dict[sub]:
                newsub.append((sub, instance))
            self.image_dict[sub] = newsub

        # checks
        # provide avail_classes
        self.avail_classes = [*self.image_dict]
        # Data augmentation/processing methods.
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
        transf_list = []

        transf_list.extend([
            transforms.RandomResizedCrop(size=224) if opt.arch in ['resnet50', 'resnet50_mcn'] else transforms.RandomResizedCrop(size=227),
            transforms.RandomHorizontalFlip(0.5)])
        transf_list.extend([transforms.ToTensor(), normalize])
        self.transform = transforms.Compose(transf_list)

        self.reshuffle()

    def ensure_3dim(self, img):
        if len(img.size) == 2:
            img = img.convert('RGB')
        return img

    def reshuffle(self):

        image_dict = copy.deepcopy(self.image_dict)
        print('shuffling data')
        for sub in image_dict:
            random.shuffle(image_dict[sub])

        classes = [*image_dict]
        random.shuffle(classes)
        total_batches = []
        batch = []
        finished = 0
        while finished == 0:
            for sub_class in classes:
                if (len(image_dict[sub_class]) >=self.samples_per_class) and (len(batch) < self.batch_size/self.samples_per_class) :
                    batch.append(image_dict[sub_class][:self.samples_per_class])
                    image_dict[sub_class] = image_dict[sub_class][self.samples_per_class:]

            if len(batch) == self.batch_size/self.samples_per_class:
                total_batches.append(batch)
                batch = []
            else:
                finished = 1

        random.shuffle(total_batches)
        self.dataset = flatten(flatten(total_batches))

    def __getitem__(self, idx):
        # we use SequentialSampler together with SuperLabelTrainDataset,
        # so idx==0 indicates the start of a new epoch
        batch_item = self.dataset[idx]

        if self.dataset_name == 'Inaturalist':
            cls = int(batch_item[0].split('/')[1])

        else:
            cls = batch_item[0]
        img = Image.open(batch_item[1])
        return cls, self.transform(self.ensure_3dim(img))

    def __len__(self):
        return len(self.dataset)