import warnings

# from sklearn.preprocessing import LabelEncoder

from typing import Tuple, Dict, Optional, Union

import anndata

import numpy as np

import pandas as pd

import scanpy as sc

from scipy import sparse

from sklearn.utils import resample

from anndata._core.anndata import AnnData

from numpy import ndarray

from scipy.sparse._base import spmatrix

from pandas.core.arrays.categorical import Categorical

from pandas.core.frame import DataFrame

from scipy.sparse._csr import csr_matrix

warnings.filterwarnings("ignore")

def preprocess(

    adata: object,

    integrated: bool=False,

    ct_col: Optional[str]=None,

    y_col: Optional[str]=None,

    pt_col: Optional[str]=None,

    class_map: Optional[Dict[str, int]]=None,

) -> AnnData:

    """standardize input data

    Parameters

    ----------

    adata : object

    integrated : bool=False

    ct_col : Optional[str]=None

    y_col : Optional[str]=None

    pt_col : Optional[str]=None

    class_map : Optional[Dict[str, int]]=None

    Returns

    -------

    AnnData

    """

    try:

        adata.__dict__["_raw"].__dict__["_var"] = (

            adata.__dict__["_raw"]

            .__dict__["_var"]

            .rename(columns={"_index": "features"})

        )

    except AttributeError as e:

        # print(f"An AttributeError occurred: {e}")

        # Handle the error or do nothing if you just want to continue

        pass

    # adata.raw = adata

    # can be recoverd by: adata_raw = adata.raw.to_adata()

    if not integrated:

        # Standardize X matrix-----------------------------------------

        ## 1. Check if X is raw count matrix

        ### If True, make X into logcount matrix

        if check_nonnegative_integers(adata.X):

            adata.layers["count"] = adata.X

            print("X is raw count matrix, adding to `count` layer......")

            sc.pp.normalize_total(adata, target_sum=1e4)

            sc.pp.log1p(adata)

        elif check_nonnegative_float(adata.X):

            adata.layers["logcount"] = adata.X

            print("X is logcount matrix, adding to `logcount` layer......")

        else:

            raise TypeError("X should be either raw counts or log-normalized counts")

    # Filter low-express genes-----------------------------------------------

    sc.pp.filter_genes(adata, min_cells=1)

    # Calculate HVGs-------------------------------------------------------

    # sc.pp.highly_variable_genes(adata, min_mean=0.0125, max_mean=3, min_disp=0)

    # Standardize Metadata-----------------------------------------

    ## cell type as ct

    ## coondition as y

    if ct_col is not None:

        adata.obs["ct"] = adata.obs[ct_col]

        adata.obs["ct"] = adata.obs["ct"].astype("category")

    else:

        adata.obs["ct"] = adata.obs["ct"].astype("category")

    if y_col is not None:

        adata.obs["y"] = adata.obs[y_col]

        adata.obs["y"] = adata.obs["y"].astype("category")

    else:

        adata.obs["y"] = adata.obs["y"].astype("category")

    if pt_col is not None:

        adata.obs["patient_id"] = adata.obs[pt_col]

        adata.obs["patient_id"] = adata.obs["patient_id"].astype("category")

    else:

        adata.obs["patient_id"] = adata.obs["patient_id"].astype("category")

    # print(class_map)

    adata.obs["label"] = adata.obs["y"].map(class_map)

    #     # add HVGs in each cell type

    #     # split each cell type into new objects.

    #     celltypes_all = adata.obs['ct'].cat.categories.tolist()

    #     alldata = {}

    #     for celltype in celltypes_all:

    #         #print(celltype)

    #         adata_tmp = adata[adata.obs['ct'] == celltype,]

    #         if adata_tmp.shape[0] > 20:

    #             sc.pp.highly_variable_genes(adata_tmp, min_mean=0.0125, max_mean=3, min_disp=1.5)

    #             adata_tmp = adata_tmp[:, adata_tmp.var.highly_variable]

    #             alldata[celltype] = adata_tmp

    #     celltype_hvgs = anndata.concat(alldata, join="outer").var_names.tolist()

    #     hvgs = adata.var.highly_variable.index[adata.var.highly_variable].tolist()

    #     hvgs_idx = adata.var_names.isin(list(set(celltype_hvgs + hvgs)))

    return adata

