# %%
from fairfetched.get import Chembl, Papyrus

# Chembl and Papyrus are 'database' instance objects
# Once instantiated, they are defined by their names (chembl, papyrus), their versions,
# and their location on your device.

# without instantiation, you can check available versions for either:
print(Chembl.get_available_versions())
Chembl.from_latest()
exit()
# Files are only downloaded if they do not exist on their default location yet
# similar to pystow, it chooses the default location based on an environment variable
# in case it is set, with the following hierarchy: FAIRFETCHED_HOME, PYSTOW_HOME, ~/.data
# it then creates a folder with the database name
# <default-root-directory>/papyrus.
# you can overwrite this with the kwarg "root_dir", which has default value f"{BASE_DIR}/papyrus"
# then, it creates the version-specific directory to ensure versions will never
# be overwritten by another
# e.g., <default-root-directory>/papyrus/05.7
# if you passed root_dir="data", it would save to data/05.7
# WARNING: this might still be subject to change in pre-release versions (pre-0.1.0)
p = Papyrus.from_latest()
# once it is loaded,

# %%
# import fairfetched.standardisation.mols as sm

# from fairfetched.get import Papyrus
# from fairfetched.standardisation import Preset

# ds = (
#     Papyrus.from_latest()
#     .to_dataset()
#     .standardise(Preset.CHEMBL)
#     .standardise(sm.remove_stereochemistry)
#     .drop_invalid(report=True)
# )

# ds.compounds  # → LazyFrame with pipeline applied
# ds.bioactivity  # → LazyFrame with pipeline applied
# ds.describe()  # → full audit trail

if __name__ == "__main__":
    import polars as pl

    # p = Papyrus.from_latest()
    # lf = p.full_data.head(100)
    # df = lf.pipe(
    #     apply_to_unique_strings,
    #     from_col="smiles",
    #     str_to_obj_fn=_safe_smiles_to_mol,
    #     # from_col="smiles",
    #     # str_to_obj_fn=_safe_smiles_to_mol,
    #     alias="mol",
    #     parallel=True,
    #     chunks_per_worker=8,
    # ).collect()
    # print("done with collect")
    # print(df.collect_schema())
    from fairfetched.get import Papyrus
    from fairfetched.standardization.compound_fns import (
        _safe_inchi_to_mol,
        _safe_smiles_to_mol,
    )
    from fairfetched.standardization.pipe import apply_to_unique_strings

    p = Papyrus.from_version("05.7")
    from fairfetched.standardization.pipe import with_cleaned_mol_descriptors

    p.full_data.head(100).pipe(with_cleaned_mol_descriptors).collect()
    # lf = pl.LazyFrame(
    #     {
    #         "smiles": ["CCO", "CCO", "c1ccccc1", "CC(=O)O", "c1ccccc1", "invalid!!"],
    #         "id": [1, 2, 3, 4, 5, 6],
    #     }
    # )

    # import fairfetched.standardization.compound_ as ms
    # import papyrus_structure_pipeline.standardizer  # ty:ignore[unresolved-import] #ty:ignore[unused-ignore-comment] #ty:ignore[unused-ignore-comment]

    # from fairfetched.standardization import standardization_function

    # standardise_ = standardisation_function(
    #     str_to_mol_fn=ms.mol_from_smiles,
    #     std_fns=[
    #         ms.remove_stereochemistry_and_return_mol,
    #         papyrus_structure_pipeline.standardizer,
    #     ],  # ty: ignore[invalid-argument-type]
    #     mol_to_repr_fns={
    #         "smiles": ms.mol_to_kekulized_smiles,
    #         "inchi": ms.mol_to_inchi,
    #     },
    #     molstring_col="smiles",
    #     output_col="kekulised_smiles",
    # )

    # result = standardise_(lf).collect()
    # print(result)
