"""
Notes
- it takes ~143 lines to define the steps including repetitive boilerplate for
  decorators, imports, and self.next calls. (only 120 if you take away docstrings).
  Whereas the functional interface takes 60 lines
- it is difficult to parse the dependency relationships between steps. I.e. 
  1. in what order do they run including joins
  2. for a given step, which steps consume its output, and which outputs?
"""


import random
from typing import Any, Dict, List, Optional

from metaflow import FlowSpec, Parameter, step, resources, current

class E2EUseCaseTrainingFlow(FlowSpec):
    """
    Model training flow.

    This is a flow that loads the data, processes it and splits
    it into train and test sets, then searches for best hyperparameters,
    trains and evaluates a model.
    """

    # Flow parameters
    model_search_space = Parameter(
        "model-search-space",
        help="Search space for hyperparameter tuning",
        default=str({
            "sgd": {
                "model_package": "sklearn.linear_model",
                "model_class": "SGDClassifier",
                "search_grid": {
                    "alpha": [0.01, 0.1, 1.0],
                    "max_iter": [1000, 2000]
                }
            },
            "rf": {
                "model_package": "sklearn.ensemble", 
                "model_class": "RandomForestClassifier",
                "search_grid": {
                    "n_estimators": [50, 100, 200],
                    "max_depth": [3, 5, 10, None]
                }
            }
        })
    )

    target_env = Parameter(
        "target-env",
        help="The environment to promote the model to",
        default="staging"
    )

    test_size = Parameter(
        "test-size",
        help="Size of holdout set for training 0.0..1.0",
        default=0.2
    )

    drop_na = Parameter(
        "drop-na", 
        help="If True NA values will be removed from dataset",
        default=None,
        type=bool,
        required=False
    )

    normalize = Parameter(
        "normalize",
        help="If True dataset will be normalized with MinMaxScaler", 
        default=None,
        type=bool,
        required=False
    )

    drop_columns = Parameter(
        "drop-columns",
        help="List of columns to drop from dataset",
        default=None,
        required=False
    )

    min_train_accuracy = Parameter(
        "min-train-accuracy",
        help="Threshold to stop execution if train set accuracy is lower",
        default=0.0
    )

    min_test_accuracy = Parameter(
        "min-test-accuracy", 
        help="Threshold to stop execution if test set accuracy is lower",
        default=0.0
    )

    fail_on_accuracy_quality_gates = Parameter(
        "fail-on-accuracy-quality-gates",
        help="If True and accuracy thresholds are not met - execution will be interrupted early",
        default=False
    )

    @step
    def start(self):
        """Initialize the flow."""
        self.next(self.data_loader)

    @resources(memory=2000, cpu=1)
    @step
    def data_loader(self):
        """Load the raw data."""
        from helpers import data_loader
        self.raw_data, self.target, self.random_state = data_loader(
            random_state=random.randint(0, 100)
        )
        self.next(self.train_data_splitter)

    @resources(memory=2000, cpu=1) 
    @step
    def train_data_splitter(self):
        """Split data into train and test sets."""
        from helpers import train_data_splitter
        self.dataset_trn, self.dataset_tst = train_data_splitter(
            dataset=self.raw_data,
            test_size=self.test_size,
        )
        self.next(self.train_data_preprocessor)

    @resources(memory=4000, cpu=2)
    @step
    def train_data_preprocessor(self):
        """Preprocess the training and test data."""
        from helpers import train_data_preprocessor
        self.dataset_trn, self.dataset_tst, self.preprocess_pipeline = train_data_preprocessor(
            dataset_trn=self.dataset_trn,
            dataset_tst=self.dataset_tst,
            drop_na=self.drop_na,
            normalize=self.normalize,
            drop_columns=self.drop_columns,
        )
        self.next(self.hp_tuning_fan_out)

    @step
    def hp_tuning_fan_out(self):
        """Fan out for hyperparameter tuning - create one branch per model configuration."""
        self.hp_configs = []
        for config_name, model_search_configuration in self.model_search_space.items():
            self.hp_configs.append({
                "config_name": config_name,
                "model_package": model_search_configuration["model_package"],
                "model_class": model_search_configuration["model_class"],
                "search_grid": model_search_configuration["search_grid"]
            })
        
        self.next(self.hp_tuning_single_search, foreach="hp_configs")

    @resources(memory=8000, cpu=4)
    @step 
    def hp_tuning_single_search(self):
        """Perform hyperparameter tuning for a single model configuration."""
        from helpers import hp_tuning_single_search
        config = self.input
        
        self.best_model, self.hp_score = hp_tuning_single_search(
            model_package=config["model_package"],
            model_class=config["model_class"], 
            search_grid=config["search_grid"],
            dataset_trn=self.dataset_trn,
            dataset_tst=self.dataset_tst,
            target=self.target,
        )
        
        self.next(self.hp_tuning_join)

    @step
    def hp_tuning_join(self, inputs):
        """Join hyperparameter tuning results and select the best model."""
        from helpers import hp_tuning_select_best_model
        # Collect all models and their scores
        models_and_scores = []
        for inp in inputs:
            models_and_scores.append((inp.best_model, inp.hp_score))
        
        self.best_model = hp_tuning_select_best_model(models_and_scores)
        
        # Keep the other attributes from one of the inputs (they should be the same)
        self.dataset_trn = inputs[0].dataset_trn
        self.dataset_tst = inputs[0].dataset_tst
        self.target = inputs[0].target
        self.preprocess_pipeline = inputs[0].preprocess_pipeline
        
        self.next(self.model_trainer)

    @resources(memory=4000, cpu=2)
    @step
    def model_trainer(self):
        """Train the best model on the full training set."""
        from helpers import model_trainer
        self.trained_model = model_trainer(
            dataset_trn=self.dataset_trn,
            model=self.best_model,
            target=self.target,
        )
        self.next(self.model_evaluator)

    @resources(memory=2000, cpu=1)
    @step
    def model_evaluator(self):
        """Evaluate the trained model."""
        from helpers import model_evaluator
        self.train_accuracy, self.test_accuracy = model_evaluator(
            model=self.trained_model,
            dataset_trn=self.dataset_trn,
            dataset_tst=self.dataset_tst,
            min_train_accuracy=self.min_train_accuracy,
            min_test_accuracy=self.min_test_accuracy,
            fail_on_accuracy_quality_gates=self.fail_on_accuracy_quality_gates,
            target=self.target,
        )
        self.next(self.compute_performance_metrics)

    @resources(memory=2000, cpu=1)
    @step
    def compute_performance_metrics(self):
        """Compute performance metrics for model comparison."""
        from helpers import compute_performance_metrics_on_current_data
        self.latest_metric, self.current_metric = compute_performance_metrics_on_current_data(
            dataset_tst=self.dataset_tst,
            target_env=self.target_env,
        )
        self.next(self.promote_with_metric_compare)

    @resources(memory=1000, cpu=1)
    @step
    def promote_with_metric_compare(self):
        """Decide whether to promote the model based on metric comparison."""
        from helpers import promote_with_metric_compare
        self.promoted = promote_with_metric_compare(
            latest_metric=self.latest_metric,
            current_metric=self.current_metric,
            target_env=self.target_env,
        )
        
        self.next(self.notify_on_success)

    @resources(memory=500, cpu=1)
    @step
    def notify_on_success(self):
        """Send success notification."""
        from helpers import notify_on_success
        notify_on_success(f"E2E training flow completed successfully! Model promoted: {self.promoted}")
        self.next(self.end)

    @step
    def end(self):
        """End the flow."""


if __name__ == "__main__":
    E2EUseCaseTrainingFlow()