import os
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
from functools import partial
from multiprocessing import Pool
from typing import Callable, Iterable, Literal

import polars as pl
from rdkit.Chem import Mol

from fairfetched.utils._track import track
from fairfetched.utils.polars import apply_to_unique, map_batches_pooled

from .compound_fns import (
    _remove_stereo_papyrus_standardise_check_inchi,
    _safe_inchi_to_mol,
    inchi_to_inchikey,
    mols_to_inchikeys,
    mols_to_inchis,
    mols_to_inchis_and_auxinfo,
    mols_to_kekulised_smiles,
    standardised_nostereo_to_smiles_inchi_aux_inchikey,
)


def apply_to_unique_strings(
    lf: pl.LazyFrame,
    function: Callable[[str], Mol | None] = _safe_inchi_to_mol,
    from_col: str = "inchi",
    to_col: str = "mol",
    parallel=True,
    chunks_per_worker: int = 8,
) -> pl.LazyFrame:
    """
    for ~1M mols: 1h07m for fully linear one, 15m for this impl. (now actually <1m)
    chunksize is calculated as max(1, len(unique_molstrings) // (workers * chunks_per_worker))
    so for smaller chunks (good for larger operations), you would need more chunks per worker (800 or so)
    """
    if parallel:
        unique_molstrings = (
            lf.select(pl.col(from_col)).unique().collect().get_column(from_col)  # ty: ignore[unresolved-attribute]
        )
        workers = round(os.cpu_count() or 1 * 0.8)
        chunksize = max(1, len(unique_molstrings) // (workers * chunks_per_worker))
        with ProcessPoolExecutor(max_workers=workers) as p:
            res = track(
                p.map(function, unique_molstrings, chunksize=chunksize),
                total=len(unique_molstrings),
            )

            mols_df = pl.DataFrame(
                [
                    pl.Series(unique_molstrings, dtype=pl.Utf8).alias(from_col),
                    pl.Series(list(res), dtype=pl.Object).alias(to_col),
                ]
            )
        return lf.join(
            mols_df.lazy(),
            on=from_col,
            how="left",
            maintain_order="left",
        )

    return lf.join(
        lf.select(pl.col(from_col))
        .unique()
        .with_columns(
            pl.col(from_col)
            .map_elements(function, return_dtype=pl.Object)
            .alias(to_col)
        ),
        on=from_col,
        how="left",
        maintain_order="left",
    )


def with_cleaned_mol_descriptors_struct(
    lf: pl.LazyFrame,
    descriptor_to_descriptors_func: Callable[[str], tuple[str | None, ...]],
    from_col: str,
    to_col: str,
    return_dtype: pl.DataTypeExpr,
    parallel=True,
    **kwargs,
) -> pl.LazyFrame:
    return lf.pipe(
        apply_to_unique,  # creates unique list, iterates over pooled,
        descriptor_to_descriptors_func,
        from_col=from_col,
        to_col=to_col,
        return_dtype=return_dtype,
        parallel=True,
    )


def with_cleaned_mol_descriptors(
    lf: pl.LazyFrame,
    smiles_col: str = "smiles",
    **kwargs,
) -> pl.LazyFrame:
    """cleans and standardises mols by:
    1) removing stereochemistry
    2) running papyrus standardisation
    3) validating inchis
    returns None if invalid, but does not drop
    """
    fields = ["smiles", "inchi", "inchi_auxinfo", "inchikey"]
    to_col = "descriptors_struct"
    lf = lf.pipe(
        with_cleaned_mol_descriptors_struct,
        descriptor_to_descriptors_func=standardised_nostereo_to_smiles_inchi_aux_inchikey,
        from_col=smiles_col,
        to_col=to_col,
        return_dtype=pl.Struct(
            {k: pl.Utf8 for k in ["smiles", "inchi", "inchi_auxinfo", "inchikey"]}
        ),
        **kwargs,
    )
    return (
        lf.cast({to_col: pl.Struct({k: pl.Utf8 for k in fields})})
        .rename({k: f"{k}_original" for k in fields if k in lf.collect_schema()})
        .with_columns(pl.col(to_col).struct.unnest())
        .drop(to_col)
    )


def with_mols_from_inchi_batched(lf, alias="mol", parallel=False):
    function = partial(
        map_batches_pooled,
        fn=_safe_inchi_to_mol,
        return_dtype=pl.Object,
        parallel=parallel,
    )
    return lf.with_columns(
        pl.col("inchi")
        .map_batches(function, return_dtype=pl.Object, is_elementwise=True)
        .alias(alias)
    )


def with_mols_from_inchi_pool(
    lf: pl.LazyFrame, alias="mol", parallel=True
) -> pl.LazyFrame:
    """
    for ~1M mols: 1h07m for fully linear one, 15m for this impl.
    """
    if parallel:
        unique_inchis = (
            lf.select(pl.col("inchi")).unique().collect().get_column("inchi")  # ty: ignore[unresolved-attribute]
        )
        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 lf.join(
            mols_df.lazy(),
            on="inchi",
            how="left",
            maintain_order="left",
        )
    return lf.join(
        lf.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_thread(
    lf: pl.LazyFrame,
    alias="mol",
    parallel=True,
    join=True,
) -> pl.LazyFrame:
    """
    Returns a new DataFrame with a column of RDKit Mol objects from InChI strings.
    """
    if parallel:
        if join:
            inchis = (
                lf.select(pl.col("inchi")).unique().collect().get_column("inchi")  # ty: ignore[unresolved-attribute]
            )
        else:
            lf = lf.sort("inchi")
            inchis = lf.select("inchi").collect().get_column("inchi")
        workers = round((os.cpu_count() or 1) * 0.9)
        chunksize = max(1, len(inchis) // (workers * 4))
        print("n_workers=", workers)
        print("chunksize=", chunksize)
        with ProcessPoolExecutor(max_workers=workers) as executor:
            res = list(
                executor.map(
                    _safe_inchi_to_mol,
                    track(inchis, total=len(inchis), desc="loading mols"),
                    chunksize=chunksize,
                )
            )
            mol_series = pl.Series(res, dtype=pl.Object).alias(alias)
        if join:
            mols_df = pl.DataFrame(
                [pl.Series(inchis, dtype=pl.Utf8).alias("inchi"), mol_series]
            )
            return lf.join(
                mols_df.lazy(),
                on="inchi",
                how="left",
                maintain_order="left",
            )
        else:
            lf.with_columns(mol_series)
    return lf.join(
        lf.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: pl.LazyFrame, alias="mol") -> pl.LazyFrame:
    return df.with_columns(
        pl.col("smiles")
        .map_elements(
            _remove_stereo_papyrus_standardise_check_inchi,
            return_dtype=pl.Object,
            strategy="threading",
        )
        .alias("mol")
    )


def _smiles_and_cleaned_mols_df(smiles: Iterable[str]) -> pl.DataFrame:
    with Pool() as p:
        res = p.map(_remove_stereo_papyrus_standardise_check_inchi, 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: pl.LazyFrame) -> pl.LazyFrame:
#     """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: pl.LazyFrame) -> pl.LazyFrame:
    """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 = _smiles_and_cleaned_mols_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(
                # sm iles
                _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 __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_)


# ====== 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",
    )
