'''
This module contains the back-end class (WholeClassAnalysis) that handles the analysis of pixel classes as-a-whole. It is accessed in the GUI
through the second half of the third tab of the program (use pixel classifier).
It is also available through the public (non-GUI) API of PalmettoBUG. 
'''
## License / derivation info
# PalmettoBUG as a whole, and therefore also this script, is licensed under the GPL3 license
#
# This class is largely based of off the Analysis class of PalmettoBUG (see Analysis.py file), which in turn was heavily inspired / derived from CATALYST
# However, this class shares very limited similarity with CATALYST besides the panel / metadata files.
# See Assets / Other_License_Details.txt for more information on 3rd-party sources of code / ideas in this package.

import os
from pathlib import Path
from typing import Union
import warnings

import scipy
import anndata as ann
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from .._vendor import sigfig

warnings.filterwarnings("ignore", message = "Transforming to str index")   ## anndata implicit modification warning that is not necessary
warnings.filterwarnings("ignore", message = "Observation names are not unique")  ## anndata UserWarning that is not necessary

__all__ = ["WholeClassAnalysis"]

class WholeClassAnalysis:     
    '''
    This class handles the whole-class Analysis, where pixel regions are treated as if they are cell segmentation masks

    It has limited options compared to the standard experiment class that handles true single cell data. This class only has a few plot options 
    and a single statistics option, and no batch correction, dropping of samples, or scaling

    Args:
        directory (string or Path): the path to a folder containing /intensities and /regionprops subfolders,
                which would have been produced by running region measurements on the pixel classification maps
                generated by a pixel classifier. 

        classifier_df (pandas dataframe): the biological_labels.csv exported from the pixel classifier whose
                output is being used. Contains "labels", "class", and/or "merge" columns,
                which help associate region numbers in the images / regionpros & intensity csvs
                with the biological labels in the classifier. 

        metadata (pandas dataframe): analogous to the metadata csv file in a standard, single-cell analysis
                Contains the same, file_name, sample_id, patient_id, condition columns

        Analysis_panel (pandas dataframe): analogous to the Analysis_panel csv file in a standard, single-cell analysis.
                For example, contains columns for antigen and marker_class.

        class_type (string): one of -- "premerge", "merged" -- whether the outputs of the classifier are before or after
                merging (relevant for what column in classifier_df is used as the class, "merging" or "class" )

    Key Attributes:
        data (anndata.AnnData): the data, with data.X being a numpy array containing the channel information per "event" (per class per image)
                data.obs being derived from the inputted metadata and data.var being derived from the Analysis_panel

        class_labels (pandas DataFrame): this is the inputted classifier_df, which associates the class numbers with biological labels

        directory (str): The path to the folder where the analysis is to be initialized. Used to set up directories 
                (such as save_dir, data_table_dir), to export some files (input_tables_to_csv) and to find the expected
                intensities / regionprops csv files when loading the data. 

        save_dir (str): The path to where plots are saved by this class (when filename is provided in plotting functions)

        data_table_dir (str): The path to where data tables are saved by this class (when filename is provided to methods that produce
                dataframes such as statistics / exports) 
    '''
    def __init__(self, 
                 directory: Union[Path, str], 
                 classifier_df: pd.DataFrame, 
                 metadata: pd.DataFrame, 
                 Analysis_panel: pd.DataFrame, 
                 csv: Union[pd.DataFrame, None] = None,
                 ) -> None:
        '''
        '''
        self.directory = str(directory)
        self.directory  = self.directory.replace("\\" , "/")

        classifier_df['class'] = classifier_df['class'].astype('int')
        classifier_df['merging'] = classifier_df['merging'].astype('int')
        self.class_labels = classifier_df
        self._metadata = metadata
        try:
            old_index = Analysis_panel.index.copy()
            Analysis_panel.index = Analysis_panel['antigen']
            Analysis_panel = Analysis_panel.drop('Object', axis = 0)
            Analysis_panel.index = old_index    ## if present, remove the Object column from the panel
        except Exception:
            pass
        self._panel = Analysis_panel

        self.save_dir = self.directory + "/python_plots"
        self.data_table_dir = self.directory + "/Data_tables"
        if not os.path.exists(self.save_dir):
            os.mkdir(self.save_dir)

        if not os.path.exists(self.directory + "/Data_tables"):
            os.mkdir(self.directory + "/Data_tables")
        
        self.percent_areas = None
        if csv is None:
            self._load()
        else:
            self._load(csv = csv)

    def _load(self, csv: Union[pd.DataFrame, None] = None, arcsinh_cofactor = 5) -> None:
        '''Helper to the __init__ method: performs the loading and shaping of data during the initial load.'''
        metadata = self._metadata
        panel = self._panel
        panel.index = panel['antigen']

        if csv is not None:
            marker_class_included = False
            marker_class = csv.copy().iloc[-1,:]
            if np.array(marker_class == "na").sum() != 0: 
                marker_class_dict_rev = {"0.0" : 'none', "1.0" : 'type', "2.0" : ' state'}
                marker_class = marker_class[marker_class != "na"].astype('str').replace(marker_class_dict_rev)
                csv = csv.iloc[:-1, :]
                marker_class_included = True

            metadata_columns = ["class", "sample_id", "patient_id", "condition","file_name","Object"]
            to_obs = csv[[i for i in csv.columns if i in metadata_columns]].astype('str')
            not_obs = csv[[i for i in csv.columns if i not in metadata_columns]].astype('float')
            self._metadata = to_obs
            self._panel = pd.DataFrame()
            self._panel['antigen'] = list(not_obs.columns)
            self._panel['name'] = list(not_obs.columns)
            self._panel.index = self._panel['antigen']
            if marker_class_included is True:
                self._panel['marker_class'] = list(marker_class) + [1] ## Areas will always be assumed type
            else:
                self._panel['marker_class'] = 'type'
                print("All antigens from the csv have been set to 'type' in the panel file! \n"
                      "Open the Analysis_panel.csv file, edit, & reload the experiment to change this!")

            if len(to_obs["Object"].unique()) == len(self.class_labels['labels'].unique()):
                self.class_type = "premerge"
            elif len(to_obs["Object"].unique()) == (len(self.class_labels['labels'].unique()) - 1):
                self.class_type = "merged"
                self.class_labels = self.class_labels[self.class_labels['labels'] != 'background']

            self.data = ann.AnnData(X = np.array(not_obs), obs = to_obs, var = self._panel)
            self._object_column = to_obs['Object']
        
        else:
            csv_directory = self.directory  + "/intensities"
            csv_dir_names = [i for i in sorted(os.listdir(csv_directory)) if i.lower().find(".csv") != -1]

            if len(csv_dir_names) != len(self._metadata):
                print("Metadata file_name column and number of csv files in analysis do not match --\n" 
                      "only keeping data present in both for analysis!")
                new_fcs_filenames = []
                truth_array = np.zeros(len(self._metadata)).astype('bool')
                for i in list(self._metadata['file_name']):

                    if i in csv_dir_names:
                        new_fcs_filenames.append(i)
                        truth_array = truth_array  + np.array(self._metadata['file_name'] == i).astype('bool')

                self._metadata = self._metadata[truth_array]
                csv_dir_names = new_fcs_filenames

            csv_path_list = ["".join([csv_directory,"/",i]) for i in csv_dir_names]

            intensities = pd.DataFrame()
            length_of_images = [0]
            length_of_images2 = []
            tally = 0
            for i in csv_path_list:
                fcs_read_in = pd.read_csv(i)
                tally = tally + len(fcs_read_in.index)
                length_of_images.append(tally)
                length_of_images2.append(len(fcs_read_in.index))
                intensities = pd.concat([intensities, fcs_read_in], axis = 0)

            self._object_column = intensities['Object']
            intensities = intensities.drop("Object", axis = 1)

            if len(self._object_column.unique()) == len(self.class_labels['labels'].unique()):
                self.class_type = "premerge"
            elif len(self._object_column.unique()) == (len(self.class_labels['labels'].unique()) - 1):
                self.class_type = "merged"
                self.class_labels = self.class_labels[self.class_labels['labels'] != 'background']
            else:
                raise Exception("The number of classes in the data and the number of classes in the biological_labels.csv do not match! \n\n"
                                "(They are neither equal innumber nor does the biological_labels.csv has +1 class [background])")

            if self.class_type == "merged":
                zip_dict = {}
                for i,ii in zip(self.class_labels['merging'],self.class_labels['labels']):
                    zip_dict[i] = ii
            else:
                zip_dict = {}
                for i,ii in zip(self.class_labels['class'],self.class_labels['labels']):
                    zip_dict[i] = ii
            if arcsinh_cofactor > 0:
                exprs = pd.DataFrame(np.arcsinh(intensities / arcsinh_cofactor))
            else:
                exprs = pd.DataFrame(intensities)
            exprs.columns = panel["antigen"]

            metadata_long = pd.DataFrame()
            sample_id_array = np.zeros([0])
            counter = 0
            for i,ii in zip(length_of_images[:-1], length_of_images[1:]):
                slicer = np.full(shape = [ii - i], fill_value = counter)
                sample_id_array = np.append(sample_id_array,slicer)
                counter += 1
                
            metadata_long['sample_id'] = sample_id_array.astype('int')

            meta_file_dict = {}
            meta_patient_dict = {}
            meta_condition_dict = {}
            for i,ii in enumerate(metadata["sample_id"]):
                meta_file_dict[ii]= list(metadata.iloc[i])[0]
                meta_patient_dict[ii]= list(metadata.iloc[i])[2]
                meta_condition_dict[ii]= list(metadata.iloc[i])[3]

            filenames = metadata_long.replace(meta_file_dict)
            patient_ids = metadata_long.replace(meta_patient_dict)
            conditions = metadata_long.replace(meta_condition_dict)

            metadata_long['file_name'] = filenames
            metadata_long['patient_id'] = patient_ids
            metadata_long['condition'] = conditions
            metadata_long['Object'] = list((self._object_column).astype('int'))
            metadata_long['class'] = metadata_long["Object"].astype('int').replace(zip_dict).astype('str')
            metadata_long.index = exprs.index

            self._metadata["number_of_cells"] = length_of_images2

            self.data = ann.AnnData(X = exprs, var = panel, obs = metadata_long)
            self.data.obs["sample_id"] =  self.data.obs["sample_id"].astype("category")
            self.data.obs["patient_id"] =  self.data.obs["patient_id"].astype("category")
            self.data.obs["condition"] =  self.data.obs["condition"].astype("category")
            self.data.obs = self.data.obs.reset_index().drop("index", axis = 1)

            self.plot_percent_areas(filename = None, calculate_only = True)   ## now this adds the "areas" column to the metadata and 
                                                                                # saves the self.percent_areas table ready for plotting

            self.data.uns['counts'] = np.array(intensities)   ## this will be the storage place for the pre-arcsinh transformed counts data 
                                                                        # (never used by GUI)

    def input_tables_to_csv(self) -> None:
        '''Allows the saving of the primary csv files within this class to the disk inside the self.directory folder'''
        self.class_labels.to_csv(self.directory + "/classifier_labels.csv", index = False)
        self._metadata.to_csv(self.directory + "/metadata.csv", index = False)
        self._panel.to_csv(self.directory + "/Analysis_panel.csv", index = False)

    def plot_percent_areas(self, filename: Union[str, None] = None, calculate_only: bool = False) -> plt.figure:
        ''' 
        Plots a boxplots of percent class in each image, showing and comparing the distributions between conditions
        
        Returns the plot as a matplotlib figure
        
        Args:
            filename (str or None):
                If filename is specified, this will export the plot as a PNG file to self.save_dir/{filename}.png

            calculate_only (bool):
                Default == False. If True (& self.percent_areas == None), this method will not return anything, 
                but instead will perform the calculation of %pixel class in each ROI. This calculation will be 
                saved to self.percent_areas, where it can easily be plotted by this function later. 
                This is implemented to save time by meaning the calculations only have to be done once

        Returns:
            matplotlib.pyplot figure or None (returns None only if calculate_only == True, and no 
            prior calculation of the % areas has been done)
        '''
        if self.percent_areas is None:
            whole_regionprops_dir = self.directory  + "/regionprops"
            data_whole = self.data.copy()
            data_whole_df = pd.DataFrame(data_whole.X, columns = data_whole.var.index, index  = data_whole.obs.index)
            data_whole_df["sample_id"] = data_whole.obs["sample_id"]
            data_whole_df["patient_id"] = data_whole.obs["patient_id"]
            data_whole_df["condition"] = data_whole.obs["condition"]
            data_whole_df['Object'] = list((self._object_column).astype('int'))
            data_whole_df['class'] = data_whole.obs["class"]

            ind_var_to_sample_id = {}
            for i,ii in zip(data_whole_df['sample_id'].astype('str'), data_whole_df["condition"]):
                ind_var_to_sample_id[i] = ii 

            data_whole_means = data_whole_df.groupby(["class","sample_id"], observed = True).mean(numeric_only = True).reset_index()
            data_whole_means['condition'] = data_whole_means['sample_id'].astype('str').replace(ind_var_to_sample_id)
            data_whole_means = pd.melt(data_whole_means, id_vars = ['condition',"class","sample_id"])
            grand_csv = pd.DataFrame()
            regionprop_file_list = [i for i in sorted(os.listdir(whole_regionprops_dir)) if i.lower().find(".csv") != -1]
            for i in regionprop_file_list:
                csv = pd.read_csv(whole_regionprops_dir + "/" + i)
                grand_csv = pd.concat([grand_csv,csv], axis = 0)
                
            grand_csv['class'] = list(data_whole_df['class'])
            grand_csv['sample_id'] = list(data_whole_df['sample_id'])
            grand_csv['condition'] = list(data_whole_df['condition'])
            grand_csv['area'] = grand_csv['area'] / grand_csv['image_area']      ## Percentile of total image area
            
            ind_var_to_sample_id = {}
            for i,ii in zip(grand_csv['sample_id'], grand_csv["condition"]):
                ind_var_to_sample_id[i] = ii 

            percent_areas = grand_csv[['area','class','sample_id']].groupby(['class','sample_id'], observed = False).mean() * 100
            percent_areas = percent_areas.reset_index()
            percent_areas['sample_class'] = [i + ii  for i, ii in zip(percent_areas['class'].astype('str'),percent_areas['sample_id'].astype('str'))]
            percent_dict = {}
            for i,ii in zip(percent_areas['area'],percent_areas['sample_class']):
                percent_dict[ii] = float(i)
            percent_areas['condition'] = percent_areas['sample_id'].replace(ind_var_to_sample_id)
            data_whole_df['sample_class'] = [i + ii  for i, ii in zip(data_whole_df['class'].astype('str'),data_whole_df['sample_id'].astype('str'))]
            data_whole_df['areas******'] = data_whole_df['sample_class'].astype('str').replace(percent_dict).astype('float')
            data_whole_df = data_whole_df.drop('sample_class', axis = 1)

            if calculate_only:
                self.data.obs['areas'] = list(data_whole_df['areas******'])
                self.percent_areas = percent_areas
                return
        else:
            percent_areas = self.percent_areas

        figure = plt.figure()
        ax = plt.gca()
        sns.boxplot(percent_areas, x = 'class', y = 'area', hue = 'condition', ax = ax)
        ax.set_title("% Area of each class type in each condition:")
        if filename is not None:
            figure.savefig(self.save_dir + f"/{filename}.png", bbox_inches = "tight") 
        plt.close()
        return figure
    
    def plot_distribution_exprs(self, 
                                unique_class: Union[str, int], 
                                plot_type: str,
                                marker_class: str = "All",
                                filename: Union[str, None] = None,
                                ) -> sns.FacetGrid:
        '''
        Plots a Bar or Violin plot from the distribution of marker expression / %class in each sample_id, comparing conditions
        
        Args: 
            unique_class (string or integer):
                Indicates which pixel class to plot antigen expressions for

            plot_type (string):
                'Violin' or 'Bar' -- determines what kind of plot is created

            marker_class (string):
                'All', 'type', 'state', or 'none' (or any other marker_class in self.data.var['marker_class']). Determines which antigens are used in the plot
                By default, every antigen, regardless of marker_class is used ('All'). 

            filename:
                If specified, this funciton will additionally export the plot as a PNG file to self.save_dir/{filename}.png

        Returns:
            the plot as a seaborn FacetGrid (FacetGrid.figure --> a matplotlib figure)
        '''
        plot_type_dict = {"Bar":sns.barplot, "Violin":sns.violinplot}
        plot_type_fx = plot_type_dict[plot_type]
        text_size = 9
        palette = 'tab20'
        legend_loc = [1.1,0.75]
        col_wrap = 4
        height = 12
        width = 7

        data_whole = self.data.copy()
        panel = self._panel.copy()
        if marker_class != "All":
            data_whole = data_whole[:, panel["marker_class"] == marker_class]
        data_whole_df = pd.DataFrame(data_whole.X, columns = data_whole.var.index, index  = data_whole.obs.index)
        data_whole_df["sample_id"] = data_whole.obs["sample_id"]
        data_whole_df["patient_id"] = data_whole.obs["patient_id"]
        data_whole_df["condition"] = data_whole.obs["condition"]
        data_whole_df['Object'] = list((self._object_column).astype('int'))
        data_whole_df['class'] = data_whole.obs["class"]
        ind_var_to_sample_id = {}
        for i,ii in zip(data_whole_df['sample_id'].astype('str'), data_whole_df["condition"]):
            ind_var_to_sample_id[i] = ii 
        data_whole_means = data_whole_df.groupby(["class","sample_id"], observed = True).mean(numeric_only = True).reset_index()
        data_whole_means['condition'] = data_whole_means['sample_id'].astype('str').replace(ind_var_to_sample_id)
        data_whole_means = pd.melt(data_whole_means, id_vars = ['condition',"class","sample_id"])
        data_whole_means_subset = data_whole_means[data_whole_means['class'] == unique_class]
        griddy = sns.FacetGrid(data_whole_means_subset,
                                palette = palette, col = 'antigen', 
                                col_wrap = col_wrap, sharey = False, sharex = False, 
                                height = 1.5, aspect = 1.75, legend_out = True)
        griddy.map_dataframe(plot_type_fx, y = "value", 
                                hue = "condition", palette = 'tab20')
        griddy.add_legend()
        griddy.tick_params("x", labelrotation = 90)
        griddy.refline(y = 0)
        griddy.figure.set_size_inches(width, height)
        griddy.figure.suptitle(f"{unique_class}", y = 1.01, size = 12)
        sns.move_legend(griddy, loc = "center right", bbox_to_anchor = legend_loc)
        for i in griddy.axes:
            i.set_title(i.get_title(), y = 1.03, size = text_size)
            i.set_ylabel(i.get_ylabel(), size = text_size)
            i.set_xlabel(i.get_xlabel(), size = text_size)
            i.set_yticks(i.get_yticks(), labels = i.get_yticklabels(), size = text_size)
            i.set_xticks(i.get_xticks(), labels = i.get_xticklabels(), size = text_size)
        self.grid_object = griddy
        if filename is not None:
            griddy.savefig(self.save_dir + f"/{filename}.png", bbox_inches = "tight") 
        plt.close()
        return griddy
    
    def whole_marker_exprs_ANOVA(self, 
                                 marker_class: str = 'All', 
                                 groupby_column: str = 'class', 
                                 variable: str = 'condition', 
                                 statistic: str = 'mean',
                                 area: bool = True,
                                 ) -> pd.DataFrame:
        '''
        Calculates statistics comparing the conditions in the experiment using ANOVA on the expression of [marker_class] markers 
        and %area of each class

        Args:
            marker_class (string): which markers / antigens to test by ANOVA. one of -- "All", "type","state", "none". 

            groupby_column (string): which column the data will be grouped by for the purposed of running separate ANOVAs for each group 
                    (as this is whole-class analysis, should always be "class")

            variable (string): which column in self.data.obs will be trated as the column containing condition / group information

            statistic (string): one of --"ANOVA", "Kruskal" -- which statistical test (ANOVA, kruskal-wallis), and what aggregate statistic 
                    (mean/std or median/IQR, respectively) is calculated & displayed in the final dataframe

            area (bool): whether to also calculate an ANOVA comparing the %area of each class between the conditions (default is True)

        Returns:
            (pandas dataframe): the pandas dataframe contianing the statistical outputs of this test.
        '''
        statistic_to_test = {"ANOVA": scipy.stats.f_oneway, "Kruskal": scipy.stats.kruskal}
        test_func = statistic_to_test[statistic]
        statistic_to_agg_func = {"ANOVA":"mean", "Kruskal":"median"}
        statistic = statistic_to_agg_func[statistic]

        data_whole = self.data.copy()
        panel = self._panel.copy()
        if marker_class != "All":
            non_type_antigens = panel[panel["marker_class"] == marker_class]['antigen']
            data_whole = data_whole[:, panel["marker_class"] == marker_class]
        else:
            non_type_antigens = panel['antigen']
    
        data_whole_df = pd.DataFrame(data_whole.X, columns = data_whole.var.index, index  = data_whole.obs.index)
        data_whole_df["sample_id"] = data_whole.obs["sample_id"]
        data_whole_df["patient_id"] = data_whole.obs["patient_id"]
        data_whole_df["condition"] = data_whole.obs["condition"]
        data_whole_df['Object'] = list((self._object_column).astype('int'))
        data_whole_df['class'] = data_whole.obs["class"]
        ind_var_to_sample_id = {}
        for i,ii in zip(data_whole_df['sample_id'].astype('str'), data_whole_df["condition"]):
            ind_var_to_sample_id[i] = ii 
        data_whole_means = data_whole_df.groupby(["class","sample_id"], observed = True).mean(numeric_only = True).reset_index()
        data_whole_means['condition'] = data_whole_means['sample_id'].astype('str').replace(ind_var_to_sample_id)
        data_whole_means = pd.melt(data_whole_means, id_vars = ['condition',"class","sample_id"])
        grand_csv = pd.DataFrame()
        whole_regionprops_dir = self.directory  + "/regionprops"
        regionprops_list = [i for i in sorted(os.listdir(whole_regionprops_dir)) if i.lower().find(".csv") != -1]
        for i in regionprops_list:
            csv = pd.read_csv(whole_regionprops_dir + "/" + i)
            grand_csv = pd.concat([grand_csv,csv], axis = 0)
        grand_csv['class'] = list(data_whole_df['class'])
        grand_csv['sample_id'] = list(data_whole_df['sample_id'])
        grand_csv['condition'] = list(data_whole_df['condition'])
        grand_csv['area'] = grand_csv['area'] / grand_csv['image_area']      ## Percentile of total image area
        ind_var_to_sample_id = {}
        for i,ii in zip(grand_csv['sample_id'], grand_csv["condition"]):
            ind_var_to_sample_id[i] = ii 
        percent_areas = grand_csv[['area','class','sample_id']].groupby(['class','sample_id'], observed = False).mean() * 100
        percent_areas = percent_areas.reset_index()
        percent_areas['sample_class'] = [i + ii  for i, ii in zip(percent_areas['class'].astype('str'),percent_areas['sample_id'].astype('str'))]
        percent_dict = {}
        for i,ii in zip(percent_areas['area'],percent_areas['sample_class']):
            percent_dict[ii] = float(i)
        percent_areas['condition'] = percent_areas['sample_id'].replace(ind_var_to_sample_id)
        data_whole_df['sample_class'] = [i + ii  for i, ii in zip(data_whole_df['class'].astype('str'),data_whole_df['sample_id'].astype('str'))]
        data_whole_df['areas******'] = data_whole_df['sample_class'].astype('str').replace(percent_dict).astype('float')
        data_whole_df = data_whole_df.drop('sample_class', axis = 1)
        data = data_whole_df.drop('Object', axis = 1)
        conditions = data[variable].unique()
        merging_clusters = data[groupby_column].unique()
        if area is True:
            non_type_antigens = list(non_type_antigens) + ['areas******']
            non_type_antigens = pd.Series(non_type_antigens, name = 'antigen')
        data_df = data.copy()
        ind_var_to_sample_id = {}
        for i,ii in zip(data_df['sample_id'].astype('str'), data_df[variable]):
            ind_var_to_sample_id[i] = ii 
        
        if statistic == "median":
            stat_column_name = "H statistic"
            data_df = data_df.groupby(['sample_id', groupby_column], observed = False).median(numeric_only = True).fillna(0).reset_index()
        elif statistic == "mean":
            stat_column_name = "F statistic"
            data_df = data_df.groupby(['sample_id', groupby_column], observed = False).mean(numeric_only = True).fillna(0).reset_index()
        
        data_df[variable] = data_df['sample_id'].astype('str').replace(ind_var_to_sample_id)
        grand_condition_list = []
        for j,jj in enumerate(merging_clusters):
            merging_data = data_df[data_df[groupby_column] == jj]
        
            condition_list = []
            for i in conditions:
                condition_data = merging_data[merging_data[variable] == i]
                condition_list.append(condition_data.drop(['sample_id',groupby_column,variable], axis = 1))
            grand_condition_list.append(condition_list)
            ANOVA_f, ANOVA_p = test_func(*condition_list)
            if j == 0:
                merging_array_F = ANOVA_f[np.newaxis,:]
                merging_array_p = ANOVA_p[np.newaxis,:]
            else:
                merging_array_F = np.concatenate([merging_array_F,ANOVA_f[np.newaxis,:]], axis = 0)
                merging_array_p = np.concatenate([merging_array_p,ANOVA_p[np.newaxis,:]], axis = 0)

        p_adj_df = pd.DataFrame(scipy.stats.false_discovery_control(np.nan_to_num(merging_array_p) + 1e-25, 
                                                                    method = 'bh'), 
                                columns = non_type_antigens)
        
        p_adj_df[groupby_column] = merging_clusters
        p_adj_df = pd.melt(p_adj_df, id_vars = groupby_column)
        p_value_df = pd.DataFrame(merging_array_p, columns = non_type_antigens)
        p_value_df[groupby_column] = merging_clusters
        p_value_df = pd.melt(p_value_df, id_vars = groupby_column)
        F_values_df = pd.DataFrame(merging_array_F, columns = non_type_antigens)
        F_values_df[groupby_column] = merging_clusters
        F_values_df = pd.melt(F_values_df, id_vars = groupby_column)
        
        final_df = pd.DataFrame()
        final_df['antigen'] = p_adj_df['antigen']
        final_df[groupby_column] = p_adj_df[groupby_column]
        final_df['p_adj'] = [sigfig.round(i, 3, warn = False) for i in p_adj_df['value']]
        final_df['p_value'] = [sigfig.round(i, 3, warn = False) for i in p_value_df['value']]     
        final_df[stat_column_name] = [sigfig.round(i, 3, warn = False) for i in F_values_df['value']]
        final_df = final_df.groupby(['antigen',groupby_column], observed = False).mean()
        condition_dict = {}
        for jj,j in enumerate(grand_condition_list):
            for i,ii in enumerate(j):
                ii[groupby_column] = data_df[groupby_column].astype('str')
                if statistic == "mean":
                    merging_mean_condition = pd.melt(ii.groupby(groupby_column, 
                                                                observed = True).mean(numeric_only = True).T.reset_index(), 
                                                     id_vars = 'antigen')
                    merging_std_condition = pd.melt(ii.groupby(groupby_column, 
                                                               observed = True).std(numeric_only = True).T.reset_index(), 
                                                    id_vars = 'antigen')
                    spread_stat = 'stdev'
                elif statistic == "median":
                    merging_mean_condition = pd.melt(ii.groupby(groupby_column, 
                                                                observed = True).median(numeric_only = True).T.reset_index(),
                                                    id_vars = 'antigen')
                    merging_std_75 = pd.melt(ii.groupby(groupby_column, 
                                                               observed = True).quantile(q = 0.75, 
                                                                                         numeric_only = True).T.reset_index(),
                                                    id_vars = 'antigen')
                    merging_std_25 = pd.melt(ii.groupby(groupby_column,
                                                        observed = True).quantile(q = 0.25,
                                                                                    numeric_only = True).T.reset_index(),
                                            id_vars = 'antigen')
                    merging_std_condition = merging_mean_condition.copy()
                    merging_std_condition['value'] = merging_std_75['value'] - merging_std_25['value']
                    spread_stat = 'IQR'
                merging_mean_condition = merging_mean_condition.groupby(['antigen',groupby_column], observed = False).mean()   
                                ## These secondary groupby lines only serve to set up the multi-index (values should not change)
                merging_std_condition = merging_std_condition.groupby(['antigen',groupby_column], observed = False).mean()
                merging_mean_condition['stdev'] = merging_std_condition['value']
                if jj == 0:
                    condition_dict[str(i)] = merging_mean_condition
                else:
                    condition_dict[str(i)] = pd.concat([condition_dict[str(i)], merging_mean_condition], axis = 0)
                
        for i,ii in enumerate(condition_dict):
            final_df[f'avg {conditions[i]} {statistic} exprs'] = condition_dict[ii]['value']   
            final_df[f'avg {conditions[i]} {statistic} exprs'] = [sigfig.round(i, 3, warn = False) for i in final_df[f'avg {conditions[i]} {statistic} exprs']]
            final_df[f'{conditions[i]} {spread_stat}'] = condition_dict[ii]['stdev']
            final_df[f'{conditions[i]} {spread_stat}']  = [sigfig.round(i, 3, warn = False) for i in final_df[f'{conditions[i]} {spread_stat}']]
        self.cluster_expression_ANOVA_df = final_df.reset_index().sort_values(stat_column_name, ascending = False)
        return final_df.reset_index().sort_values(stat_column_name, ascending = False)

    def plot_heatmap(self, 
                     type_of_stat: str = "F statistic",
                     filename: Union[str, None] = None,
                     ) -> plt.figure:
        '''
        Plots a statistics heatmap. -Neg log(statistic) if the statistic is a p value instead of an F statistic
        '''
        heatmap_df = self.cluster_expression_ANOVA_df.pivot(index = 'antigen', columns = 'class', values = type_of_stat)
        heatmap_df = heatmap_df.astype('float')
        title_helper = ''
        if not ((type_of_stat == "F statistic") or (type_of_stat == "H statistic")):
            heatmap_df = heatmap_df.transform(np.log) * (-1)
            title_helper = "Neg log "
        figure2 = plt.figure()
        ax = plt.gca()
        sns.heatmap(heatmap_df, cmap = "coolwarm", ax = ax, linewidths = 0.01, vmin = 0, vmax = 7, xticklabels = True, yticklabels = True)
        ax.set_title(f"{title_helper}{type_of_stat} Heatmap")
        if filename is not None:
            figure2.savefig(self.save_dir + f"/{filename}.png", bbox_inches = "tight") 
        plt.close()
        return figure2
    
    #### shared with the main experiment class
    def export_data(self, 
                    subset_columns: Union[list[str], None] = None, 
                    subset_types: Union[list[list[str]], None] = None, 
                    groupby_columns: Union[list[str], None] = None, 
                    statistic: str = 'mean',
                    groupby_nan_handling: str = 'zero',
                    include_marker_class_row: bool = False,
                    untransformed: bool = False,
                    filename: Union[str, None] = None, 
                    ) -> pd.DataFrame:
        '''
        Exports currently loaded data from the Analysis, from self.data. 
        
        Preserves any previously performed scaling, dropped categories, & batch correction. Always of arcsinh(data / 5) transformed data. Can
        export the entirety of relevant self.data information, or export subsets of self.data, and/or export aggregate summary statistics for 
        groups within the data. 

        Args:
            subset_columns (list[str] or None): 
                a list of strings denoting the columns to subset the dataframe's rows on (here and in other arguments, non-string input is attepmted 
                to be cast to strings inside the function, as well as the corresponding column of the data). if this or subset_types is None, no subsetting occurs. 

            subset_types (list[list[str]] or None): 
                a list contains sub-lists for strings. The length of the upper list must be the length of
                the subset_columns list, as each sub-list contains strings corresponding to the rows to keep. 

                    As in: if subset_columns = ['column1', 'column3'] and subset_types = [['type2', 'type6'],['typeB', 'typeZ']],
                    then rows of type2 / type6 in column1 will be kept, and similarly rows of typeB / typeZ in column2.

                When > 1 columns / conditions are subsetted on, as in the above example, the rows that are kept are the union of 
                all the subsetting conditions WITHIN a given column, but the intersection BETWEEN what is kept from each column. 
                So in the above example, all rows of column1 == type2/6 that also have column2 == typeB/Z are the rows that are maintained.
                                                        
            groupby_columns (list[str] or None): 
                A list of strings indicating what columns of the data to groupby. If None, then grouping is not performed.
                Used like this:    self.data.obs.groupby(groupby_columns)              but on a dataframe containing the data.X values as well

            statistic (str): 
                Possible values: 'mean','median','sum','std','count'. Denotes the pandas groupby method to be used after grouping (ignored if groupby_columns is None).
                Numeric methods (mean, median, sum, std) are only applied to numeric columns, so only those columns + the groupby columns 
                will be in the final dataframe / csv
            
            groupby_nan_handling(str):
                'zero' or 'drop' -- when grouping the data whether to drop (nans), which usually represent non-existent category combinations or to 
                convert nans to zeros. Any other values of this parameter will cause NaNs to be left as-is in the data export
                Note that the default (and only option available in GUI) is 'zero', which converts ALL NaN values to 0, while the 'drop' option only drops
                rows where EVERY numerical value is NaN.
                By default, all possible groupby_columns combinations are included in the export (even if they are not present in the data, such cell types 
                not present in every ROI), This is the source of most NaN values. Notably, columnns in the metadata (not data.obs!) of the Analysis are given special 
                treatment to try to prevent non-existent experimental categories from having data exported (for example, each ROI / sample_id should have been 
                with a single condition, not every possible condition in the dataset). 

            include_marker_class_row (bool): 
                Whether to include the marker_class information as a row at the bottom of the table --> True to 
                include this row -- useful for reimport into PalmettoBUG.
                False to not include this row -- this is probably better for import into non-PalmettoBUG software for analysis,
                or at the least the user will need to remember to remove this row before analyzing!
                When the marker_class row is included, it is encoded as integers (to prevent mixed dtype issues/warnings on reload)
                
                    >>> 0 = 'none', 1 = 'type', 2 = 'state'

                metadata columns (which have no marker_class) have this row filled with 'na'. 
                NOT USED IN COMBINATION WITH GROUPING!

            untransformed (bool):
                if True, will export the untransformed (pre-arcsinh, pre-scaling, etc., etc.) data, from self.data.uns['count'].
                Provided so that the raw data is not difficult to recover, although not expected to be used frequently. Default == False. 

            filename: (str, or None): 
                the name of the csv file to save the exported dataframe inside the self.data_table_dir folder. If None, no export occurs, and the data table is only returned. 

        Returns:
            (pandas DataFrame) -- the pandas dataframe representing the exported data. 

        Inputs/Outputs:
            Outputs: 
                If filename is provided (is not None), then exports the data table to self.data_table_dir/filename.csv
        '''
        if filename is None:
            output_path =  None
        else:
            output_path = "".join([self.data_table_dir, "/", str(filename), ".csv"])
        data = self.data.copy()
        if untransformed:
            data.X = self.data.uns['counts'].copy()
        data.obs = data.obs.reset_index()    ## cell index included as 'index' column --> useful if dropping / filtering cells out of the dataset
        ## anndata to pd.DataFrame:
        data_points = pd.DataFrame(data.X)
        data_points.columns = data.var.index.astype('str')
        data.obs.columns = data.obs.columns.astype('str')
        data_col_list = [i for i in data.obs.columns]
        for i in data_col_list:
            if (include_marker_class_row is True) & (groupby_columns is None):
                data_points[str(i)] = list(data.obs[str(i)]) + ["na"]
            else:
                data_points[str(i)] = list(data.obs[str(i)])
                data_points[str(i)] = data_points[str(i)].astype(data.obs[str(i)].dtype)
        
        # subset:
        if (subset_types is not None) and (subset_columns is not None):
            counter1 = 0
            for i,ii in zip(subset_columns, subset_types):
                counter2 = 0
                for j in ii:
                    data_slice = data_points[data_points[str(i)].astype('str') == str(j)]    ## subsets for
                    if counter2 == 0:
                        new_data = data_slice.copy()
                        counter2 = 1
                    else:
                        new_data = new_data.merge(data_slice, how = 'outer')
                if counter1 == 0 :
                    counter1 = 1
                    output_df = new_data.copy()
                else:
                    output_df = output_df.merge(new_data, how = 'inner')
            data_points = output_df
            
        if groupby_columns is None:
            if output_path is not None:
                data_points.to_csv(str(output_path), index = False)
            return data_points

        else:
            extra_columns = None
            if len(groupby_columns) > 1:
                extra_columns = [i for i in groupby_columns if i in self.metadata.columns]
                if len(extra_columns) > 1:
                    extra_data_points = data_points[extra_columns]
                    def concat(*args):
                        return "_|_|_".join(*args)
                    data_points['use'] = extra_data_points.T.apply(concat)
                    groupby_columns = ['use'] + [i for i in groupby_columns if i not in self.metadata.columns]
                else:
                    extra_columns = None
      
            groupby_object = data_points.groupby(groupby_columns, observed = False)
    
            if statistic == 'mean':
                groupby_object = groupby_object.mean(numeric_only = True)
            if statistic =='median':
                groupby_object = groupby_object.median(numeric_only = True)
            if statistic =='sum':
                groupby_object = groupby_object.sum(numeric_only = True)
            if statistic =='std':
                groupby_object = groupby_object.std(numeric_only = True)
            if statistic =='count':
                groupby_object = groupby_object.count()
                groupby_object = pd.DataFrame(groupby_object[groupby_object.columns[0]])
                groupby_object.columns = ['count']
    
            groupby_object = groupby_object.reset_index()
            if statistic == 'count':
                pass
            else:
                backup_groupby = pd.DataFrame(groupby_object[groupby_columns], index = groupby_object.index)
                if groupby_nan_handling == 'drop':
                    groupby_object = groupby_object.drop(groupby_columns, axis = 1).dropna(how = 'all')
                elif groupby_nan_handling == 'zero':
                    groupby_object = groupby_object.drop(groupby_columns, axis = 1).fillna(0)

                groupby_object = pd.concat([backup_groupby, groupby_object], axis = 1)

            for i in data_col_list:
                if (i in groupby_object.columns.astype('str')) and (i not in groupby_columns):
                    groupby_object = groupby_object.drop(i, axis = 1)
            if extra_columns:
                def split(value, part = 0):
                    return value.split("_|_|_")[part]
                interim = pd.DataFrame()
                for i,ii in enumerate(extra_columns):
                    interim[ii] = groupby_object['use'].apply(split, part = i)
                groupby_object = pd.concat([interim, groupby_object], axis = 1)
                groupby_object = groupby_object.drop('use', axis = 1)
    
    
            if output_path is not None:
                groupby_object.to_csv(str(output_path), index = False)
    
            return groupby_object