import logging as lg
import os
from curses import COLOR_MAGENTA
from multiprocessing import Pool
from typing import Any, Callable, Hashable, Iterable, Literal

import polars as pl
from rdkit.Chem import Mol

from fairfetched.utils.polars import (
    map_batches_unique,
    map_elements_cached,
    map_elements_pooled_cached,
    map_elements_unique,
    map_elements_unique_pooled,
)

from .compounds import (
    _safe_inchi_to_mol,
    _safe_smiles_to_mol,
    fuse_pipeline,
    inchi_to_inchikey,
    mol_to_strings,
    mols_to_inchikeys,
    mols_to_inchis,
    mols_to_inchis_and_auxinfo,
    mols_to_kekulised_smiles,
    standardise_aeffect,
    standardise_papyrus,
)

DFLike = pl.DataFrame | pl.LazyFrame


def with_standardised_mols(
    df,
    *standardisation_functions,
    start_col: str = "smiles",
    str_to_mol_function: Callable[[str], Mol | None] | None = None,
    mol_to_repr_function: Callable[[Mol | None], str | tuple[str, ...] | Any]
    | None = None,
    new_col_name: str = "standardised",
    final_dtype: pl.DataType | pl.DataTypeExpr | None = pl.Object,  # ty: ignore[invalid-parameter-default]
):
    """Example usage:
    ```
        from fairfetched.standardisation import compounds as ms
        from fairfetched.standardisation.pipe import smiles_inchi_auxinfo_inchikey_to_three_cols
        df.pipe(
            with_standardised_mols,
            ms.standardise_papyrus,
            start_col="inchi",
            new_col_name="smiles_inchi_auxinfo_inchikey",
            final_dtype=pl.Array(inner=pl.Utf8, shape=3),
            str_to_mol_function=ms._safe_inchi_to_mol,
            mol_to_repr_function=ms.mol_to_strings,
        ).pipe(smiles_inchi_auxinfo_inchikey_to_three_cols, )

        #is roughly equivalent to
        df.with_columns(
            smiles_inchi_auxinfo_inchikey=pl.col("inchi")
            .map_elements(ms._safe_inchi_to_mol, return_dtype=pl.Object)
            .map_elements(ms.standardise_papyrus, return_dtype=pl.Object)
            .map_elements(ms.mol_to_strings, return_dtype=pl.Array(inner=pl.Utf8, shape=3))
        )
        #but should be faster as it reduces multiple function call overhead
        # and evaluates only for unique molecules.
    ```
    """
    if str_to_mol_function:
        standardisation_functions = (str_to_mol_function, *standardisation_functions)
    if mol_to_repr_function:
        standardisation_functions = (*standardisation_functions, mol_to_repr_function)
    return df.pipe(
        map_elements_unique,
        fuse_pipeline(*standardisation_functions),  # ty: ignore[invalid-argument-type]
        col=start_col,
        alias=new_col_name,
        return_dtype=final_dtype,
    )


def map_pipeline(
    df: pl.LazyFrame,
    *functions: Callable[[Mol | Any], Mol | Any],
    col="smiles",
    alias="standardised_form",
    return_dtype: pl.DataType | pl.DataTypeExpr | None = pl.Object,  # ty: ignore[invalid-parameter-default]
    pooled=True,
) -> pl.LazyFrame:
    """element-wise map a pipeline of functions to unique elements. uses a join at the end to combine into the lazyframe"""
    if not isinstance(df.select(col).collect_schema().dtypes()[0], Hashable):
        return df.with_columns(
            pl.col(col)
            .map_elements(
                fuse_pipeline(*functions),
                return_dtype=return_dtype,
            )
            .alias(alias)
        )
    map_fx = map_elements_cached
    if pooled:
        map_fx = map_elements_pooled_cached
    return df.pipe(
        map_fx,
        fuse_pipeline(*functions),
        col=col,
        alias=alias,
        return_dtype=return_dtype,
    )


