#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date: Created on 26 Oct 2023 16:53
# @Author: Yao LI
# @File: spagrn/simulator.py
"""
SIMULATION
*magic, do not touch*
"""

import pandas as pd
import scanpy as sc
import anndata as an
import numpy as np
from random import sample
from typing import Union, Optional, List


class Simulator(object):
    """
    Down-stream
    handle scMultiSim outputs
    """

    def __init__(self):
        self.data = None
        self.tfs = None
        self.tf_names = None
        self.motif_names = None
        self.receptors = None
        self.ligands = None
        self.targets = None
        self.total_genes = None

        self.lr_gt = None
        self.grn_gt = None

        self.counts = None
        self.noise_ids = None  # noise genes
        self.real_ids = None  # real genes
        self.total_name = None
        self.total_id = None

        self.name_df = None

        self.coor = None
        self.celltypes = None

    def get_lr(self, df):
        """
        Load Ligand-Receptor ground truth if available
        :param df:
        :return:
        """
        self.lr_gt = df
        self.receptors = list(set(df.receptor))
        self.ligands = list(set(df.ligand))

    def get_tg(self, df):
        """
        Load TF-Target ground truth (must have)
        :param df:
        :return:
        """
        self.grn_gt = df
        self.targets = list(set(df['regulated.gene']))
        self.tfs = list(set(df['regulator.gene']))

    def recognize_ids(self):
        """
        for counts gene index
        :return:
        """
        if self.counts:
            self.total_genes = list(self.counts.index)
            self.noise_ids = [i for i in self.total_genes if 'gene' in i]
            self.real_ids = [i for i in self.total_genes if 'gene' not in i]
        else:
            raise ValueError("Expression matrix is not available, load expression matrix first")

    def add_coor(self, value: Union[str, pd.DataFrame]):
        if isinstance(value, pd.DataFrame):
            self.coor = value
        elif isinstance(value, str):
            self.coor = pd.read_csv(value, index_col=0)

    def add_celltypes(self, value: Union[str, pd.DataFrame]):
        if isinstance(value, pd.DataFrame):
            self.celltypes = value
        elif isinstance(value, str):
            self.celltypes = pd.read_csv(value, index_col=0)

    def load_exp(self, df: pd.DataFrame):
        """
        when loading only one csv file created by scMultiSim
        :param df:
        :return:
        """
        self.counts = df
        self.recognize_ids()

    def load_multi_samples(self):
        """
        when loading multiple csv files generated by scMultiSim
        :return:
        """
        # Load in needed data
        df_list = []
        n = 0
        for num, tf in enumerate(self.tfs[:5]):
            df2 = pd.read_csv(f'counts_{tf}.csv', index_col=0)
            new_cell_num = df2.shape[1]
            new_cell_names = [f'cell{x}' for x in list(range(n + 1, n + new_cell_num + 1))]
            df2.columns = new_cell_names
            n = n + new_cell_num
            df_list.append(df2)
        # merge into one dataframe
        df = pd.concat(df_list, axis=1).fillna(0).astype(int)
        df3 = pd.read_csv('counts_addition.csv', index_col=0)
        gene_index = [i if 'gene' not in i else f'{i}3' for i in list(df3.index)]
        df3.index = gene_index
        df = pd.concat([df, df3]).fillna(0).astype(int)
        # self.noise_ids = [i for i in list(df.index) if 'gene' in i]
        self.counts = df
        self.recognize_ids()
        return df

    @staticmethod
    def load_multi(lists):
        """
        when loading multiple csv files generated by scMultiSim
        :return:
        """
        # Load in needed data
        df_list = []
        n = 0
        for num, i in enumerate(lists):
            df2 = pd.read_csv(f'counts_{i}.csv', index_col=0)
            new_cell_num = df2.shape[1]
            new_cell_names = [f'cell{x}' for x in list(range(n + 1, n + new_cell_num + 1))]
            df2.columns = new_cell_names
            n = n + new_cell_num
            df_list.append(df2)
        # merge into one dataframe
        df = pd.concat(df_list, axis=1).fillna(0).astype(int)
        df3 = pd.read_csv('counts_addition.csv', index_col=0)
        gene_index = [i if 'gene' not in i else f'{i}3' for i in list(df3.index)]
        df3.index = gene_index
        df = pd.concat([df, df3]).fillna(0).astype(int)
        return df

    def load_multiple_gts(self):
        """
        multiple GRN gts
        :return:
        """
        pass

    @staticmethod
    def get_dir(gt: pd.DataFrame, regulators, regulator_key='regulator.gene', value_key='regulator.effect',
                regulated_key='regulated.gene'):
        """
        Create dictionaries to map TF names to motif names and TFs to TF names
        :param gt:
        :param regulators: regulator: TF or ligand
        :param regulator_key:
        :param value_key:
        :param regulated_key:
        :return:
        """
        dir = {}
        for reg in regulators:
            sub = gt[gt[regulator_key] == reg]
            sorted_sub = sub.sort_values(by=value_key, ascending=False)
            dir[reg] = list(sorted_sub[regulated_key])
        return dir

    def assign_gene_names(self, rdb: pd.DataFrame, tf_id_dir, tf_motif_dir):
        grn_dir = Simulator.get_dir(self.grn_gt, self.tfs)
        # 1. for targets:
        # tf targets names, assign real names to tf targets. choose top n genes for each motif(TF).
        for tf in self.tfs:
            current_len = len(self.total_name)  # find current total_name number
            one_motif = tf_motif_dir[tf_id_dir[tf]]  # for each motif, find its gene rankings
            # find top genes for the motif
            sub = rdb[rdb.motifs == one_motif].drop('motifs', axis=1)  # remove motif column, only sort gene columns
            sorted_sub = sub.sort_values(by=(sub.index.values[0]), axis=1)  # sort columns (genes) by rank numbers
            top_num = len(grn_dir[tf])
            self.total_id += grn_dir[tf]
            # match real gene name with gene id
            for tg in sorted_sub.columns:
                if tg not in self.total_name:
                    self.total_name.append(tg)
                if len(self.total_name) == (current_len + top_num):
                    break

        # at this point, total names are unique, total ids have duplicates
        # remove duplicates
        name_df1 = pd.DataFrame({'id': self.total_id,
                                 'name': list(self.total_name)}).drop_duplicates(subset='id', keep='first').astype(str)
        self.total_name = list(name_df1.name)
        self.total_id = list(name_df1.id)
        # at this point, len(set(total_name)) should equal to len(set(total_id))
        assert name_df1.duplicated().sum().sum() == 0
        self.name_df = name_df1
        return name_df1

    def assign_name_addition(self, rdb, gene_id_list: list):
        """
        based on rest_rdb_names, when rest_rdb_names =
        :param rdb:
        :param gene_id_list: e.g noise_ids
        :return:
        """
        # PART TWO
        # some of the gene names are used. find out the rest unused gene names,
        # to avoid assign used gene name to noise genes and ligand and receptors
        total_rdb_names = set(rdb.columns) - set('motifs')
        rest_rdb_names = list(total_rdb_names - set(self.total_name))
        # 2. for noise genes
        noise_genes_names = sample(rest_rdb_names, k=len(gene_id_list))
        # !! random.choices returns duplicated items, use random.sample instead
        # add noise names into total_name
        self.total_name += noise_genes_names
        self.total_id += gene_id_list

        # create names df
        name_df = pd.DataFrame({'id': self.total_id,
                                'name': list(self.total_name)}).drop_duplicates(subset='id', keep='first').astype(str)
        assert name_df.duplicated().sum().sum() == 0
        self.name_df = name_df
        return name_df

    def to_anndata(self):
        """

        :param celltypes: cell types
        :param coor: cell coordinates
        :param name_df: simulated ID - gene name
        :return: anndata.Anndata
        """
        df = self.counts.T
        adata = sc.AnnData(df)
        adata.obs['cells'] = list(df.index)
        adata.var['genes'] = list(df.columns)
        if self.celltypes:
            adata.obs['celltype'] = self.celltypes['cell.type']
        if self.coor:
            adata.obsm['spatial'] = self.coor.iloc[:int(df.shape[0])]
        data_genes = adata.var.copy()
        data_genes = data_genes['genes'].replace(list(self.name_df['id']), list(self.name_df['name']))
        data_genes = data_genes.to_frame()
        adata.var = data_genes
        adata.var_names = adata.var['genes']
        return adata

    @staticmethod
    def subset(adata: an.AnnData, n_samples: int = 100, label: str = 'celltype',
               sub_clusters: Optional[list] = None) -> an.AnnData:
        """
        Subset samples
        :param adata:
        :param n_samples: number of samples to keep. should be smaller than total sample size.
        :param label:
        :param sub_clusters:
        :return:

        Example:
            subset(adata, n_samples=300, label='annotation', sub_clusters=['ct1','ct2']
        """
        l = []
        if sub_clusters:
            adata_cluster_sampled = adata[adata.obs[label].isin(sub_clusters)].copy()
        else:
            adata_cluster_sampled = adata.copy()

        for k, v in adata_cluster_sampled.obs.groupby(label).indices.items():
            l.append(np.random.choice(v, n_samples, replace=False))

        return adata_cluster_sampled[np.concatenate(l)]
