import logging

import mlflow
import pandas as pd
import typer
from mlflow.models import infer_signature
from sklearn import linear_model
from sklearn import metrics

from .. import constants

app = typer.Typer()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@app.command()
def main(
    strategy: constants.ModellingStrategy,
    register: bool = typer.Option(
        False,
        "--register/--no-register",
        help="Enregistrer le modele dans le Model Registry (cree une nouvelle version).",
    ),
):
    training_file_path = constants.DATASET_DIR / "train.parquet"
    validation_file_path = constants.DATASET_DIR / "validation.parquet"

    features = constants.MODELLING_FEATURES[strategy]

    train_df = pd.read_parquet(training_file_path)
    validation_df = pd.read_parquet(validation_file_path)

    train_df = train_df.dropna(
        subset=features + [constants.TARGET]
    )
    validation_df = validation_df.dropna(
        subset=features + [constants.TARGET]
    )

    X_train = train_df[features]
    y_train = train_df[constants.TARGET]

    X_validation = validation_df[features]
    y_validation = validation_df[constants.TARGET]

    logger.info(f"Training model with strategy '{strategy}' and features {features}")

    with mlflow.start_run(run_name=f"modelling_{strategy.value}"):
        mlflow.log_param("model_type", "linear")
        mlflow.log_param("strategy", strategy.value)
        mlflow.log_param("split_strategy", constants.CHOSEN_SPLIT_STRATEGY.value)

        mlflow.log_param("features", ",".join(features))
        model = linear_model.LinearRegression()
        model.fit(X_train, y_train)

        train_predictions = model.predict(X_train)
        validation_predictions = model.predict(X_validation)

        train_rmse = metrics.root_mean_squared_error(y_train, train_predictions)
        validation_rmse = metrics.root_mean_squared_error(y_validation, validation_predictions)

        train_mae = metrics.mean_absolute_error(y_train, train_predictions)
        validation_mae = metrics.mean_absolute_error(y_validation, validation_predictions)

        mlflow.log_metric("train_rmse", train_rmse)
        mlflow.log_metric("validation_rmse", validation_rmse)

        mlflow.log_metric("train_mae", train_mae)
        mlflow.log_metric("validation_mae", validation_mae)

        for feature_name, coefficient in zip(
            features,
            model.coef_,
            strict=True,
        ):
            mlflow.log_metric(f"coef_{feature_name}", float(coefficient))

        # Sauvegarde de l'artefact du modele (poids + signature + environnement
        # d'execution : requirements.txt, conda.yaml, MLmodel). --register empile
        # une nouvelle version dans le Model Registry pour les meilleures experiences.
        signature = infer_signature(X_train, train_predictions)
        mlflow.sklearn.log_model(
            sk_model=model,
            name="model",
            signature=signature,
            input_example=X_train.iloc[:5],
            registered_model_name=(
                constants.REGISTERED_MODEL_NAME if register else None
            ),
        )
        if register:
            logger.info(
                f"Model registered as '{constants.REGISTERED_MODEL_NAME}' (nouvelle version)"
            )


if __name__ == "__main__":
    app()
