"""
DCC Experiment Extractor module
    This module takes care of extracting experiment data from XML DCC files to  Spark DataFrames.
    In the XML files we can find two experiment classes: specimen level experiments and line level experiments.

    The files are expected to be organized by data source in the following directory structure <DCC_XML_PATH>/<DATASOURCE>/*experiment*.xml
    Each directory containing the raw XML in the XML schema defined by the DCC.
"""
import luigi
from luigi.contrib.spark import PySparkTask
from pyspark import SparkContext
from pyspark.sql import DataFrame, SparkSession

from impc_etl.jobs.extract.xml_extraction_helper import (
    extract_dcc_xml_files,
    get_entity_by_type,
)
from impc_etl.shared.exceptions import UnsupportedEntityError
from impc_etl.workflow.config import ImpcConfig


class ExperimentExtractor(PySparkTask):
    """
    PySpark Task class to extract experimental data from the XML files compliant with the DCC XML format.
    """

    #: Name of the Spark task
    name: str = "IMPC_Experiment_Extractor"

    #: Path in the filesystem (local or HDFS) to the experiment XML files
    dcc_experiment_xml_path: luigi.Parameter = luigi.Parameter()

    #: Type of experiment can be "specimen_level" or "line_level"
    experiment_type: luigi.Parameter = luigi.Parameter()

    #:  Path of the output directory where the new parquet file will be generated.
    output_path: luigi.Parameter = luigi.Parameter()

    def output(self):
        """
        Returns the full parquet path as an output for the Luigi Task
        (e.g. impc/dr15.2/parquet/specimen_level_experiment_raw_parquet)
        """
        return ImpcConfig().get_target(
            f"{self.output_path}{self.experiment_type}_experiment_raw_parquet"
        )

    def app_options(self):
        """
        Generates the options pass to the PySpark job
        """
        return [
            self.dcc_experiment_xml_path,
            self.experiment_type,
            self.output().path,
        ]

    def main(self, sc: SparkContext, *args):
        """
        Takes in a SparkContext and the list of arguments generated by `app_options` and executes the PySpark job.
        """
        spark = SparkSession(sc)
        # TODO update time parsing strategy
        # Setting time parser policy to legacy so it behaves as expected
        spark.sql("set spark.sql.legacy.timeParserPolicy=LEGACY")
        spark.sql("set spark.sql.legacy.parquet.datetimeRebaseModeInWrite=LEGACY")

        # Parsing app options
        dcc_experiment_xml_path = args[0]
        experiment_type = args[1]
        output_path = args[2]

        dcc_df = extract_dcc_xml_files(spark, dcc_experiment_xml_path, "experiment")
        experiment_df = self.get_experiments_by_type(dcc_df, experiment_type)
        experiment_df.write.mode("overwrite").parquet(output_path)

    def get_experiments_by_type(self, dcc_df: DataFrame, entity_type: str) -> DataFrame:
        """
        Takes a DataFrame generated by `impc_etl.jobs.extract.dcc_extractor_helper.extract_dcc_xml_files`, an entity_type
        (can be 'specimen_level' or 'line_level') and gets the experiment data entities.
        It also expands the procedure struct column
        and adds all its fields to the top level. raises UnsupportedEntityError when the given entity_type is not supported.
        """
        if entity_type not in ["specimen_level", "line_level"]:
            raise UnsupportedEntityError
        input_to_xml_entity_map = {"specimen_level": "experiment", "line_level": "line"}
        entity_type = input_to_xml_entity_map[entity_type]
        experiment_df = get_entity_by_type(
            dcc_df,
            entity_type,
            [
                "_centreID",
                "_pipeline",
                "_project",
                "_sourceFile",
                "_dataSource",
                "_sourcePhenotypingStatus",
            ],
        )
        return experiment_df.select(
            ["procedure.*"]
            + [column for column in experiment_df.columns if column is not "procedure"]
        ).drop("procedure")


class SpecimenLevelExperimentExtractor(ExperimentExtractor):
    name: str = "IMPC_Specimen_Level_Experiment_Extractor"
    experiment_type: str = "specimen_level"


class LineLevelExperimentExtractor(ExperimentExtractor):
    name: str = "IMPC_Line_Level_Experiment_Extractor"
    experiment_type: str = "line_level"