def get_X_y_from_ann(adata: AnnData, return_adj: bool=False, n_neigh: int=10) -> Union[Tuple[ndarray, Categorical], Tuple[ndarray, Categorical, csr_matrix]]:

    """

    Parameters

    ----------

    adata

    return_adj

    n_neigh

    Returns

    -------

    """

    # scale and store results in layer

    # adata.layers['scaled'] = sc.pp.scale(adata, max_value=10, copy=True).X

    # This X matrix should be scaled value

    X = adata.to_df().values

    y = adata.obs["label"].values

    if return_adj:

        if adata.n_vars >= 50:

            sc.pp.pca(adata, n_comps=30, use_highly_variable=False)

            sc.pl.pca_variance_ratio(adata, log=True, n_pcs=30)

            knn = sc.Neighbors(adata)

            knn.compute_neighbors(n_neighbors=n_neigh, n_pcs=30)

        else:

            print(

                "Number of input genes < 50, use original space to get adjacency matrix..."

            )

            knn = sc.Neighbors(adata)

            knn.compute_neighbors(n_neighbors=n_neigh, n_pcs=0)

        adj = knn.connectivities + sparse.diags([1] * adata.shape[0]).tocsr()

        return X, y, adj

    else:

        return X, y

def check_nonnegative_integers(X: Union[np.ndarray, sparse.spmatrix]) -> bool:

    """Checks values of X to ensure it is count data"""

    from numbers import Integral

    data = X[np.random.choice(X.shape[0], 10, replace=False), :]

    data = data.todense() if sparse.issparse(data) else data

    # Check no negatives

    if np.signbit(data).any():

        return False

    # Check all are integers

    elif issubclass(data.dtype.type, Integral):

        return True

    elif np.any(~np.equal(np.mod(data, 1), 0)):

        return False

    else:

        return True

def check_nonnegative_float(X: Union[np.ndarray, sparse.spmatrix]) -> bool:

    """Checks values of X to ensure it is logcount data"""

    from numbers import Integral

    data = X[np.random.choice(X.shape[0], 10, replace=False), :]

    data = data.todense() if sparse.issparse(data) else data

    # Check no negatives

    if np.signbit(data).any():  # one negative(if True), return False

        return False  # non negative(if False), go to next condition

    # Check all are integers

    elif np.issubdtype(data.dtype, np.floating):

        return True

def compute_cell_weight(data: Union[AnnData, DataFrame]) -> ndarray:

    ## data: anndata object | DataFrame from anndata.obs

    if isinstance(data, anndata.AnnData):

        data = data.obs

    from sklearn.utils.class_weight import compute_sample_weight

    ## class weight

    w_c = compute_sample_weight(class_weight="balanced", y=data["label"])

    w_c_df = pd.DataFrame({"patient_id": data["patient_id"], "w_c": w_c})

    ## patient weight

    gped = data.groupby(["label"])

    w_pt_df = pd.DataFrame([])

    for name, group in gped:

        w_pt = compute_sample_weight(class_weight="balanced", y=group["patient_id"])

        w_pt_df = pd.concat(

            [

                w_pt_df,

                pd.DataFrame(

                    {

                        "label": group["label"],

                        "patient_id": group["patient_id"],

                        "w_pt": w_pt,

                    }

                ),

            ],

            axis=0,

        )

    ## cell_weight = class weight*patient_weight

    w_df = w_c_df.merge(w_pt_df, left_index=True, right_index=True)

    w_df["w"] = w_df["w_c"] * w_df["w_pt"]

    ## Add results to adata metadata or save csv???

    # adata_train.obs = adata_train.obs.merge(w_df[['w_pt','w_c','w']], left_index=True, right_index=True)

    return w_df["w"].values

def downsample_adata(adata, downsample_size=4000, random_state=1):

    #######################

    ## stratified downsampling by patient_id, y, and ct (patient, disease status, and cell type)

    # metadata for stratified downsampling

    adata.obs["downsample_stratify"] = adata.obs[["patient_id", "y", "ct"]].apply(

        lambda x: "_".join(x.dropna().astype(str)), axis=1

    )

    adata.obs["downsample_stratify"] = adata.obs["downsample_stratify"].astype(

        "category"

    )

    down_index = resample(

        adata.obs_names,

        replace=False,

        n_samples=downsample_size,

        stratify=adata.obs["downsample_stratify"],

        random_state=random_state,

    )

    adata = adata[down_index,]

    return adata