# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# hide
# %load_ext autoreload
# %autoreload 2
# %load_ext nb_black
# %load_ext lab_black

# +
# default_exp model_pipeline
# -

# # ModelPipeline
# > Putting it all together.

# ## Overview
#
# The functionality below uses the `NumerFrame`, `PreProcessor`, `Model` and `PostProcessor` objects to easily propagate
# data, generate predictions and postprocess them in one go.
#
# Specifically, this section introduces two objects:
# 1. `ModelPipeline`: Run all preprocessing, models and postprocessing that you define and return a `NumerFrame`.
# 2. `ModelPipelineCollection`: Manage and run multiple `ModelPipeline` objects.

# hide
from nbdev.showdoc import *

# +
#export
import uuid
import pandas as pd
from tqdm.auto import tqdm
from typeguard import typechecked
from typing import List, Union, Dict
from rich import print as rich_print

from numerblox.numerframe import NumerFrame, create_numerframe
from numerblox.preprocessing import BaseProcessor, CopyPreProcessor, GroupStatsPreProcessor, FeatureSelectionPreProcessor
from numerblox.model import BaseModel, ConstantModel, RandomModel
from numerblox.postprocessing import Standardizer, MeanEnsembler, FeatureNeutralizer


# -

# ## 1. ModelPipeline

# `ModelPipeline` handles all preprocessing, model prediction and postprocessing. It returns a `NumerFrame` with the preprocessed data, metadata and postprocessed prediction columns.

#export
@typechecked
class ModelPipeline:
    """
    Execute all preprocessing, prediction and postprocessing for a given setup.

    :param models: Initiliazed numerai-blocks Models (Objects inheriting from BaseModel) \n
    :param preprocessors: List of initialized Preprocessors. \n
    :param postprocessors: List of initialized Postprocessors. \n
    :param copy_first: Whether to copy the NumerFrame as a first preprocessing step. \n
    Highly recommended in order to avoid surprise behaviour by manipulating the original dataset. \n
    :param pipeline_name: Unique name for pipeline. Only used for display purposes.
    """
    def __init__(self,
                 models: List[BaseModel],
                 preprocessors: List[BaseProcessor] = [],
                 postprocessors: List[BaseProcessor] = [],
                 copy_first = True,
                 standardize = True,
                 pipeline_name: str = None):
        self.pipeline_name = pipeline_name if pipeline_name else uuid.uuid4().hex
        self.models = models
        self.copy_first = copy_first
        self.standardize = standardize
        self.preprocessors = preprocessors
        self.postprocessors = postprocessors

    def preprocess(self, dataf: Union[pd.DataFrame, NumerFrame]) -> NumerFrame:
        """ Run all preprocessing steps. Copies input by default. """
        if self.copy_first:
            dataf = CopyPreProcessor()(dataf)
        for preprocessor in tqdm(self.preprocessors,
                                 desc=f"{self.pipeline_name} Preprocessing:",
                                 position=0):
            rich_print(f":construction: Applying preprocessing: '[bold]{preprocessor.__class__.__name__}[/bold]' :construction:")
            dataf = preprocessor(dataf)
        return NumerFrame(dataf)

    def postprocess(self, dataf: Union[pd.DataFrame, NumerFrame]) -> NumerFrame:
        """ Run all postprocessing steps. Standardizes model prediction by default. """
        if self.standardize:
            dataf = Standardizer()(dataf)
        for postprocessor in tqdm(self.postprocessors,
                                  desc=f"{self.pipeline_name} Postprocessing: ",
                                  position=0):
            rich_print(f":construction: Applying postprocessing: '[bold]{postprocessor.__class__.__name__}[/bold]' :construction:")
            dataf = postprocessor(dataf)
        return NumerFrame(dataf)

    def process_models(self, dataf: Union[pd.DataFrame, NumerFrame]) -> NumerFrame:
        """ Run all models. """
        for model in tqdm(self.models,
                                  desc=f"{self.pipeline_name} Model prediction: ",
                                  position=0):
            rich_print(f":robot: Generating model predictions with '[bold]{model.__class__.__name__}[/bold]'. :robot:")
            dataf = model(dataf)
        return NumerFrame(dataf)

    def pipeline(self, dataf: Union[pd.DataFrame, NumerFrame]) -> NumerFrame:
        """ Process full pipeline and return resulting NumerFrame. """
        preprocessed_dataf = self.preprocess(dataf)
        prediction_dataf = self.process_models(preprocessed_dataf)
        processed_prediction_dataf = self.postprocess(prediction_dataf)
        rich_print(f":checkered_flag: [green]Finished pipeline:[green] [bold blue]'{self.pipeline_name}'[bold blue]! :checkered_flag:")
        return processed_prediction_dataf

    def __call__(self, dataf: Union[pd.DataFrame, NumerFrame]) -> NumerFrame:
        return self.pipeline(dataf)


# Example using several preprocessor, dummy models and postprocessors

