"""Parser integrating FoldX data from OpenTargets project OTAR2081."""

from __future__ import annotations

import pyspark.sql.functions as f
from pyspark.sql import Column, DataFrame
from pyspark.sql import types as t

from gentropy.common.spark import enforce_schema
from gentropy.dataset.amino_acid_variants import AminoAcidVariants
from gentropy.dataset.variant_index import VariantEffectNormaliser


class OpenTargetsFoldX:
    """Class to parser FoldX dataset generated by the OTAR2081 project."""

    VARIANT_EFFECT_SCHEMA = AminoAcidVariants.get_schema()[
        "variantEffect"
    ].dataType.elementType

    @staticmethod
    @enforce_schema(VARIANT_EFFECT_SCHEMA)
    def get_foldx_prediction(score_column: Column) -> Column:
        """Generate variantEffect object from ddG column.

        Args:
            score_column (Column): ddG column from the FoldX dataset.

        Returns:
            Column: struct with the right shape of the variantEffect field.
        """
        return f.struct(
            f.lit("FoldX").alias("method"),
            score_column.cast(t.FloatType()).alias("score"),
        )

    @classmethod
    def ingest_foldx_data(
        cls: type[OpenTargetsFoldX], foldx_input: DataFrame, plddt_threshold: float
    ) -> AminoAcidVariants:
        """Ingest FoldX dataset and convert into a AminoAcidVariants object.

        Args:
            foldx_input (DataFrame): Input dataframe provided by the FoldX project.
            plddt_threshold (float): lower threshold for filtering confident residues from structural models.

        Returns:
            AminoAcidVariants: _description_
        """
        excluded_identifiers = cls._uniprot_ids_to_exclude(foldx_input)
        return AminoAcidVariants(
            _df=(
                foldx_input.filter(f.col("plddt") > plddt_threshold)
                .join(excluded_identifiers, on="protein_acc", how="left_anti")
                .select(
                    f.col("protein_acc").alias("uniprotAccession"),
                    f.concat(
                        f.col("wild_type"), f.col("position"), f.col("mutated_type")
                    ).alias("aminoAcidChange"),
                    cls.get_foldx_prediction(f.col("foldx_ddg")).alias(
                        "foldx_prediction"
                    ),
                )
                # Collapse all predictors for a single array object to avoid variant explosions:
                .groupBy("uniprotAccession", "aminoAcidChange")
                .agg(f.collect_set(f.col("foldx_prediction")).alias("variantEffect"))
                # Normalise FoldX free energy changes:
                .withColumn(
                    "variantEffect",
                    VariantEffectNormaliser.normalise_variant_effect(
                        f.col("variantEffect")
                    ),
                )
            ),
            _schema=AminoAcidVariants.get_schema(),
        )

    @staticmethod
    def _uniprot_ids_to_exclude(foldx_input: DataFrame) -> DataFrame:
        """Compute distinct set of UniprotIDs to drop from the input dataset.

        Exclude UniprotIds, where one position in the structure corresponds to multiple positions in the original sequence.
        Such cases are impossible to disambiguate.

        Args:
            foldx_input (DataFrame): raw dataset.

        Returns:
            DataFrame: one column with uniprot ids.
        """
        return (
            foldx_input.groupby("protein_acc", "position", "wild_type")
            .agg(f.collect_set("plddt").alias("plddts"))
            .filter(f.size("plddts") > 1)
            .select(
                "protein_acc",
            )
            .distinct()
        )
