import pandas as pd
import numpy as np

from celescope.tools.plotly_plot import Tsne_plot
from celescope.tools.step import Step
from celescope.tools.capture.__init__ import SUM_UMI_COLNAME, LOG2_COLNAME
from celescope.tools import analysis_wrapper

def get_opts_analysis(parser, sub_program):
    if sub_program:
        parser.add_argument('--filter_umi_file', help='filter umi file', required=True)
        analysis_wrapper.get_opts_analysis_match(parser, sub_program)


class Analysis(Step):

    def __init__(self, args, display_title='Analysis'):
        super().__init__(args, display_title)
        self.df_filter_umi = pd.read_csv(args.filter_umi_file, index_col=0)

        # out
        self.df_tsne_file = f'{self.out_prefix}_UMI_tsne.csv'
        self.df_tsne = None

    def add_tsne_info(self):
        """
        add tsne coord and log2 column
        """
        report_runner = analysis_wrapper.Report_runner(self.args)
        df_tsne, _df_marker = report_runner.get_df()
        self.df_tsne = pd.merge(df_tsne, self.df_filter_umi, left_index=True, right_index=True)
        self.df_tsne[LOG2_COLNAME] = np.log2(self.df_tsne[SUM_UMI_COLNAME] + 1)
        self.df_tsne.to_csv(self.df_tsne_file)

    def add_cluster_metrics(self):
        self.add_help_content('cluster 1,2,3...', 'number of positive cells(sum_UMI > 0) in each cluster after filtering')

        df_cluster_all = self.df_tsne.groupby("cluster").count()

        df_positive = self.df_tsne[self.df_tsne[SUM_UMI_COLNAME] > 0]
        df_cluster_positive = df_positive.groupby("cluster").count()

        for index, row in df_cluster_positive.iterrows():
            self.add_metric(
                name=f'cluster {index}',
                value=int(row['tSNE_1']),
                total=int(df_cluster_all.loc[index, 'tSNE_1']),
            )

    def run(self):
        self.add_tsne_info()
        self.add_cluster_metrics()

        tsne_cluster = Tsne_plot(self.df_tsne, 'cluster').get_plotly_div()
        self.add_data(tsne_cluster=tsne_cluster)

        tsne_plot = Tsne_plot(self.df_tsne, LOG2_COLNAME, discrete=False)
        tsne_plot.set_color_scale(['LightGrey', 'Orange', 'Red'])
        tsne_feature = tsne_plot.get_plotly_div()
        self.add_data(tsne_feature=tsne_feature)