sientia-dataops-model-manager/model_manager/utils/models/train_model_params.py

from dataclasses import dataclass
from typing import Any

from jsonschema import Draft202012Validator, ValidationError  # type: ignore[import-untyped]

# Allowed frontend date formats and their strftime equivalents (single source of truth)
FRONTEND_DATE_FORMAT_TO_STRFTIME = {
    'dd/MM/yyyy HH:mm:ss': '%d/%m/%Y %H:%M:%S',
    'MM/dd/yyyy HH:mm:ss': '%m/%d/%Y %H:%M:%S',
    'yyyy/MM/dd HH:mm:ss': '%Y/%m/%d %H:%M:%S',
    'dd-MM-yyyy HH:mm:ss': '%d-%m-%Y %H:%M:%S',
    'MM-dd-yyyy HH:mm:ss': '%m-%d-%Y %H:%M:%S',
    'yyyy-MM-dd HH:mm:ss': '%Y-%m-%d %H:%M:%S',
}
ALLOWED_FRONTEND_DATE_FORMATS = frozenset(FRONTEND_DATE_FORMAT_TO_STRFTIME.keys())

# When the client omits date_format (or sends null/blank), parsing uses this frontend format.
DEFAULT_TRAIN_DATE_FORMAT = 'yyyy-MM-dd HH:mm:ss'


def validate_frontend_date_format(fmt: str | None) -> None:
    """Raise ValueError if fmt is set and not one of the allowed frontend date formats."""
    if not fmt or not fmt.strip():
        return
    if fmt not in ALLOWED_FRONTEND_DATE_FORMATS:
        allowed = ', '.join(sorted(ALLOWED_FRONTEND_DATE_FORMATS))
        raise ValueError(f'Invalid date_format "{fmt}". Allowed formats: {allowed}')


# Model name constants
MODEL_LINEAR_REGRESSION = 'Linear Regression'
MODEL_POLYNOMIAL_REGRESSION = 'Polynomial Regression'