def map_pipeline_unique_join(
    df: pl.LazyFrame,
    *functions: Callable[[Mol | Any], Mol | Any],
    col="smiles",
    alias="standardised_form",
    return_dtype: pl.DataType | pl.DataTypeExpr | None = pl.Object,  # ty: ignore[invalid-parameter-default]
) -> pl.LazyFrame:
    """element-wise map a pipeline of functions to unique elements. uses a join at the end to combine into the lazyframe"""

    return df.pipe(
        map_elements_unique,
        fuse_pipeline(*functions),
        col=col,
        alias=alias,
        return_dtype=return_dtype,
    )


def map_pipeline_unique_join_pooled(
    df: pl.LazyFrame,
    *functions: Callable[[Mol | Any], Mol | Any],
    col="smiles",
    alias="standardised_form",
    return_dtype: pl.DataType | pl.DataTypeExpr | None = pl.Object,  # ty: ignore[invalid-parameter-default]
) -> pl.LazyFrame:
    """element-wise map a pipeline of functions to unique elements. uses a join at the end to combine into the lazyframe"""
    return df.pipe(
        map_elements_unique_pooled,
        fuse_pipeline(*functions),
        col=col,
        alias=alias,
        return_dtype=return_dtype,
    )


def polars_pipe_pipeline(df, *pipe_functions):
    for fnx in pipe_functions:
        df = df.pipe(fnx)
    return df


def unnest_molstrings(
    df: pl.LazyFrame,
    col: str = "smiles_inchi_auxinfo_inchikey",
    rename_old_suffix: str = "_original",
) -> pl.LazyFrame:
    return (
        df.rename(
            {
                c: f"{c}{rename_old_suffix}"
                for c in ["smiles", "inchi", "inchi_auxinfo", "inchikey"]
                if c in df.collect_schema().keys()
            },
            strict=False,
        )
        .with_columns(
            pl.col(col)
            .arr.to_struct(fields=["smiles", "inchi", "inchi_auxinfo", "inchikey"])
            .struct.unnest()
        )
        .select(
            pl.col(df.collect_schema().names()),
            pl.exclude(df.collect_schema().names()),
        )
        .drop(col)
    )


def with_mols_from_inchi_pool(df: DFLike, alias="mol", parallel=True) -> DFLike:
    """
    for ~1M mols: 1h07m for fully linear one, 15m for this impl.
    """
    if parallel:
        unique_inchis = df.lazy().select(pl.col("inchi")).unique().collect()["inchi"]  # ty: ignore
        with Pool() as p:
            res = p.map(_safe_inchi_to_mol, unique_inchis)
        mols_df = pl.DataFrame(
            [
                pl.Series(unique_inchis, dtype=pl.Utf8).alias("inchi"),
                pl.Series(res, dtype=pl.Object).alias(alias),
            ]
        )
        return df.lazy().join(
            mols_df.lazy(),
            on="inchi",
            how="left",
            maintain_order="left",
        )
    return df.lazy().join(
        df.lazy()
        .select(pl.col("inchi"))
        .unique()
        .with_columns(
            pl.col("inchi")
            .map_elements(_safe_inchi_to_mol, return_dtype=pl.Object)
            .alias(alias)
        ),
        on="inchi",
        how="left",
        maintain_order="left",
    )


def with_mols_from_inchi(df: pl.DataFrame, alias="mol", parallel=True) -> DFLike:
    """
    Parallel InChI-to-mol conversion using ThreadPoolExecutor (safe with JAX).
    Returns a new DataFrame with a column of RDKit Mol objects from InChI strings.
    """
    if parallel:
        from concurrent.futures import ThreadPoolExecutor

        unique_inchis = (
            df.select(pl.col("inchi"))
            .unique()
            .get_column("inchi")  # ignore ty[not-subscriptable]
        )
        with ThreadPoolExecutor(
            max_workers=round((os.cpu_count() or 1) * 0.9)
        ) as executor:
            res = list(executor.map(_safe_inchi_to_mol, unique_inchis))
        mols_df = pl.DataFrame(
            [
                pl.Series(unique_inchis, dtype=pl.Utf8).alias("inchi"),
                pl.Series(res, dtype=pl.Object).alias(alias),
            ]
        )
        return df.lazy().join(
            mols_df.lazy(),
            on="inchi",
            how="left",
            maintain_order="left",
        )
    return df.lazy().join(
        df.lazy()
        .select(pl.col("inchi"))
        .unique()
        .with_columns(
            pl.col("inchi")
            .map_elements(_safe_inchi_to_mol, return_dtype=pl.Object)
            .alias(alias)
        ),
        on="inchi",
        how="left",
        maintain_order="left",
    )


