from typing import Optional
import gc
import anndata
import numpy as np
import pandas as pd
from scipy import sparse
from scipy.spatial import distance_matrix

from .._optimal_transport import cot_combine_sparse


### MetaChat cell communication
class CellCommunication(object):

    def __init__(self,
        adata,
        df_MetaSen,
        LRC_type,
        dis_thr,
        cost_scale,
        cost_type
    ):
        
        # Find overlap metabolites and sensors in df_MetaSen
        data_var = set(adata.var_names)
        self.mets = list(set(df_MetaSen["Metabolite"]).intersection(data_var))
        self.sens = list(set(df_MetaSen["Sensor"]).intersection(data_var))

        # Generate an infinite matrix A. If the metabolite and the sensor can interact, 
        # let the corresponding position be 1
        A = np.inf * np.ones([len(self.mets), len(self.sens)], float)
        
        LRC = {}
        for i in range(len(df_MetaSen)):
            tmp_met = df_MetaSen.loc[i,"Metabolite"]
            tmp_sen = df_MetaSen.loc[i,"Sensor"]
            ## plus
            LRC[tmp_met] = df_MetaSen.loc[i,"Long.Range.Channel"]
            if tmp_met in self.mets and tmp_sen in self.sens:
                if cost_scale is None:
                    A[self.mets.index(tmp_met), self.sens.index(tmp_sen)] = 1.0
                else:
                    A[self.mets.index(tmp_met), self.sens.index(tmp_sen)] = cost_scale[(tmp_met, tmp_sen)]
        self.A = A.copy()
        self.LRC = LRC.copy()

        LRC_type = ["No"] + LRC_type
        self.LRC_type = LRC_type.copy()

        # Generate expression matrices of metabolites and sensors for all spots 
        self.S = adata[:,self.mets].X.toarray()
        self.D = adata[:,self.sens].X.toarray()
        
        # The dictionary approach to storing the distance matrix, since there are multiple LRC channels,
        # allows you to store the new distances generated by the corresponding channel in dmat according
        # to the name of the channel to be inferred.   
        if cost_type == 'euc':
            dmat = {}
            for tLRC in LRC_type:
                dmat[tLRC] = adata.obsp['spatial_distance_LRC_' + tLRC].copy()
        elif cost_type == 'euc_square':
            dmat = {}
            for tLRC in LRC_type:
                dmat[tLRC] = adata.obsp['spatial_distance_LRC_' + tLRC].copy() ** 2   
        self.M = dmat.copy()
        
        if np.isscalar(dis_thr):
            if cost_type == 'euc_square':
                dis_thr = dis_thr ** 2
            self.cutoff = float(dis_thr) * np.ones_like(A)
        elif type(dis_thr) is dict:
            self.cutoff = np.zeros_like(A)
            for i in range(A.shape[0]):
                for j in range(A.shape[1]):
                    if A[i,j] > 0:
                        if cost_type == 'euc_square':
                            self.cutoff[i,j] = dis_thr[(self.mets[i], self.sens[j])] ** 2
                        else:
                            self.cutoff[i,j] = dis_thr[(self.mets[i], self.sens[j])]
        self.nmet = self.S.shape[1]; self.nsen = self.D.shape[1]
        self.npts = adata.shape[0]

    def run_cot_signaling(self,
        cot_eps_p=1e-1, 
        cot_eps_mu=None, 
        cot_eps_nu=None, 
        cot_rho=1e1, 
        cot_nitermax=1e4, 
        cot_weights=(0.25,0.25,0.25,0.25)
    ):
        self.comm_network = cot_combine_sparse(self.S, self.mets, self.D, self.A, self.M, self.LRC, self.LRC_type, self.cutoff, \
            eps_p=cot_eps_p, eps_mu=cot_eps_mu, eps_nu=cot_eps_nu, rho=cot_rho, weights=cot_weights, nitermax=cot_nitermax)


def assign_distance(adata, dmat=None):
    if dmat is None:
        adata.obsp["spatial_distance"] = distance_matrix(adata.obsm["spatial"], adata.obsm["spatial"])
    else:
        adata.obsp["spatial_distance"] = dmat

