import pandas as pd

from .types import DirectionalTrainingOutcome, TrainingOutcome
from training.train_model import train_model
from training.walk_forward import walk_forward_process_transformations

from typing import Optional
from config.types import Config
from models.base import Model
from models.model_map import default_feature_selector_classification

from transformations.scaler import get_scaler
from transformations.rfe import RFETransformation
from transformations.pca import PCATransformation


def train_directional_model(
    X: pd.DataFrame,
    y: pd.Series,
    forward_returns: pd.Series,
    config: Config,
    model: Model,
    from_index: Optional[pd.Timestamp],
    preloaded_training_step: Optional[DirectionalTrainingOutcome] = None,
) -> DirectionalTrainingOutcome:

    if preloaded_training_step is None:
        print("Preprocess transformations")
        transformations_over_time = walk_forward_process_transformations(
            X=X,
            y=y,
            forward_returns=forward_returns,
            expanding_window=config.expanding_window_base,
            window_size=config.sliding_window_size_base,
            retrain_every=config.retrain_every,
            from_index=from_index,
            transformations=[
                get_scaler(config.scaler),
                PCATransformation(
                    ratio_components_to_keep=0.5,
                    sliding_window_size=config.sliding_window_size_base,
                ),
                RFETransformation(
                    n_feature_to_select=40,
                    model=default_feature_selector_classification,
                ),
            ],
        )
    else:
        transformations_over_time = preloaded_training_step.transformations

    training_outcome = train_model(
        ticker_to_predict=config.target_asset[1],
        X=X,
        y=y,
        forward_returns=forward_returns,
        model=model,
        expanding_window=config.expanding_window_base,
        sliding_window_size=config.sliding_window_size_base,
        retrain_every=config.retrain_every,
        from_index=from_index,
        no_of_classes=config.no_of_classes,
        level="primary",
        output_stats=config.mode == "training",
        transformations_over_time=transformations_over_time,
        model_over_time=preloaded_training_step.training.model_over_time
        if preloaded_training_step
        else None,
    )
    if config.mode == "training":
        print(training_outcome.stats)

    return DirectionalTrainingOutcome(training_outcome, transformations_over_time)