# -*- coding: utf-8 -*-
import dataiku
import pandas as pd, numpy as np
from dataiku import pandasutils as pdu
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score
from sklearn.metrics import mean_absolute_error

# Read recipe inputs: Load the scored test dataset from Dataiku
test_ts_scored = dataiku.Dataset("test_ts_scored")
test_ts_scored_df = test_ts_scored.get_dataframe()

# Define the ground truth and predicted values for evaluation
y_true = test_ts_scored_df[["output1", "output2"]]
y_pred = test_ts_scored_df[["mlflow_out_0", "mlflow_out_1"]]

# Compute the evaluation metrics: R2, MSE, and MAE
evaluation_metrics_df = pd.DataFrame({
    "R2": [r2_score(y_true, y_pred)], 
    "MSE": [mean_squared_error(y_true, y_pred)],
    "MAE": [mean_absolute_error(y_true, y_pred)]
})

# Write recipe outputs: Save the evaluation metrics to a Dataiku dataset
evaluation_metrics = dataiku.Dataset("evaluation_metrics")
evaluation_metrics.write_with_schema(evaluation_metrics_df)
