import pandas as pd
import joblib
from sklearn.pipeline import Pipeline

from lr_customer_value.config import config
from lr_customer_value import __version__ as _version

import logging


_logger = logging.getLogger(__name__)


def load_dataset(*, files_list: str) -> pd.DataFrame:
    data = pd.DataFrame()
    for file_name in files_list:
        tmp = pd.read_csv(f"{file_name}")
        if data.empty:
            data = tmp
        else:
            data = data.merge(right=tmp, on='people_id', how='inner')

    return data


def save_pipeline(*, pipeline_to_persist) -> None:
    """
    Persist the pipeline.
    Saves the versioned model, and overwrites any previous
    saved models. So when a model is published, there is only
    one trained model that can be called, and we know exactly
    how it was built.
    """

    save_file_name = f"{config.PIPELINE_SAVE_FILE}_{_version}.pkl"
    save_path = config.TRAINED_MODEL_DIR / save_file_name

    remove_old_pipelines(files_to_keep=save_file_name)
    joblib.dump(pipeline_to_persist, save_path)

    _logger.info(f"saved pipeline: {save_file_name}")


def load_pipeline(*, file_name: str) -> Pipeline:
    """Load a persisted pipeline."""

    file_path = config.TRAINED_MODEL_DIR / file_name
    saved_pipeline = joblib.load(filename=file_path)

    # _logger.info(f"loaded the pipeline: {saved_pipeline}")

    return saved_pipeline


def remove_old_pipelines(*, files_to_keep) -> None:
    """
    Only keeps the most recent version of the
    model we used in the training. It ensures there is a
    simple one-to-one mapping between the package version
    and the model version to be imported and used by other
    applications.
    """

    for model_file in config.TRAINED_MODEL_DIR.iterdir():
        if model_file.name not in [files_to_keep, "__init__.py"]:
            model_file.unlink()