# +
model_names = ["test_0.5", "test_0.8"]

dataf = create_numerframe("test_assets/mini_numerai_version_1_data.csv", metadata={'version': 1})
preprocessors = [GroupStatsPreProcessor(), FeatureSelectionPreProcessor(feature_cols=['feature_intelligence_mean', 'feature_intelligence_std'])]
models = [ConstantModel(constant=0.5, model_name=model_names[0]), ConstantModel(constant=0.8, model_name=model_names[1])]
postprocessors = [MeanEnsembler(cols=[f"prediction_{name}" for name in model_names], final_col_name='prediction_ensembled'),
                  FeatureNeutralizer(feature_names=['feature_intelligence_mean', 'feature_intelligence_std'],
                                     pred_name='prediction_ensembled', proportion=0.8)]
# -

test_pipeline = ModelPipeline(preprocessors=preprocessors, models=models,
                              postprocessors=postprocessors, pipeline_name="test_pipeline",
                              standardize=False)
processed_dataf = test_pipeline(dataf)

assert processed_dataf.meta == dataf.meta
assert isinstance(processed_dataf, NumerFrame)
processed_dataf.head(2)


# ## 2. ModelPipelineCollection

# `ModelPipelineCollection` can be used to manage and run multiple `ModelPipeline` objects.
#
# `ModelPipelineCollection` simply takes a list of `ModelPipeline` objects as input.

#export
@typechecked
class ModelPipelineCollection:
    """
    Execute multiple initialized ModelPipelines in a sequence.

    :param pipelines: List of initialized ModelPipelines.
    """
    def __init__(self, pipelines: List[ModelPipeline]):
        self.pipelines = {pipe.pipeline_name: pipe for pipe in pipelines}
        self.pipeline_names = list(self.pipelines.keys())

    def process_all_pipelines(self, dataf: Union[pd.DataFrame, NumerFrame]) -> Dict[str, NumerFrame]:
        """ Process all pipelines and return Dictionary mapping pipeline names to resulting NumerFrames. """
        result_datafs = dict()
        for name, pipeline in tqdm(self.pipelines.items(),
                                   desc="Processing Pipeline Collection"):
            result_datafs[name] = self.process_single_pipeline(dataf, name)
        return result_datafs

    def process_single_pipeline(self, dataf: Union[pd.DataFrame, NumerFrame], pipeline_name: str) -> NumerFrame:
        """ Run full model pipeline for given name in collection. """
        rich_print(f":construction_worker: [bold green]Processing model pipeline:[/bold green] '{pipeline_name}' :construction_worker:")
        pipeline = self.get_pipeline(pipeline_name)
        dataf = pipeline(dataf)
        return NumerFrame(dataf)

    def get_pipeline(self, pipeline_name: str) -> ModelPipeline:
        """ Retrieve model pipeline for given name. """
        available_pipelines = self.pipeline_names
        assert pipeline_name in available_pipelines, f"Requested pipeline '{pipeline_name}', but only the following models are in the collection: '{available_pipelines}'."
        return self.pipelines[pipeline_name]

    def __call__(self, dataf: Union[pd.DataFrame, NumerFrame]) -> Dict[str, NumerFrame]:
        return self.process_all_pipelines(dataf=dataf)


# We introduce a different pipeline with no preprocessing or postprocessing. Only a `RandomModel`.

test_pipeline2 = ModelPipeline(models=[RandomModel()], pipeline_name="test_pipeline2")

# + [markdown] pycharm={"name": "#%% md\n"}
# We process two `ModelPipeline`s with different characteristics on the same data.
# -

collection = ModelPipelineCollection([test_pipeline, test_pipeline2])
assert collection.get_pipeline("test_pipeline2").pipeline_name == 'test_pipeline2'

result_datasets = collection(dataf=dataf)

# The `ModelPipelineCollection` returns a dictionary mapping pipeline names to `NumerFrame` objects, retaining all metadata and added prediction columns for each. Note that in this example, the 1st `NumerFrame` had a feature selection step, so it did not retain all columns. However, the second dataset retained all feature columns, because no preprocessing was done.

# + pycharm={"name": "#%%\n"}
result_datasets.keys()
# -

result_datasets['test_pipeline'].head(2)

# + pycharm={"name": "#%%\n"}
result_datasets['test_pipeline2'].head(2)
# -

# Since metadata is not manipulated in these pipelines, metadata should be the same as the original `NumerFrame` for all resulting `NumerFrame` objects.

# + pycharm={"name": "#%%\n"}
for _, result in result_datasets.items():
    assert dataf.meta == result.meta

# + pycharm={"name": "#%%\n"}
result_datasets['test_pipeline'].meta
# -

# -----------------------------------------------------------------------------

# +
# hide
# Run this cell to sync all changes with library
from nbdev.export import notebook2script

notebook2script()

# + pycharm={"name": "#%%\n"}


StarcoderdataJupyterScriptsDedupFiltered