@dataclass
class TrainModelParams:
    """
    Parameters for machine learning model training.

    This class encapsulates all configuration parameters required for the training
    pipeline, including data processing settings, model configuration, and experiment
    tracking information. All parameters are validated upon initialization to ensure
    data integrity and prevent runtime errors.

    Use the `from_dict()` class method to create instances from dictionaries with
    automatic validation of all fields.

    Attributes:
        variable_columns (list[str]): List of variable column names to use as features.
        target_variable (str): Name of the target variable to predict.
        bucket_name (str): Name of the MinIO bucket containing training data.
        file_name (str): Name of the file in the MinIO bucket.
        line_separator (str): Line separator used in the CSV file.
        decimal_separator (str): Decimal separator used in the CSV file.
        date_column (str): Name of the date/time column in the dataset (required).
        date_format (str): Format of the date column (allowed frontend strings). If omitted or blank
            in the input dict, defaults to DEFAULT_TRAIN_DATE_FORMAT.
        train_size (int): Percentage of data to use for training (0-100).
        shuffle (bool): Whether to shuffle the data during train/test split.
        experiment_run_id (int): Unique identifier for the experiment run.
        model_name (str): Name of the model type ('Linear Regression' or 'Polynomial Regression').
        val_file_name (str | None): Name of the validation file in the MinIO bucket.
        data_model_kwargs (dict | None): Keyword arguments for the data model.
        model_kwargs (dict | None): Keyword arguments for the model.
        opt_params (dict | None): Keyword optimazation arguments for the wrapper.
        model_type (str): Type of the model to use (ex.: 'Linear Regression', 'XGBoost').
    """

    # Old Parameters (keep)

    variable_columns: list[str]
    target_variable: str
    bucket_name: str
    file_name: str
    line_separator: str
    decimal_separator: str
    date_column: str
    date_format: str
    train_size: int
    shuffle: bool
    random_state: int
    experiment_run_id: int
    model_name: str
    experiment_name: str

    # New Parameters
    val_file_name: str | None
    data_model_kwargs: dict | None  # Removed params used in DataPreprocessor here
    model_kwargs: dict | None  # Removed params used in Linear Regression Model here
    opt_params: dict | None
    model_type: str
    model_id: str | None

    # Context Parameters
    model_metadata: dict | None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> 'TrainModelParams':
        """
        Create TrainModelParams from dictionary with validation.

        This factory method creates a TrainModelParams instance from a dictionary,
        applying validation to ensure all required fields are present and have
        the correct types. This is the recommended way to create instances from
        workflow input data.

        Args:
            data: Dictionary containing training parameters with keys matching the
                  attribute names (e.g. variable_columns, date_column, data_model_kwargs, model_kwargs, opt_params).
                  Unknown keys are ignored by from_dict; missing required snake_case keys raise.
                  model_metadata may be omitted or None until load_model_metadata fills it.
                  experiment_run_id may be an int or numeric string.

        Returns:
            TrainModelParams: Validated instance with all fields populated

        Raises:
            ValueError: If any required field is missing or None
            TypeError: If any field has an incorrect type
            KeyError: If any required key is missing from the dictionary
        """
        # `from_dict()` should only build the "raw" object from the input dict.
        # Semantic validation and defaults must be handled by `validate_business_rules()`
        # (using `model_metadata` JSON Schemas).

        model_name = cls._check_none(data.get('model_name'), str, 'model_name')

        return cls(
            variable_columns=cls._check_none(
                data.get('variable_columns'), list, 'variable_columns'
            ),
            target_variable=cls._check_none(data.get('target_variable'), str, 'target_variable'),
            bucket_name=cls._check_none(data.get('bucket_name'), str, 'bucket_name'),
            file_name=cls._check_none(data.get('file_name'), str, 'file_name'),
            line_separator=cls._check_none(data.get('line_separator'), str, 'line_separator'),
            decimal_separator=cls._check_none(
                data.get('decimal_separator'), str, 'decimal_separator'
            ),
            date_column=cls._check_none(data.get('date_column'), str, 'date_column'),
            date_format=cls._resolve_date_format(data.get('date_format')),
            train_size=cls._check_none(data.get('train_size'), int, 'train_size'),
            shuffle=cls._check_none(data.get('shuffle'), bool, 'shuffle'),
            random_state=cls._check_none(data.get('random_state', 42), int, 'random_state'),
            experiment_run_id=cls._coerce_experiment_run_id(data.get('experiment_run_id')),
            model_name=model_name,
            experiment_name=model_name,
            val_file_name=data.get('val_file_name'),
            data_model_kwargs=cls._check_none(
                data.get('data_model_kwargs'), dict, 'data_model_kwargs'
            ),
            model_kwargs=cls._check_none(data.get('model_kwargs'), dict, 'model_kwargs'),
            opt_params=cls._check_none(data.get('opt_params'), dict, 'opt_params'),
            model_type=cls._check_none(data.get('model_type'), str, 'model_type'),
            model_id=data.get('model_id'),
            model_metadata=cls._parse_optional_model_metadata(data.get('model_metadata')),
        )

    @staticmethod
    def _resolve_date_format(raw: Any) -> str:
        """
        Resolve date_format from workflow input.

        Omitted, null, or blank values use DEFAULT_TRAIN_DATE_FORMAT. Non-string types raise.

        Args:
            raw: Raw date_format from the payload, or None if absent.

        Return:
            str: Canonical frontend date format string.
        """
        if raw is None:
            return DEFAULT_TRAIN_DATE_FORMAT
        if isinstance(raw, str) and not raw.strip():
            return DEFAULT_TRAIN_DATE_FORMAT
        if not isinstance(raw, str):
            raise TypeError(
                f'date_format must be a string or omitted, but got {type(raw).__name__}.'
            )
        return raw.strip()

    def to_dict(self) -> dict[str, Any]:
        """
        Convert TrainModelParams to a dictionary.
        """
        return self.__dict__

    @staticmethod
    def _check_none(value: Any | None, expected_type: type, field_name: str) -> Any:
        """
        Validate that a value is not None and check its type.

        This method ensures that required parameters are provided and have the
        correct type, raising descriptive errors if validation fails.

        Args:
            value (Any | None): The value to validate.
            expected_type (type): The expected type of the value.
            field_name (str): The name of the field being validated (for error messages).

        Returns:
            Any: The validated value if it is not None and matches the expected type.

        Raises:
            ValueError: If the value is None.
            TypeError: If the value is not of the expected type.
        """
        if value is None:
            error = f'{field_name} is required and cannot be None.'
            raise ValueError(error)

        return TrainModelParams._check_type(value, expected_type, field_name)

    @staticmethod
    def _check_type(value: Any | None, expected_type: type, field_name: str) -> Any:
        """
        Validate that a value matches the expected type.

        This method checks type compatibility and raises a descriptive error
        if the value does not match the expected type.

        Args:
            value (Any | None): The value to validate.
            expected_type (type): The expected type of the value.
            field_name (str): The name of the field being validated (for error messages).

        Returns:
            Any: The validated value if it matches the expected type.

        Raises:
            TypeError: If the value is not of the expected type.
        """
        if value is not None and not isinstance(value, expected_type):
            error = f'{field_name} must be of type {expected_type.__name__}, but got {type(value).__name__}.'
            raise TypeError(error)

        return value

    @staticmethod
    def _coerce_experiment_run_id(value: Any) -> int:
        """
        Coerce experiment_run_id to int.

        Workflow clients may send numeric strings; this keeps from_dict aligned with
        workflow validation.

        Args:
            value: Raw experiment_run_id from the payload.

        Returns:
            int: Parsed experiment run id.

        Raises:
            ValueError: If the value is None.
            TypeError: If the value cannot be coerced to a non-boolean integer.
        """
        if value is None:
            raise ValueError('experiment_run_id is required and cannot be None.')
        if isinstance(value, bool):
            raise TypeError('experiment_run_id must be an integer, got bool.')
        if isinstance(value, int):
            return value
        if isinstance(value, str) and value.strip().isdigit():
            return int(value.strip())
        if isinstance(value, float) and value.is_integer():
            return int(value)
        raise TypeError(
            f'experiment_run_id must be an integer or numeric string, but got {type(value).__name__}.'
        )

    @staticmethod
    def _parse_optional_model_metadata(value: Any) -> dict | None:
        """
        Parse model_metadata for from_dict before load_model_metadata fills the index.

        Args:
            value: model_metadata from the payload, or None if not sent yet.

        Returns:
            dict | None: Dict when provided; None when absent (filled later by load_model_metadata).

        Raises:
            TypeError: If value is neither None nor a dict.
        """
        if value is None:
            return None
        if isinstance(value, dict):
            return value
        raise TypeError(f'model_metadata must be a dict or None, but got {type(value).__name__}.')

    def validate_business_rules(self) -> None:
        """
        Validate business rules and constraints for training parameters.

        This method performs additional validation beyond type checking to ensure
        that parameter values are within acceptable ranges and logically consistent.
        It implements defense-in-depth validation to catch configuration errors
        early in the workflow.

        Raises:
            ValueError: If any business rule is violated
        """
        self._validate_numeric_ranges()
        self._validate_model_params()
        self._validate_required_strings()
        self._validate_date_format()

    def _validate_numeric_ranges(self) -> None:
        """Validate numeric parameters are within acceptable ranges."""
        if not 10 <= self.train_size <= 100:
            raise ValueError(f'train_size must be between 10 and 100, got {self.train_size}')

        if not self.variable_columns:
            raise ValueError('variable_columns cannot be empty')

    def _validate_model_params(self) -> None:
        """Validate model-related parameters."""
        if not self.model_metadata:
            raise ValueError('model_metadata is required')

        schemas = self.model_metadata.get('schemas', {}).get('components', {}).get('schemas')

        if not schemas:
            return

        data_model_schema = schemas.get('data_model')
        model_schema = schemas.get('model')
        opt_params_schema = schemas.get('opt_params')

        if data_model_schema:
            self._validate_model_param(data_model_schema, self.data_model_kwargs)
        if model_schema:
            self._validate_model_param(model_schema, self.model_kwargs)
        if opt_params_schema:
            self._validate_model_param(opt_params_schema, self.opt_params)

    def _validate_model_param(self, schema: dict[str, Any], value: Any) -> None:
        """Validate model parameter against schema."""
        try:
            validator = Draft202012Validator(schema)
            validator.validate(value)
        except ValidationError as e:
            raise ValueError(f'Model parameters validation failed: {e.message}') from e

    def _validate_required_strings(self) -> None:
        """Validate required string fields are not empty."""
        if not self.target_variable.strip():
            raise ValueError('target_variable cannot be empty or whitespace')

        if not self.bucket_name.strip():
            raise ValueError('bucket_name cannot be empty or whitespace')

        if not self.file_name.strip():
            raise ValueError('file_name cannot be empty or whitespace')

        if not self.model_name.strip():
            raise ValueError('model_name cannot be empty or whitespace')

        if not self.date_column.strip():
            raise ValueError('date_column cannot be empty or whitespace')

    def _validate_date_format(self) -> None:
        """Validate date_format is one of the allowed frontend formats."""
        validate_frontend_date_format(self.date_format)