"""
In Machine Leaning, Training Data is the Data set on top of which Machine Learning Models are trained.

Training Data is generated by joining historical features created using Create Raw & Derived Feature Modules. Generation of Training Data implements point-in-time joins to ensure there's no data leakage.
Function:
    - print
    - register
    - generate
"""

from gru.utils.utils import check_data_type
from gru.utils.entity_type import EntityType
from typing import Dict, List, Optional
from gru.features.processing_engine_configs import BatchProcessingEngineConfigs
from gru.utils.processing_engine import ProcessingEngine


class TrainingData:
    """
    Training data is the join of historical data and the features for that historical data. It is used to train our Machine Learning Models

    Parameters
    ----------
    self: str
        Python self object
    name: str
        Name of the training data object
    description:
        description of the training data
    historical_data_source: str
        It'll be the name of historical data source registered in our registry. We'll fetch the timestamp field and timestamp format from the registry. So the end DS/DE doesn't have to mention it multiple times.
    entities:
        List of columns to be treated as entities in the historical data
    features:
        List of features to be joined with Historical data
    ttl:
        TTL (in minutes) used to determine how far back point-in-time joins should look
    owners:
        list of owners for the training data
    timeout:
        timeout value set for the job's execution to complete. By default value set to 30 mins.


    Returns
    -------
    None

    Notes
    -----
    Time to register the Training data will be generated in Stuart and then will be written to DB. And same will be passed on to Jerry and Bob

    Examples
    --------
    """

    def __init__(
        self,
        name: str,
        description: str,
        historical_data_source: str,
        entities: List[str],
        features: str,
        ttl: int,
        owners: List[str],
        timeout: int = 30,
        k8s_configs: Dict = {},
        processing_engine: Optional[ProcessingEngine] = ProcessingEngine.PYSPARK_K8S,
        processing_engine_configs: Optional[BatchProcessingEngineConfigs] = None,
    ):
        airflow_configs = {"timeout": timeout}

        self.type = EntityType.TRAINING_DATA.value
        self.entity_type = self.type
        self.ttl = ttl
        self.name = name
        self.owners = owners
        self.k8s_configs = k8s_configs
        self.entities = entities
        self.features = features
        self.description = description
        self.historical_data_source = historical_data_source
        self.airflow_configs = airflow_configs
        self.processing_engine = processing_engine
        self.processing_engine_configs = (
            processing_engine_configs or BatchProcessingEngineConfigs()
        )

        self.field_datatypes_dict = {
            int: [
                self.ttl,
                self.airflow_configs["timeout"],
            ],
            str: [
                self.name,
                self.description,
                self.historical_data_source,
            ],
            list: [
                self.owners,
                self.features,
            ],
        }

        self.args_type_check = check_data_type(self.field_datatypes_dict)

        if not self.args_type_check:
            return

    def to_json(self) -> Dict:
        return {
            "type": self.type,
            "ttl": self.ttl,
            "name": self.name,
            "owners": self.owners,
            "airflow_configs": self.airflow_configs,
            "entities": self.entities,
            "features": self.features,
            "description": self.description,
            "k8s_configs": self.k8s_configs,
            "processing_engine": self.processing_engine.value,
            "processing_engine_configs": self.processing_engine_configs.to_json(),
            "historical_data_source": self.historical_data_source,
        }
