"""
    Colony tracking extraction task. Extracts information from GenTar.
"""
import luigi
from luigi.contrib.spark import PySparkTask
from pyspark import SparkContext
from pyspark.sql.session import SparkSession

from impc_etl.shared import utils
from impc_etl.workflow.config import ImpcConfig


class ColonyTrackingExtractor(PySparkTask):
    """
    PySpark Task class to extract colony tracking information from reports.
    """

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

    #: Path in the filesystem (local or HDFS) to the GenTar colonies report
    gentar_colonies_tsv_path: luigi.Parameter = luigi.Parameter()

    #: Path of the output directory where ethe 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/colonies_tracking_parquet)
        """
        return ImpcConfig().get_target(f"{self.output_path}colonies_tracking_parquet")

    def app_options(self):
        """
        Generates the options pass to the PySpark job
        """
        return [
            self.gentar_colonies_tsv_path,
            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)
        # Parsing app options
        gentar_tsv_path = args[0]
        output_path = args[1]

        # Load the data form TSV to a Spark DataFrame
        gentar_df = utils.extract_tsv(spark, gentar_tsv_path)

        # Map GenTar Column names to match the ones on the Observations Schema Report
        gentar_col_mapping = {
            "Phenotyping External Reference": "colony_name",
            "Background Strain": "colony_background_strain",
            "Mutation Symbol": "allele_symbol",
            "Gene Symbol": "marker_symbol",
            "MGI Gene Accession ID": "mgi_accession_id",
            "MGI Strain Accession ID": "mgi_strain_accession_id",
            "Phenotyping Work Unit": "phenotyping_centre",
            "Phenotyping Work Group": "phenotyping_consortium",
            "Production Work Unit": "production_centre",
            "Production Work Group": "production_consortium",
        }
        new_col_names = []
        for col_name in gentar_df.columns:
            if col_name in gentar_col_mapping:
                new_col_names.append(gentar_col_mapping[col_name])
            else:
                new_col_names.append(col_name.replace(" ", "_").lower())
        gentar_df = gentar_df.toDF(*new_col_names)
        gentar_df.write.parquet(output_path)
