"""
Created on: 1/21/2025
Original author: Adil Zaheer
"""

# Built-Ins
import logging
from pathlib import Path

# Third Party
import pandas as pd

# Local Imports
from caf.brain.ml._functions._ml_inputs import (
    DataClassificationInputs,
    ModellingInputs,
    Paths,
    TransformingInputDataInputs,
)
from caf.brain.ml._functions.feature_selection.functions import (
    analyse_feature_importance,
    combine_results,
    rf_feature_selection,
)
from caf.brain.ml._functions.process_data_functions.main import (
    main_input_data,
)

LOG = logging.getLogger(__name__)


def main_feature_selection(
    paths: Paths,
    data_classification: DataClassificationInputs,
    transforming_inputs: TransformingInputDataInputs,
    modelling: ModellingInputs,
    train: pd.DataFrame = None,
    test: pd.DataFrame = None,
    initialised_model=None,
    output: Path | None = None,
):
    """
    Feature selection function.

    Parameters
    ----------
    paths: Path inputs from the PredictionModelInputs class. These inputs
           define paths to external files. See
           caf/brain/ml/main_models/prediction_model/_ml_inputs.py
           for available options.
    data_classification: Data classification inputs from the PredictionModelInputs
                         class. These inputs help define and outline the
                         structure of the input data. See
                         caf/brain/ml/main_models/prediction_model/_ml_inputs.py
                         for available options.
    transforming_inputs: Transforming inputs from the PredictionModelInputs
                         class. These inputs dictate how the data is transformed
                         for machine learning modelling. See
                         caf/brain/ml/main_models/prediction_model/_ml_inputs.py
                         for available options.
    modelling: Modelling inputs from the PredictionModelInputs
               class. These inputs control the machine learning modelling
               pipeline and _functions. See
               caf/brain/ml/main_models/prediction_model/_ml_inputs.py
               for available options.
    train: Transformed input data split into training set.
    test: Transformed input data split into training set.
    initialised_model: Initialised model algorithm from Models enum class.
    output: Path to output location.

    Returns
    -------
    train_final: Dataframe of final training data post feature selection.
    test_final: Dataframe of final test data post feature selection.
    cols_dropped_by_feat_select: These are the columns removed due to
                                 feature selection.
    """

    if modelling.skip_feature_selection:
        return train, test, None

    if output is None:
        if paths.output_path is None:
            raise ValueError("paths.output_path must not be None when output is not provided.")
        output = Path(paths.output_path) / "output"
        output.mkdir(parents=True, exist_ok=True)

    if train is None and test is None:
        LOG.info("Train and test are none so they are being generated by the \
                  main_input_data function")
        data_dict, _, _ = main_input_data(
            paths, data_classification, transforming_inputs, output_path=output
        )

        train = pd.DataFrame.from_dict(data_dict["train_scaled"])
        test = pd.DataFrame.from_dict(data_dict["test_scaled"])

    if modelling.intensive_feature_selection and initialised_model is None:
        LOG.warning("In order to use intensive feature selection, you must \
                     provide an initialised model from the Models class. This \
                     this can be done using the main_model_selection function. \
                     Standard feature selection being ran instead")
        train_final = analyse_feature_importance(
            train_transformed=train,
            target_column=data_classification.target_column,
            weight_column=data_classification.weight_column,
            output_path=output,
            is_time_series=data_classification.is_time_series,
        )

        test_final, cols_dropped_by_feat_select = combine_results(
            train_final=train_final,
            target_column=data_classification.target_column,
            weight_column=data_classification.weight_column,
            test=test,
        )
        return train_final, test_final, cols_dropped_by_feat_select

    if modelling.intensive_feature_selection:
        # eval no. samples
        n_features = (
            len(train.columns)
            - (1 if data_classification.target_column in train.columns else 0)
            - (1 if data_classification.weight_column else 0)
        )
        n_samples = len(train)

        min_samples_per_feature = 15
        required_samples = n_features * min_samples_per_feature

        if n_samples < required_samples:
            LOG.warning("Insufficient data for intensive feature selection. \
                 Falling back to standard feature importance analysis.")
            train_final = analyse_feature_importance(
                train_transformed=train,
                target_column=data_classification.target_column,
                weight_column=data_classification.weight_column,
                output_path=output,
                is_time_series=data_classification.is_time_series,
            )

            test_final, cols_dropped_by_feat_select = combine_results(
                train_final=train_final,
                target_column=data_classification.target_column,
                weight_column=data_classification.weight_column,
                test=test,
            )
            return train_final, test_final, cols_dropped_by_feat_select

        LOG.info("Intensive feature selection running")
        train_final = rf_feature_selection(
            data=train,
            target_column=data_classification.target_column,
            cv=modelling.cv,
            regression_method=initialised_model,
            weight_column=data_classification.weight_column,
            classification_prediction=transforming_inputs.classification_prediction,
            is_time_series=data_classification.is_time_series,
        )

        test_final, cols_dropped_by_feat_select = combine_results(
            train_final=train_final,
            target_column=data_classification.target_column,
            weight_column=data_classification.weight_column,
            test=test,
        )

        return train_final, test_final, cols_dropped_by_feat_select

    LOG.info("Standard feature importance running")
    train_final = analyse_feature_importance(
        train_transformed=train,
        target_column=data_classification.target_column,
        weight_column=data_classification.weight_column,
        output_path=output,
        is_time_series=data_classification.is_time_series,
    )

    test_final, cols_dropped_by_feat_select = combine_results(
        train_final=train_final,
        target_column=data_classification.target_column,
        weight_column=data_classification.weight_column,
        test=test,
    )

    return train_final, test_final, cols_dropped_by_feat_select
