"""Ingest FoldX data from OTAR2081 project."""

from __future__ import annotations

from typing import TYPE_CHECKING

from gentropy.common.session import Session
from gentropy.datasource.open_targets.foldex_integration import OpenTargetsFoldX

if TYPE_CHECKING:
    from gentropy.dataset.amino_acid_variants import AminoAcidVariants


class FoldXIngestionStep:
    """Step to ingest proteome-wide FoldX dataset generated by the OTAR2081 project."""

    def __init__(
        self,
        session: Session,
        foldx_dataset_path: str,
        plddt_threshold: float,
        annotation_path: str,
    ) -> None:
        """Initialize step.

        Args:
            session (Session): Session object.
            foldx_dataset_path (str): path to the FoldX dataset.
            plddt_threshold (float): plddt threshold to filter amino acids in the structural model.
            annotation_path (str): path of the output dataset.
        """
        fold_x_data = session.spark.read.parquet(foldx_dataset_path)
        # Transform
        gene_index: AminoAcidVariants = OpenTargetsFoldX.ingest_foldx_data(
            fold_x_data, plddt_threshold
        )
        # Load
        gene_index.df.coalesce(session.output_partitions).write.mode(
            session.write_mode
        ).parquet(annotation_path)