def with_clean_mols_from_smiles(df: DFLike, alias="mol") -> DFLike:
    return df.with_columns(
        pl.col("smiles")
        .map_elements(_safe_smiles_to_mol, return_dtype=pl.Object, strategy="threading")
        .alias("mol")
    )


def _clean_mols_og_smiles_df(smiles: Iterable[str]) -> pl.DataFrame:
    with Pool() as p:
        res = p.map(_safe_smiles_to_mol, smiles)
    return pl.DataFrame(
        [
            pl.Series(smiles, dtype=pl.Utf8).alias("smiles"),
            pl.Series(res, dtype=pl.Object).alias("mol"),
        ]
    )


def _series_mols_to_string_representation(
    mols: pl.Series,
    type: Literal["smiles", "inchi", "inchi_and_auxinfo", "inchikey"] = "smiles",
) -> pl.Series:
    match type:
        case "smiles":
            func = mols_to_kekulised_smiles
            dtype = pl.Utf8
        case "inchi":
            func = mols_to_inchis
            dtype = pl.Utf8
        case "inchi_and_auxinfo":
            func = mols_to_inchis_and_auxinfo
            dtype = pl.Struct({"inchi": pl.Utf8, "inchi_auxinfo": pl.Utf8})
        case "inchikey":
            func = mols_to_inchikeys
            dtype = pl.Utf8
        case _:
            raise NotImplementedError
    return pl.Series(name=type, values=func(mols), dtype=dtype)


def _unique_with_clean_string_representations(df: DFLike) -> DFLike:
    """cleans and standardises molecules, replacing original string
    representations (which are tagged _papyrus) and adding new ones
    without the suffix. This is to ensure that all string representations
    are consistent and can be used for downstream tasks without confusion."""
    df = df.pipe(with_clean_mols_from_smiles).rename(
        {k: f"{k}_papyrus" for k in ["smiles", "inchi", "inchi_auxinfo", "inchikey"]},
        strict=False,
    )
    mols = df.lazy().select("mol").collect().get_column("mol")
    return df.with_columns(
        _series_mols_to_string_representation(mols, type="smiles"),
        _series_mols_to_string_representation(
            mols, type="inchi_and_auxinfo"
        ).struct.unnest(),
    ).with_columns(
        pl.col("inchi")
        .map_elements(inchi_to_inchikey, return_dtype=pl.Utf8)
        .alias("inchikey")
    )


def with_clean_string_representations(df: DFLike) -> DFLike:
    """Variant of with_clean_string_representations that uses a join on unique mols
    to avoid rerunning similar mols. Throws away all invalid mols"""
    unique_smiles_df: pl.DataFrame = _clean_mols_og_smiles_df(
        df.select("smiles").unique().collect()["smiles"]  # ty: ignore
    ).drop_nulls("mol")

    original_col_order = list(df.collect_schema().keys())
    mols: pl.Series = unique_smiles_df["mol"]
    return (
        (
            unique_smiles_df.rename({"smiles": "smiles_"})
            .with_columns(
                # smiles
                _series_mols_to_string_representation(mols, type="smiles"),
                # inchis + auxinfo
                *(
                    _series_mols_to_string_representation(
                        mols, type="inchi_and_auxinfo"
                    )
                    .struct.unnest()
                    .iter_columns()
                ),
            )
            # inchikeys directly from the clean inchis
            .with_columns(
                pl.col("inchi")
                .map_elements(inchi_to_inchikey, return_dtype=pl.Utf8)
                .alias("inchikey")
            )
        )
        .lazy()
        .join(
            df.lazy(),
            left_on="smiles_",
            right_on="smiles",
            how="right",
            maintain_order="right",
            suffix="_old",
        )
        .select(*original_col_order, pl.exclude(original_col_order).exclude("mol"))
    )


