"""Ingest Loss-of-Function variant data generated by OTAR2075."""

from gentropy.common.session import Session
from gentropy.datasource.open_targets.lof_curation import OpenTargetsLOF


class LOFIngestionStep:
    """Step to ingest the Loss-of-Function dataset generated by OTAR2075."""

    def __init__(
        self,
        session: Session,
        lof_curation_dataset_path: str,
        lof_curation_variant_annotations_path: str,
    ) -> None:
        """Initialize step.

        Args:
            session (Session): Session object.
            lof_curation_dataset_path (str): path of the curated LOF dataset.
            lof_curation_variant_annotations_path (str): path of the resulting variant annotations.
        """
        # Read in data:
        lof_dataset = session.spark.read.csv(
            lof_curation_dataset_path,
            sep=",",
            header=True,
            multiLine=True
        )
        # Extract relevant information to a VariantIndex
        lof_variant_annotations = OpenTargetsLOF.as_variant_index(lof_dataset)
        # Write to file:
        (
            lof_variant_annotations.df
            .coalesce(session.output_partitions)
            .write.mode(session.write_mode)
            .parquet(lof_curation_variant_annotations_path)
        )