def metabolic_communication(
    adata: anndata.AnnData, 
    database_name: str = None, 
    df_MetaSen: pd.DataFrame = None,
    LRC_type: list = None,
    dis_thr: Optional[float] = None, 
    cost_scale: Optional[dict] = None, 
    cost_type: str = 'euc',
    cot_eps_p: float = 1e-1, 
    cot_eps_mu: Optional[float] = None, 
    cot_eps_nu: Optional[float] = None, 
    cot_rho: float =1e1, 
    cot_nitermax: int = 10000, 
    cot_weights: tuple = (0.25,0.25,0.25,0.25),
    copy: bool = False
):
    """
    Infer spatial communication.

    Parameters
    ----------
    adata
        The data matrix of shape ``n_obs`` × ``n_var``.
        Rows correspond to cells or spots and columns to genes.
        If the spatial distance is absent in ``.obsp['spatial_distance']``, Euclidean distance determined from ``.obsm['spatial']`` will be used.
    database_name
        Name of the ligand-receptor interaction database. Will be included in the keywords for anndata slots.
    df_ligrec
        A data frame where each row corresponds to a ligand-receptor pair with ligands, receptors, and the associated signaling pathways in the three columns, respectively.
    pathway_sum
        Whether to sum over all ligand-receptor pairs of each pathway.
    heteromeric
        Whether the ligands or receptors are made of heteromeric complexes.
    heteromeric_rule
        Use either 'min' (minimum) or 'ave' (average) expression of the components as the level for the heteromeric complex.
    heteromeric_delimiter
        The character in ligand and receptor names separating individual components.
    dis_thr
        The threshold of spatial distance of signaling.
    cost_scale
        Weight coefficients of the cost matrix for each ligand-receptor pair, e.g. cost_scale[('ligA','recA')] specifies weight for the pair ligA and recA.
        If None, all pairs have the same weight. 
    cost_type
        If 'euc', the original Euclidean distance will be used as cost matrix. If 'euc_square', the square of the Euclidean distance will be used.
    cot_eps_p
        The coefficient of entropy regularization for transport plan.
    cot_eps_mu
        The coefficient of entropy regularization for untransported source (ligand). Set to equal to cot_eps_p for fast algorithm.
    cot_eps_nu
        The coefficient of entropy regularization for unfulfilled target (receptor). Set to equal to cot_eps_p for fast algorithm.
    cot_rho
        The coefficient of penalty for unmatched mass.
    cot_nitermax
        Maximum iteration for collective optimal transport algorithm.
    cot_weights
        A tuple of four weights that add up to one. The weights corresponds to four setups of collective optimal transport: 
        1) all ligands-all receptors, 2) each ligand-all receptors, 3) all ligands-each receptor, 4) each ligand-each receptor.
    smooth
        Whether to (spatially) smooth the gene expression for identifying more global signaling trend.
    smth_eta
        Kernel bandwidth for smoothing
    smth_nu
        Kernel sharpness for smoothing
    smth_kernel
        'exp' exponential kernel. 'lorentz' Lorentz kernel.
    copy
        Whether to return a copy of the :class:`anndata.AnnData`.

    Returns
    -------
    adata : anndata.AnnData
        Signaling matrices are added to ``.obsp``, e.g., for a LR interaction database named "databaseX", 
        ``.obsp['commot-databaseX-ligA-recA']``
        is a ``n_obs`` × ``n_obs`` matrix with the *ij* th entry being the "score" of 
        cell *i* sending signal to cell *j* through ligA and recA. If ``pathway_sum==True``, the cell-by-cell signaling matrix for a pathway "pathwayX" 
        will be stored in ``.obsp['commot-databaseX-pathwayX']``.
        The marginal sums (sender and receiver) of the signaling matrices are stored in ``.obsm['commot-databaseX-sum-sender']`` and ``.obsm['commot-databaseX-sum-receiver']``.
        Metadata of the analysis is added to ``.uns['commot-databaseX-info']``.
        If copy=True, return the AnnData object and return None otherwise.
        
    Examples
    --------
    >>> import anndata
    >>> import commot as ct
    >>> import pandas as pd
    >>> import numpy as np
    >>> # create a toy data with four cells and four genes
    >>> X = np.array([[1,0,0,0],[2,0,0,0],[0,1,1,4],[0,3,4,1]], float)
    >>> adata = anndata.AnnData(X=X, var=pd.DataFrame(index=['LigA','RecA','RecB','RecC']))
    >>> adata.obsm['spatial'] = np.array([[0,0],[1,0],[0,1],[1,2]], float)
    >>> # create a toy ligand-receptor database with two pairs sharing the same ligand
    >>> df_ligrec = pd.DataFrame([['LigA','RecA_RecB'],['LigA','RecC']])
    >>> df_ligrec = pd.DataFrame([['LigA','RecA_RecB','PathwayA'],['LigA','RecC','PathwayA']])
    >>> # infer CCC with a distance threshold of 1.1. Due the threshold, the last cell won't be communicating with other cells.
    >>> ct.tl.spatial_communication(adata, database_name='toyDB',df_ligrec=df_ligrec, pathway_sum=True, heteromeric=True, dis_thr=1.1)
    >>> adata
    AnnData object with n_obs × n_vars = 4 × 4
        uns: 'commot-toyDB-info'
        obsm: 'spatial', 'commot-toyDB-sum-sender', 'commot-toyDB-sum-receiver'
        obsp: 'commot-toyDB-LigA-RecA_RecB', 'commot-toyDB-LigA-RecC', 'commot-toyDB-PathwayA', 'commot-toyDB-total-total'
    >>> print(adata.obsp['commot-toyDB-LigA-RecA_RecB'])
    (0, 2)	0.600000000000005
    >>> print(adata.obsp['commot-toyDB-LigA-RecC'])
    (0, 2)	0.9000000000039632
    >>> print(adata.obsm['commot-toyDB-sum-receiver'])
    r-LigA-RecA_RecB  r-LigA-RecC  r-total-total  r-PathwayA
    0               0.0          0.0            0.0         0.0
    1               0.0          0.0            0.0         0.0
    2               0.6          0.9            1.5         1.5
    3               0.0          0.0            0.0         0.0
    >>> print(adata.obsm['commot-toyDB-sum-sender'])
    s-LigA-RecA_RecB  s-LigA-RecC  s-total-total  s-PathwayA
    0               0.6          0.9            1.5         1.5
    1               0.0          0.0            0.0         0.0
    2               0.0          0.0            0.0         0.0
    3               0.0          0.0            0.0         0.0
    >>> print(adata.obsp['commot-toyDB-PathwayA'])
    (0, 2)	1.5000000000039682
    >>> print(adata.obsp['commot-toyDB-total-total'])
    (0, 2)	1.5000000000039682
    >>> adata.uns['commot-toyDB-info']
    {'df_ligrec':   ligand   receptor   pathway
    0   LigA  RecA_RecB  PathwayA
    1   LigA       RecC  PathwayA, 'distance_threshold': 1.1}
    
    
    """

    assert database_name is not None, "Please give a database_name"
    assert df_MetaSen is not None, "Please give a Metabolite-Sensor database"

    # remove unavailable genes or metabolites from df_MetaSen
    data_var = list(adata.var_names)
    tmp_MetaSen = []
    for i in range(df_MetaSen.shape[0]):
        if df_MetaSen.loc[i,"Metabolite"] in data_var and df_MetaSen.loc[i,"Sensor"] in data_var:
            tmp_MetaSen.append(df_MetaSen.loc[i,:])
    tmp_MetaSen = np.array(tmp_MetaSen, str)
    df_MetaSen_filtered = pd.DataFrame(data = tmp_MetaSen)
    df_MetaSen_filtered.columns = df_MetaSen.columns.copy()

    # Drop duplicate pairs
    df_MetaSen_filtered = df_MetaSen_filtered.drop_duplicates()
    adata.uns["Metabolite_Sensor_filtered"] = df_MetaSen_filtered.copy()
    print("There are %d pairs were found from the spatial data." %df_MetaSen_filtered.shape[0])

    model = CellCommunication(adata,
        df_MetaSen_filtered,
        LRC_type,
        dis_thr,
        cost_scale, 
        cost_type
    )
    model.run_cot_signaling(cot_eps_p=cot_eps_p, 
        cot_eps_mu = cot_eps_mu, 
        cot_eps_nu = cot_eps_nu, 
        cot_rho = cot_rho, 
        cot_nitermax = cot_nitermax, 
        cot_weights = cot_weights
    )

    adata.uns['MetaChat-'+database_name+'-info'] = {}
    adata.uns['MetaChat-'+database_name+'-info']['distance_threshold'] = dis_thr

    ncell = adata.shape[0]
    X_sender = np.empty([ncell,0], float)
    X_receiver = np.empty([ncell,0], float)
    col_names_sender = []
    col_names_receiver = []
    tmp_mets = model.mets
    tmp_sens = model.sens
    S_total = sparse.csr_matrix((ncell, ncell), dtype=float)

    for (i,j) in model.comm_network.keys():
        S = model.comm_network[(i,j)]
        adata.obsp['MetaChat-'+database_name+'-'+tmp_mets[i]+'-'+tmp_sens[j]] = S
        S_total = S_total + S
        lig_sum = np.array(S.sum(axis=1))
        rec_sum = np.array(S.sum(axis=0).T)
        X_sender = np.concatenate((X_sender, lig_sum), axis=1)
        X_receiver = np.concatenate((X_receiver, rec_sum), axis=1)
        col_names_sender.append("s-%s-%s" % (tmp_mets[i], tmp_sens[j]))
        col_names_receiver.append("r-%s-%s" % (tmp_mets[i], tmp_sens[j]))

    X_sender = np.concatenate((X_sender, X_sender.sum(axis=1).reshape(-1,1)), axis=1)
    X_receiver = np.concatenate((X_receiver, X_receiver.sum(axis=1).reshape(-1,1)), axis=1)
    col_names_sender.append("s-total-total")
    col_names_receiver.append("r-total-total")
    
    adata.obsp['MetaChat-'+database_name+'-total-total'] = S_total

    df_sender = pd.DataFrame(data=X_sender, columns=col_names_sender, index=adata.obs_names)
    df_receiver = pd.DataFrame(data=X_receiver, columns=col_names_receiver, index=adata.obs_names)
    adata.obsm['MetaChat-'+database_name+'-sum-sender'] = df_sender
    adata.obsm['MetaChat-'+database_name+'-sum-receiver'] = df_receiver

    del model
    gc.collect()

    return adata if copy else None