def with_chembl_approval_year(
    df: pl.LazyFrame | pl.DataFrame,
    chembl_id_col: str = "chembl_id",
    max_calls: int = 10,
    sleep_time: float = 1.0,
) -> pl.LazyFrame | pl.DataFrame:
    """
    Adds a column 'approval_year' to the DataFrame by fetching approval years from ChEMBL API.
    Uses batching and rate limiting to avoid overloading the API.
    """
    # Extract unique, non-null ChEMBL IDs
    uniques = (
        df.lazy()
        .select(pl.col(chembl_id_col))
        .unique()
        .drop_nulls()
        .collect()[chembl_id_col]  # ty: ignore[not-subscriptable]
        .to_list()
    )

    results = chembls_to_approval_years(
        uniques, max_calls=max_calls, sleep_time=sleep_time
    )
    print(results)
    approval_df = (
        pl.DataFrame(
            results, schema={"chembl_id": pl.String, "approval_year": pl.Int32}
        )
        .rename({"chembl_id": chembl_id_col})
        .lazy()
    )

    return df.lazy().join(approval_df, on=chembl_id_col, how="left")


def with_cluster_idx(
    df: DFLike,
    descriptor_col: str = "morgan",
    method="bitbirch",
    alias: str | None = None,
    **kwargs,
):
    # @TODO: check if it is fine to take all unique descriptors or if it should still happen at the mol level
    descriptors = df.lazy().select(descriptor_col).unique().collect()[descriptor_col]  # ty: ignore
    clusters = cluster(
        descriptors=descriptors.to_numpy(),
        method=method,
        **kwargs,
    )
    res = pl.DataFrame(
        [
            descriptors.alias(descriptor_col),
            pl.Series(clusters.tolist(), dtype=pl.Int64).alias(
                alias if alias else method
            ),
        ]
    ).lazy()
    return df.lazy().join(res, on=descriptor_col, how="left", maintain_order="left")


def with_dimred(
    df: DFLike,
    descriptor_col: str = "morgan",
    method: Literal["umap", "tsne"] = "umap",
    metric: str = "auto",
    alias: str | None = None,
    **kwargs,
) -> DFLike:
    descriptors = df.lazy().select(descriptor_col).unique().collect()[descriptor_col]  # ty: ignore
    dimred = dimensionality_reduction(
        descriptors.to_numpy(),  # ty:ignore  [invalid-argument-type]
        method=method,
        metric=metric,
        **kwargs,
    )
    width = dimred.shape[1]
    res = pl.DataFrame(
        [
            descriptors.alias(descriptor_col),
            pl.Series(
                name=alias if alias else method,
                values=dimred.tolist(),
                dtype=pl.Array(pl.Float32, width),
            ),
        ]
    ).lazy()
    return df.lazy().join(res, on=descriptor_col, how="left", maintain_order="left")


def eager_generate_and_filter_by_protein_similarity(
    df: pl.LazyFrame | pl.DataFrame,
    query_df: pl.LazyFrame,
    method: Literal[
        "alphafold_cosine",
        "esm_cosine",
        "global_alignment_blosum62",
        "identity",
        "lehvenstein",
    ] = "esm_cosine",
    threshold: float = 0.7,
) -> pl.LazyFrame | pl.DataFrame:
    """will materialise both dataframes for protein similarities.
    keeps similarities for easy further filtering"""

    max_sims: pl.LazyFrame = (
        protein_similarity_matrix(df.lazy(), query_df.lazy(), method=method)
        .lazy()
        .select(
            pl.col("sequence"),
            pl.max_horizontal(pl.exclude("sequence")).alias("max"),
        )
    )
    return df.lazy().join(
        max_sims.filter(pl.col("max") >= threshold),
        on="sequence",
        how="semi",  # keep df if in max_sims["sequence"]
        maintain_order="left",
    )


