from __future__ import annotations

# import re
import urllib

import anndata
import numpy as np
import pandas as pd
import protdata
import scipy.sparse

# def read_alphastats(
#     loader: alphastats.BaseLoader,
#     x_dtype: Union[np.dtype, type, int, float, None] = None,
#     proteins_as_obs: bool = False,
# ) -> anndata.AnnData:
#     """Read proteomics data into an AnnData object.

#     Parameters
#     ----------
#     loader
#         A loader object from alphastats that contains the raw proteomics data and metadata
#     x_dtype
#         Data type to use for the intensity matrix, by default None
#     proteins_as_obs
#         If True, proteins will be stored in obs (rows) rather than var (columns), by default False


#     Notes
#     -----
#     The loader object must contain:
#     - rawinput: DataFrame with protein data
#     - software: Name of proteomics software used
#     - index_column: Column name containing protein identifiers
#     - intensity_column: Column name pattern for intensity values
#     - filter_columns: Columns used for filtering
#     - gene_names: Gene name mapping information
#     """

#     try:
#         import alphastats
#     except ImportError:
#         raise Exception(
#             "To read alphastats, please install the `alphastats` python package (pip install alphastats)."
#         )

#     alphastats.DataSet._check_loader(
#         1, loader
#     )  # need to put ugly 1 because this is not specified as a staticmethod in alphapeptstats
#     rawinput = loader.rawinput
#     software = loader.software
#     index_column = loader.index_column
#     intensity_column = loader.intensity_column
#     intensity_regex = re.compile(intensity_column.replace("[sample]", ".*"))
#     filter_columns = loader.filter_columns
#     # evidence_df = loader.evidence_df
#     gene_names = loader.gene_names

#     df = rawinput.copy()

#     # get the intensity columns
#     if isinstance(intensity_column, str):
#         intensity_regex = re.compile(intensity_column.replace("[sample]", ".*"))
#         intensity_col_mask = df.columns.map(lambda x: intensity_regex.search(x) is not None)
#     else:
#         intensity_col_mask = df.columns.isin(intensity_column)

#     # Convert to anndata object
#     var = df.loc[:, ~intensity_col_mask]
#     X = df.loc[:, intensity_col_mask]
#     X = X.replace(np.nan, 0)
#     obs = pd.DataFrame(index=X.columns)
#     var.set_index(index_column, inplace=True)
#     var.index = var.index.astype(str)
#     adata = anndata.AnnData(X=X.to_numpy(dtype=x_dtype).T, var=var, obs=obs)
#     adata.obs["Intensity_col"] = adata.obs.index
#     sample_regex = re.compile(intensity_column.replace("[sample]", ""))
#     adata.obs["Sample_name"] = adata.obs.index.str.replace(sample_regex, "", regex=True)
#     adata.obs.set_index(keys="Sample_name", drop=False, inplace=True)
#     obs.index = obs.index.astype(str)

#     # Proteins could either be in the rows or columns
#     if proteins_as_obs:
#         adata = adata.T

#     # Add properties of the experiment to uns
#     adata.uns["RawInfo"] = {
#         "software": software,
#         "filter_columns": filter_columns,
#         "gene_names": gene_names,
#     }
#     adata.uns["proteins_as_obs"] = proteins_as_obs
#     return adata


def _preprocess_adata(adata: anndata.AnnData) -> anndata.AnnData:
    """Preprocess an AnnData object."""

    # Replace NaNs with 0 in .X
    if isinstance(adata.X, np.ndarray):
        adata.X = np.nan_to_num(adata.X, nan=0)
    elif isinstance(adata.X, scipy.sparse.spmatrix):
        adata.X.data = np.nan_to_num(adata.X.data, nan=0, copy=False)

    # Replace NaNs with 0 in all layers
    for layer in list(adata.layers.keys()):
        arr = adata.layers[layer]
        if isinstance(arr, np.ndarray):
            adata.layers[layer] = np.nan_to_num(arr, nan=0, copy=False)
        elif isinstance(arr, scipy.sparse.spmatrix):
            adata.layers[layer].data = np.nan_to_num(arr.data, nan=0, copy=False)

    return adata


