import logging

import mlflow
import pandas as pd
import typer
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 = constants.ModellingStrategy.MIXED,
):
    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 Ridge with strategy '{strategy}' and features {features}")

    for alpha in constants.RIDGE_ALPHAS:
        with mlflow.start_run(run_name=f"modelling_{strategy.value}_ridge_alpha_{alpha:g}"):
            mlflow.log_param("model_type", "ridge")
            mlflow.log_param("strategy", strategy.value)
            mlflow.log_param("split_strategy", constants.CHOSEN_SPLIT_STRATEGY.value)

            mlflow.log_param("features", ",".join(features))
            mlflow.log_param("alpha", alpha)
            model = linear_model.Ridge(alpha=alpha)
            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))


if __name__ == "__main__":
    app()