import time

from feature import feat_eng
from load_data import load_data
from model01 import model_workflow
from sklearn.model_selection import train_test_split


def log_execution_results(start_time, val_pred, test_pred, hypers, execution_label):
    """Log the results of a single model execution."""
    feedback_str = f"{execution_label} end.\n"
    feedback_str += f"Validation predictions shape: {val_pred.shape if val_pred is not None else 'None'}\n"
    feedback_str += f"Test predictions shape: {test_pred.shape if test_pred is not None else 'None'}\n"
    feedback_str += f"Hyperparameters: {hypers if hypers is not None else 'None'}\n"
    feedback_str += f"Execution time: {time.time() - start_time:.2f} seconds.\n"
    print(feedback_str)


# Load and preprocess data
X, y, test_X, test_ids = load_data()
X, y, test_X = feat_eng(X, y, test_X)
train_X, val_X, train_y, val_y = train_test_split(X, y, test_size=0.8, random_state=42)
print(f"train_X.shape: {train_X.shape}")
print(f"train_y.shape: {train_y.shape}" if not isinstance(train_y, list) else f"train_y(list)'s length: {len(train_y)}")
print(f"val_X.shape: {val_X.shape}")
print(f"val_y.shape: {val_y.shape}" if not isinstance(val_y, list) else f"val_y(list)'s length: {len(val_y)}")

# First execution
print("The first execution begins.\n")
start_time = time.time()
val_pred, test_pred, hypers = model_workflow(
    X=train_X,
    y=train_y,
    val_X=val_X,
    val_y=val_y,
    test_X=None,
)
log_execution_results(start_time, val_pred, test_pred, hypers, "The first execution")

# Second execution
print("The second execution begins.\n")
start_time = time.time()
val_pred, test_pred, final_hypers = model_workflow(
    X=train_X,
    y=train_y,
    val_X=None,
    val_y=None,
    test_X=test_X,
    hyper_params=hypers,
)
log_execution_results(start_time, val_pred, test_pred, final_hypers, "The second execution")

print("Model code test end.")
