import argparse
import os
from data_loader import load_and_process_data, CATEGORICAL_COLUMNS
from model_trainer import train_models
from model_manager import save_models, load_models
from model_predictor import predict
from config import MODEL_DIR
## ===========================
#  MAIN FUNCTION
# ===========================

def main(train=True, retrain=False):
    """ Main entry point to train, retrain or predict """
    # Create model directory if it doesn't exist
    if not os.path.exists(MODEL_DIR):
        os.makedirs(MODEL_DIR)
    print("\n🚀 Loading data...")
    X_train, X_val, y_train, y_val, test_df = load_and_process_data()

    if train or retrain:
        print("\n🚀 Training models...")
        models = train_models(X_train, y_train, CATEGORICAL_COLUMNS)
        save_models(models)

    else:
        print("\n🚀 Loading existing models...")
        models = load_models()

    print("\n🔍 Making predictions...")
    predictions = predict(models, test_df)

    # Save final predictions
    predictions.to_csv("final_predictions.csv", index=False)
    print("\n✅ Predictions saved successfully as 'final_predictions.csv'!")

# ===========================
#  COMMAND-LINE EXECUTION
# ===========================
if __name__ == "__main__":
    # parser = argparse.ArgumentParser(description="Train, retrain or make predictions")
    # parser.add_argument("--train", action="store_true", help="Train new models")
    # parser.add_argument("--retrain", action="store_true", help="Retrain models with updated data")
    #
    # args = parser.parse_args()
    main(train=False, retrain=False)