def filter_protein_similarity(
    df: pl.LazyFrame,
    prot_with_avoidome_sims: pl.LazyFrame,
    avoidome_df: pl.LazyFrame | None,
    threshold: float = 0.7,
    subset: Literal[
        "all",
        "ion channels",
        "G-protein coupled receptors",
        "nuclear hormone receptors",
        "transporters",
        "enzymes",
        "kinases",
        "first",  # for testing mostly
    ] = "all",
    avoidome_sim_col="avoidome_esm_mean_cosine",
) -> pl.LazyFrame:
    if subset == "first":
        idx = [0]
    elif subset != "all":
        assert avoidome_df is not None, "subset != 'all': avoidome_df must be provided"
        select_ = "" if subset == "all" else subset
        idx = avoidome_df.collect().filter(  # ty: ignore[unresolved-attribute]
            pl.col("type").cast(pl.String).str.contains(select_)
        )["index"]
    else:
        idx = list(
            range(prot_with_avoidome_sims.collect_schema()[avoidome_sim_col].shape[0])  # ty:ignore
        )
    return df.join(
        prot_with_avoidome_sims.filter(
            pl.col(avoidome_sim_col)
            .cast(pl.List)  # workaround for until polars issue #26008 is solved
            .list.gather(idx)
            .list.max()
            >= threshold
        ),
        on="target_id",
        how="semi",  # no info from prot df, only selection
        maintain_order="left",
    )


def with_split_index(
    df,
    X_col: str = "inchikey",
    y_cols: str | list = "pchembl_value_mean",
    cluster_col: str = "cluster",
    method: Literal["tricario", "sklearn"] = "tricario",
    n_splits: int = 5,
    alias: str = "split",
    **kwargs,
):
    if isinstance(y_cols, str):
        y_cols = [y_cols]
    clusters, split_idxs = split(
        df=df,
        X_col=X_col,
        y_cols=y_cols,
        cluster_col=cluster_col,
        n_splits=n_splits,
        method=method,
        **kwargs,
    )
    assert len(clusters) == len(split_idxs)
    mapping = pl.DataFrame(
        [
            pl.Series(name=cluster_col, values=clusters),
            pl.Series(name=alias, values=split_idxs),
        ]
    ).unique()  # to handle the
    return df.lazy().join(
        mapping.lazy(),
        on=cluster_col,
        how="left",
    )


def __conditional_rename(
    df: pl.DataFrame | pl.LazyFrame, dict_: dict[str, str] | None
) -> pl.DataFrame | pl.LazyFrame:
    if not dict_:
        return df
    dict_ = {k: v for k, v in dict_.items() if k in df.collect_schema()}
    return df.rename(dict_)


def with_query_sim_array(
    df: pl.DataFrame | pl.LazyFrame,
    query_df: pl.DataFrame | pl.LazyFrame,
    method: str,
    alias: str = "protein_similarity",
    use_col_as: dict[str, str] | None = None,
) -> pl.LazyFrame:
    """compresses the similarity matrix df to pl.Array, keeping the order

    Args:
        df (pl.DataFrame | pl.LazyFrame): dataframe with 'sequence' column
        query_df (pl.DataFrame | pl.LazyFrame): dataframe with 'sequence' column
        method (str): method for protein similarity
        alias (str, optional): name for the similarity array column. Defaults to "protein_similarity".
        use_col_as (dict[str, str] | None, optional): mapping of column names to use as standard
            names. useful when using embeddings. Example could be `{"esm_max_pooling":"esm"}`.
            Will not raise on missing items and will try to apply on both. In case the columns
            do not exist and embeddings are needed, will just calculate the embeddings from scratch.
            Defaults to None.
    Returns:
        pl.LazyFrame: original df with added similarity array column

    """

    idx_sequence = query_df.select(["index", "sequence"]).sort("index").lazy().collect()

    matrix = protein_similarity_matrix(
        df.unique("sequence").pipe(__conditional_rename, use_col_as).lazy(),
        query_df.pipe(__conditional_rename, use_col_as).lazy(),
        method=method,
    )
    n = matrix.drop("sequence").width
    matrix = matrix.select(
        pl.col("sequence"),
        (
            pl.concat_list(pl.col(idx_sequence.get_column("sequence")))  # ty: ignore[unresolved-attribute]
            .list.to_array(n)
            .alias(alias)
        ),
    )
    return df.lazy().join(matrix.lazy(), on="sequence", how="left")


# ====== utils =======


