import numpy as np
import pandas as pd
import scipy.io
import scipy.sparse
import glob
import scanpy as sc
import anndata as ad

from celescope.tools import utils
from celescope.tools.step import Step, s_common
from celescope.__init__ import HELP_DICT
from celescope.tools.matrix import CountMatrix, Features


TAG_COL = "tag_name"


def filter_fun(df):
    """
    Filter data according to mean and standard deviation
    """
    df["filter_num"] = np.mean(df, axis=1) + np.std(df, axis=1)
    for index, _ in df.iterrows():
        df.loc[index] = df.loc[index].mask(
            df.loc[index] < df.loc[index]["filter_num"], 0
        )
    df.drop(columns="filter_num", inplace=True)
    return df


class CountMatrix2(CountMatrix):
    @classmethod
    def from_dataframe2(cls, df, features: Features, barcodes=None):
        matrix = df
        mtx = scipy.sparse.coo_matrix(matrix, matrix.shape, int)
        return cls(features, barcodes, mtx)


class Count_cite(Step):
    def __init__(self, args, display_title):
        super().__init__(args, display_title)

        self.df_read_count = pd.read_csv(args.read_count_file, sep="\t", index_col=0)

        self.match_dict = utils.parse_match_dir(args.match_dir)
        self.match_barcode = self.match_dict["match_barcode"]
        self.match_matrix_dir = self.match_dict["matrix_dir"]

        # input
        self.tsne_coord = glob.glob(f"{args.match_dir}/*analysis*/*tsne_coord.tsv")[0]

        # out
        self.mtx = f"{self.out_prefix}_citeseq.mtx.gz"
        self.raw_matrix_dir = f"{self.out_prefix}_raw_citeseq_matrix"
        self.filtered_matrix_dir = f"{self.out_prefix}_filtered_citeseq_matrix"
        self.filtered_tsne_coord = f"{self.out_prefix}_filtered_tsne_coord.tsv"

    @utils.add_log
    def run(self):
        mapped_read = int(self.df_read_count["read_count"].sum())

        # in cell
        df_read_count_in_cell = self.df_read_count[
            self.df_read_count.index.isin(self.match_barcode)
        ]
        mapped_read_in_cell = int(df_read_count_in_cell["read_count"].sum())
        self.add_metric(
            name="Mapped Reads in Cells",
            value=mapped_read_in_cell,
            total=mapped_read,
        )

        tag_names = df_read_count_in_cell[TAG_COL].unique()
        features_raw = Features(tag_names)

        # raw_matrix
        raw_citeseq_matrix = CountMatrix.from_dataframe(
            df_read_count_in_cell,
            features_raw,
            barcodes=self.match_barcode,
            row=TAG_COL,
            column="barcode",
            value="UMI",
        )
        rna_matrix = CountMatrix.from_matrix_dir(matrix_dir=self.match_matrix_dir)
        raw_merged_matrix = rna_matrix.concat_by_barcodes(raw_citeseq_matrix)
        raw_merged_matrix.to_matrix_dir(self.raw_matrix_dir)

        # UMI
        df_UMI_in_cell = (
            df_read_count_in_cell.reset_index()
            .groupby(["barcode", TAG_COL])
            .agg({"UMI": "count"})
        )

        df_UMI_in_cell = df_UMI_in_cell.reset_index()
        df_UMI_in_cell = df_UMI_in_cell.pivot(
            index="barcode", columns=TAG_COL, values="UMI"
        )
        df_cell = pd.DataFrame(index=self.match_barcode)
        df_UMI_cell = pd.merge(
            df_cell, df_UMI_in_cell, how="left", left_index=True, right_index=True
        )

        # fillna
        df_UMI_cell.fillna(0, inplace=True)
        df_UMI_cell = df_UMI_cell.astype(int)
        df_UMI_cell_out = df_UMI_cell.T
        df_UMI_cell_out.to_csv(self.mtx, sep="\t", compression="gzip")

        # fix
        ## 1.normalize;2.filter
        ### normalize
        obs = pd.DataFrame(index=df_UMI_cell_out.columns)
        var = pd.DataFrame(
            df_UMI_cell_out.index, index=df_UMI_cell_out.index, columns=["ADT"]
        )
        mdata_citeseq = ad.AnnData(np.array(df_UMI_cell_out.T), obs=obs, var=var)
        sc.pp.normalize_total(
            mdata_citeseq,
            target_sum=1e4,
            inplace=True,
        )
        sc.pp.log1p(mdata_citeseq)
        ### filter
        df_filtered = filter_fun(mdata_citeseq.to_df().T)
        features_filtered = Features(df_filtered.index.to_list())
        filtered_citeseq_matrix = CountMatrix2.from_dataframe2(
            df_filtered, features_filtered, barcodes=self.match_barcode
        )
        filtered_merged_matrix = rna_matrix.concat_by_barcodes(filtered_citeseq_matrix)
        filtered_merged_matrix.to_matrix_dir(self.filtered_matrix_dir)

        # filtered_tsne.csv
        df_tsne = pd.read_csv(self.tsne_coord, sep="\t")
        if "Unnamed: 0" in df_tsne.columns:
            df_tsne.rename(columns={"Unnamed: 0": "barcode"}, inplace=True)
            df_tsne = df_tsne.set_index("barcode")

        df_citeseq = mdata_citeseq.to_df()
        df_tsne = pd.concat([df_tsne, df_citeseq], axis=1)
        df_tsne.fillna(0, inplace=True)
        df_tsne.to_csv(self.filtered_tsne_coord, sep="\t")

        # UMI
        UMIs = df_UMI_cell.apply(sum, axis=1)
        median_umi = round(np.median(UMIs), 2)
        mean_umi = round(np.mean(UMIs), 2)
        self.add_metric(
            name="Median UMI per Cell",
            value=float(median_umi),
        )
        self.add_metric(
            name="Mean UMI per Cell",
            value=float(mean_umi),
        )


def get_opts_count_cite(parser, sub_program):
    if sub_program:
        parser.add_argument("--match_dir", help=HELP_DICT["match_dir"], required=True)
        parser.add_argument("--read_count_file", help="tag read count file")
        s_common(parser)


def count_cite(args):
    with Count_cite(args, display_title="Count") as runner:
        runner.run()