def read_prolocdata(file_name: str, allow_nullable_strings: bool = False) -> anndata.AnnData:
    """Read a prolocdata file and return an AnnData object.

    Parameters
    ----------
    file_name : str
        The path to the prolocdata file or a URL.
    allow_nullable_strings : bool, default False
        If False, convert pandas nullable StringDtype columns in obs/var to
        regular Python object-dtype strings for compatibility with older
        anndata writers (anndata<0.11). If True, keep nullable string dtype.

    Returns
    -------
    adata : AnnData
    """
    try:
        import rdata
    except ImportError:
        raise Exception(
            "To read prolocdata, please install the `rdata` python package (pip install rdata)."
        )

    parsed_url = urllib.parse.urlparse(file_name)
    if parsed_url.scheme != "":
        with urllib.request.urlopen(file_name) as dataset:
            pdata = rdata.parser.parse_data(dataset.read(), extension="rda")
    else:
        pdata = rdata.parser.parse_file(file_name)
    proloc_classes = {
        "AnnotatedDataFrame": lambda x, y: x,
        "Versions": lambda x, y: x,
        "MSnProcess": lambda x, y: x,
        "MSnSet": lambda x, y: x,
        "MIAPE": lambda x, y: x,
    }
    pdata = rdata.conversion.convert(
        pdata, constructor_dict={**proloc_classes, **rdata.conversion.DEFAULT_CLASS_MAP}
    )
    # Handle both cases: pdata is a dict (container) or already the dataset
    if isinstance(pdata, dict):
        dataset_name = next(iter(pdata.keys()))  # Reads the first dataset in the dictionary
        pdata = pdata[dataset_name]
    else:
        dataset_name = file_name
    # Create AnnData object with robust dtype handling
    obs_raw = pdata.featureData.data
    var_raw = pdata.phenoData.data
    # Ensure pandas DataFrames
    obs = (obs_raw if isinstance(obs_raw, pd.DataFrame) else pd.DataFrame(obs_raw)).copy()
    var = (var_raw if isinstance(var_raw, pd.DataFrame) else pd.DataFrame(var_raw)).copy()

    # Normalize column/index names to plain Python strings (avoid numpy.str_)
    obs.columns = obs.columns.map(str)
    obs.index = obs.index.map(str)
    var.columns = var.columns.map(str)
    var.index = var.index.map(str)

    # Infer better dtypes; control whether strings become pandas StringDtype
    # obs = obs.convert_dtypes(convert_string=allow_nullable_strings)
    # var = var.convert_dtypes(convert_string=allow_nullable_strings)

    if not allow_nullable_strings:
        for df in (obs, var):
            for c in df.columns:
                dt = df[c].dtype
                if pd.api.types.is_extension_array_dtype(dt) and pd.api.types.is_string_dtype(
                    dt
                ):
                    df[c] = df[c].astype(object)

    # Expression matrix
    X = np.asarray(pdata.assayData.maps[0]["exprs"], dtype=float)

    # Construct AnnData (expects shape: (n_obs, n_vars) == X.shape)
    adata = anndata.AnnData(obs=obs, var=var, X=X)

    # Add metadata
    adata.uns["dataset_name"] = dataset_name
    adata.uns["file_name"] = file_name
    metadata = {
        k: v
        for k, v in vars(pdata.experimentData).items()
        if hasattr(v, "__len__") and len(v) > 0
    }
    # Remove class version key if present
    metadata.pop(".__classVersion__", None)
    adata.uns["MIAPE_metadata"] = metadata
    _preprocess_adata(adata)

    return adata


def read_maxquant(*args, **kwargs) -> anndata.AnnData:
    """Read MaxQuant proteinGroups.txt file into an AnnData object.

    This function serves as a wrapper around :func:`~protdata.io.read_maxquant`,
    automatically transposing the data so that proteins are stored in ``.obs``
    (rows) and samples in ``.var`` (columns), following grassp conventions.

    Parameters
    ----------
    *args
        Positional arguments passed to :func:`protdata.io.read_maxquant`.
    **kwargs
        Keyword arguments passed to :func:`protdata.io.read_maxquant`.

    Returns
    -------
    AnnData
        Annotated data object with proteins as observations (rows) and samples
        as variables (columns).

    See Also
    --------
    protdata.io.read_maxquant : Original protdata function for reading MaxQuant files.

    Examples
    --------
    >>> import grassp as gr  # doctest: +SKIP
    >>> adata = gr.io.read_maxquant('proteinGroups.txt')  # doctest: +SKIP
    """
    adata = protdata.io.read_maxquant(*args, **kwargs)
    _preprocess_adata(adata)
    return adata.T


def read_fragpipe(*args, **kwargs) -> anndata.AnnData:
    """Read FragPipe combined_protein.tsv file into an AnnData object.

    This function serves as a wrapper around :func:`~protdata.io.read_fragpipe`,
    automatically transposing the data so that proteins are stored in ``.obs``
    (rows) and samples in ``.var`` (columns), following grassp conventions.

    Parameters
    ----------
    *args
        Positional arguments passed to :func:`protdata.io.read_fragpipe`.
    **kwargs
        Keyword arguments passed to :func:`protdata.io.read_fragpipe`.

    Returns
    -------
    AnnData
        Annotated data object with proteins as observations (rows) and samples
        as variables (columns).

    See Also
    --------
    protdata.io.read_fragpipe : Original protdata function for reading FragPipe files.

    Examples
    --------
    >>> import grassp as gr  # doctest: +SKIP
    >>> adata = gr.io.read_fragpipe('combined_protein.tsv')  # doctest: +SKIP
    """
    adata = protdata.io.read_fragpipe(*args, **kwargs)
    _preprocess_adata(adata)
    return adata.T


def read_diann(*args, **kwargs) -> anndata.AnnData:
    """Read DIA-NN report.tsv file into an AnnData object.

    This function serves as a wrapper around :func:`~protdata.io.read_diann`,
    automatically transposing the data so that proteins are stored in ``.obs``
    (rows) and samples in ``.var`` (columns), following grassp conventions.

    Parameters
    ----------
    *args
        Positional arguments passed to :func:`protdata.io.read_diann`.
    **kwargs
        Keyword arguments passed to :func:`protdata.io.read_diann`.

    Returns
    -------
    AnnData
        Annotated data object with proteins as observations (rows) and samples
        as variables (columns).

    See Also
    --------
    protdata.io.read_diann : Original protdata function for reading DIA-NN files.

    Examples
    --------
    >>> import grassp as gr  # doctest: +SKIP
    >>> adata = gr.io.read_diann('report.tsv')  # doctest: +SKIP
    """
    adata = protdata.io.read_diann(*args, **kwargs)
    _preprocess_adata(adata)
    return adata.T
