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 fot_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_fot_signaling(self,
        fot_eps_p=1e-1, 
        fot_eps_mu=None, 
        fot_eps_nu=None, 
        fot_rho=1e1, 
        fot_nitermax=1e4, 
        fot_weights=(0.25,0.25,0.25,0.25)
    ):
        self.comm_network = fot_combine_sparse(self.S, self.mets, self.D, self.A, self.M, self.LRC, self.LRC_type, self.cutoff, \
            eps_p=fot_eps_p, eps_mu=fot_eps_mu, eps_nu=fot_eps_nu, rho=fot_rho, weights=fot_weights, nitermax=fot_nitermax)

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',
    fot_eps_p: float = 1e-1, 
    fot_eps_mu: Optional[float] = None, 
    fot_eps_nu: Optional[float] = None, 
    fot_rho: float =1e1, 
    fot_nitermax: int = 10000, 
    fot_weights: tuple = (0.25,0.25,0.25,0.25),
    copy: bool = False
):
    """
    Function for inferring spatial metabolic cell communication.

    Parameters
    ----------
    adata
        The data matrix of shape ``n_obs`` × ``n_var``.
        Rows correspond to cells or spots and columns to genes.
    database_name
        Name of the Metabolite-Sensor interaction database. Will be included in the keywords for anndata slots.
    df_metasen
        A data frame where each row corresponds to a metabolite-sensor pair with Metabolite, Sensor, Metabolite.Pathway, Sensor.Pathway, Metabolite.Names, Long.Range.Channel, respectively.
    LRC_type
        The name of all possible long-range channel, provided as a `list`, such as ["Blood"] or ["CSF"] or ["Blood","CSF"].
    dis_thr
        The farthest distance from a nearby cell to the LRC indicates the range of cells in which long-range communication can occur, provided as a `float`.
    cost_scale
        Weight coefficients of the cost matrix for each metabolite-sensor pair, e.g. cost_scale[('metA','senA')] specifies weight for the pair metA and senA.
        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.
    fot_eps_p
        The coefficient of entropy regularization for transport plan.
    fot_eps_mu
        The coefficient of entropy regularization for untransported source (metabolite). Set to equal to fot_eps_p for fast algorithm.
    fot_eps_nu
        The coefficient of entropy regularization for unfulfilled target (sensor). Set to equal to fot_eps_p for fast algorithm.
    fot_rho
        The coefficient of penalty for unmatched mass.
    fot_nitermax
        Maximum iteration for flow optimal transport algorithm.
    fot_weights
        A tuple of four weights that add up to one. The weights corresponds to four setups of flow optimal transport: 
        1) all metabolite-all sensors, 2) each metabolite-all sensors, 3) all metabolite-each sensor, 4) each metabolite-each sensor.
    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 MS interaction database named "databaseX", 
        ``.obsp['metachat-databaseX-metA-senA']``
        is a ``n_obs`` × ``n_obs`` matrix with the *ij* th entry being the "score" of 
        cell *i* sending signal to cell *j* through metA and senA.
        The marginal sums (sender and receiver) of the signaling matrices are stored in ``.obsm['metachat-databaseX-sum-sender']`` and ``.obsm['metachat-databaseX-sum-receiver']``.
        Metadata of the analysis is added to ``.uns['metachat-databaseX-info']``.
        If copy=True, return the AnnData object and return None otherwise.
        
    """
    # Check inputs
    assert database_name is not None, "Please give a database_name"
    assert df_metasen is not None, "Please give a Metabolite-Sensor database"
    if LRC_type is None: 
        print("You didn't input LRC_type, so long-range communication will not be consider in inference")
    assert dis_thr is not None, "Please give a dis_thr"

    # 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_fot_signaling(fot_eps_p=fot_eps_p, 
        fot_eps_mu = fot_eps_mu, 
        fot_eps_nu = fot_eps_nu, 
        fot_rho = fot_rho, 
        fot_nitermax = fot_nitermax, 
        fot_weights = fot_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