def task_readout_per_activity(
    flat_df: pl.DataFrame | pl.LazyFrame,
    readout_fmt: str = "^type_.*$",
    group_by="activity_id",
) -> pl.DataFrame | pl.LazyFrame:
    """calculates the sum of readouts per activity. useful for stratification for splitting.
    Args:
        - flat_df (pl.DataFrame | pl.LazyFrame): dataframe with single datapoints
            that have one-hot-encodings for each activity column
        - readout_fmt (str): regex expression (or single string) for selecting the readout columns
        - group_by (str): column name for grouping by, e.g. `activity_id`,`target_id`,`smiles`

    Returns:
        - pl.LazyFrame with columns [group_by, readout_fmt] where readout_fmt
        is the sum of the readouts for each group_by unit.

    Example usage:
        `pl.scan_parquet('papyrus.parquet').pipe(semicolon_to_flat).pipe(task_readout_per_activity)`
    """
    return (
        flat_df.select(pl.col(readout_fmt), group_by)
        .group_by(group_by)
        .agg(pl.col(readout_fmt).cast(pl.Int8).sum())
    )


def task_protein_per_mol(
    flat_df, target_col="target_id", mol_col="inchi"
) -> pl.LazyFrame:
    return flat_df.select(target_col, mol_col).pivot(
        on=target_col,
        on_columns=flat_df.select(target_col).unique().collect().to_series(),
        index=mol_col,
        values=target_col,
        aggregate_function="len",
    )


def filter_is_in(lf, lf_series_filter):
    """join-based filter is_in for a lazyframe.
    Args:
        lf (pl.LazyFrame): lazyframe to filter
        lf_series_filter (pl.LazyFrame): lf for which values to keep. if not single-column, will use all columns. Usually only single-columned LF.

    Note: could be used with pl.DataFrames too but I generally avoid joins with eager df's due to memory concerns
    """
    return lf.lazy().join(
        lf_series_filter.unique().lazy(),
        on=lf_series_filter.collect_schema().names(),  # all cols
        how="semi",  # keep left if in right
        maintain_order="left",
    )


# ======== prep


def with_chemprop_datapoints_from_mols(
    df: pl.DataFrame,
    X_d: str | None,
    mol_col: str = "mol",
    y: str = "pchembl_value_mean",
) -> pl.DataFrame:
    assert X_d is None or X_d in df.collect_schema()
    assert y is None or y in df.collect_schema()
    df = df.drop("index", strict=False).with_row_index().sort("index")
    mols = df.lazy().select(mol_col).collect().get_column(mol_col).to_list()  # ty: ignore[unresolved-attribute]
    prot_descs = (
        df.lazy().select(X_d).collect().get_column(X_d).to_numpy() if X_d else None  # ty: ignore[unresolved-attribute]
    )
    ys = df.lazy().select(y).collect().get_column(y).to_list() if y else None  # ty: ignore[unresolved-attribute]
    return df.sort("index").with_columns(
        pl.Series(
            create_molecule_datapoints(mols, prot_descs, ys), dtype=pl.Object
        ).alias("chemprop_datapoints")
    )


def with_chemprop_datapoints_from_inchi(
    df: pl.DataFrame,
    X_d: str | None,
    y: str = "pchembl_value_mean",
    parallel_mols: bool = True,
) -> pl.DataFrame:
    assert X_d is None or X_d in df.collect_schema()
    assert y is None or y in df.collect_schema()
    df = df.pipe(with_mols_from_inchi, parallel=parallel_mols)  # ty:ignore[invalid-assignment]
    return with_chemprop_datapoints_from_mols(
        df, X_d, mol_col="mol", y="pchembl_value_mean"
    ).drop("mol")


def with_chemprop_avoidome_datapoints(
    df: pl.DataFrame,
    avoidome: pl.DataFrame,
    prot_desc: str,
    parallel_mols: bool = False,
) -> pl.DataFrame:
    assert prot_desc in avoidome.collect_schema()
    df = (  # ty: ignore[invalid-assignment]
        df.join(avoidome.select("accession", "sequence", prot_desc), how="cross")
        .with_row_index()
        .pipe(with_mols_from_inchi, parallel=parallel_mols)
        .sort("index")
        .collect()  # ty: ignore[ unresolved-attribute]
    )
    mols = df["mol"]
    prot_descs = df[prot_desc].to_numpy() if prot_desc else None
    return (
        df.sort("index")
        .with_columns(
            pl.Series(
                create_molecule_datapoints(mols, prot_descs, ys=None), dtype=pl.Object
            ).alias("chemprop_datapoints")
        )
        .drop("mol")
    )
