main.py

import sys
import os
import json
import numpy as np
import pandas as pd
from datetime import date
import matplotlib.pyplot as plt

# Import modules
from monk_model import MonkMLP
from cup_model import train_mlp3_flexible, predict_mlp3, save_cup_model_npz
from utils import (load_monk_data, load_cup_data, save_config, 
                   plot_monk_combined, plot_cup_curve, create_dirs)

# ================= MONK SECTION  =================
def run_monk():
    """
    Orchestrates the MONK Classification Tasks pipeline.
    1. Defines configurations for Monk 1, 2, and 3.
    2. Loads the dataset for each task.
    3. Trains the MonkMLP model.
    4. Plots MSE and Accuracy learning curves.
    5. Saves the final model and hyperparameters.
    """
    
    print("\n--- MONK Tasks ---")
    tasks = [
        ("monks-1", "monks-1", {"l2_lambda": 0.0, "lr": 0.1, "n_hidden": 3, "activation": "sigmoid", "init_type": "small_random"}),
        ("monks-2", "monks-2", {"l2_lambda": 0.0, "lr": 0.1, "n_hidden": 3, "activation": "sigmoid", "init_type": "small_random"}),
        ("monks-3", "monks-3-no-reg", {"l2_lambda": 0.0, "lr": 0.1, "n_hidden": 3, "activation": "sigmoid", "init_type": "small_random"}),
        ("monks-3", "monks-3-l2", {"l2_lambda": 0.1, "lr": 0.1, "n_hidden": 3, "activation": "sigmoid", "init_type": "small_random"})
    ]
    defaults = {"epochs": 400, "seed": 42}

    for ds_name, folder_name, specific_params in tasks:
        print(f"\nProcessing {folder_name}...")
        train_path = f"datasets/monk/{ds_name}.train"
        test_path = f"datasets/monk/{ds_name}.test"
        
        hyp_folder = os.path.join("monk_hyperparameters", folder_name)
        create_dirs([hyp_folder])
        
        run_config = defaults.copy()
        run_config.update(specific_params)
        save_config(run_config, hyp_folder, "config.json")

        if os.path.exists(train_path):
            X_train, y_train, X_test, y_test = load_monk_data(train_path, test_path)
            model = MonkMLP(
                n_inputs=17, n_hidden=run_config["n_hidden"], n_output=1, 
                activation=run_config["activation"], init_type=run_config["init_type"],
                seed=run_config["seed"]
            )
            hist = model.train(
                X_train, y_train, epochs=run_config["epochs"], 
                lr=run_config["lr"], l2=run_config["l2_lambda"], 
                X_test=X_test, y_test=y_test
            )
            plot_dir = f"plot/monk/{folder_name}"
            create_dirs([plot_dir])
            plot_monk_combined(hist, folder_name.upper(), f"{plot_dir}/best-{folder_name}.png")
            
            final_dir = "final_model"
            create_dirs([final_dir])
            save_config(run_config, final_dir, f"best_{folder_name}_hyperparameters.json")
            model.save_npz(f"{final_dir}/best_{folder_name}_model.npz")
            print(f"Finished {folder_name}. Final Test Acc: {hist['acc_test'][-1]:.4f}")
        else:
            print(f"Dataset not found: {train_path}")

# ================= CUP SECTION =================
def run_cup():
    """
    Orchestrates the ML-CUP Regression pipeline.
    1. Loads the CUP training and test datasets.
    2. Splits the training data into Train (70%), Validation (15%), and Internal Test (15%) sets.
    3. Defines the top 3 hyperparameter configurations.
    4. Iterates through the top 3 configs, training each and tracking the best one based on Validation MEE.
    5. Plots learning curves for all 3 configs and a specific curve for the best model.
    6. Retrains the best configuration on the 'Development Set' (Train + Val) and evaluates on Internal Test.
    7. Retrains the best configuration on ALL available data (Train + Val + Int Test) for the final Blind Test.
    8. Generates predictions for the Blind Test and saves them to a CSV file.
    """
    
    print("\n--- ML-CUP25 Pipeline ---")
    TR_PATH = "datasets/cup/ML-CUP25-TR.csv"
    TS_PATH = "datasets/cup/ML-CUP25-TS.csv"
    
    if not os.path.exists(TR_PATH):
        print(f"Error: {TR_PATH} not found.")
        return

    # 1. Load Data
    X, Y, X_blind, ids_blind = load_cup_data(TR_PATH, TS_PATH)

    # 2. EXACT NOTEBOOK SPLITTING LOGIC
    np.random.seed(42)
    indices = np.random.permutation(X.shape[0])

    X = X[indices]
    Y = Y[indices]

    N = X.shape[0]
    N_train = int(0.70 * N)
    N_val   = int(0.15 * N)

    X_train = X[:N_train]
    Y_train = Y[:N_train]

    X_val = X[N_train:N_train+N_val]
    Y_val = Y[N_train:N_train+N_val]

    X_test_int = X[N_train + N_val:]
    Y_test_int = Y[N_train + N_val:]
    
    print(f"Splits -> Train: {len(X_train)}, Val: {len(X_val)}, Int Test: {len(X_test_int)}")

    # 3. Top 3 Configs
    top_3_configs = [
        {"h1": 128, "h2": 64, "h3": 32, "lr": 1e-3, "momentum": 0.8, "l2": 1e-4},
        {"h1": 128, "h2": 64, "h3": 32, "lr": 5e-4, "momentum": 0.9, "l2": 1e-5},
        {"h1": 256, "h2": 128, "h3": 64, "lr": 1e-3, "momentum": 0.8, "l2": 0.001}
    ]
    
    top3_folder = "cup_hyperparameters/top_3"
    create_dirs([top3_folder])
    save_config(top_3_configs, top3_folder, "top_3_configs.json")
    
    print("\n--- Evaluating Top 3 Configs ---")
    best_mee = float('inf')
    best_config = None
    best_hist = None # Variable to store the history of the winning model
    
    for i, cfg in enumerate(top_3_configs):
        print(f"Config {i+1}: {cfg}")
        # Train
        _, _, val_mee, _, hist = train_mlp3_flexible(
            X_train, Y_train, X_val, Y_val,
            **cfg, epochs=4000, patience=200, verbose_n=500, plot=False
        )
        print(f"  -> Validation MEE: {val_mee:.4f}")
        
        # Plot individual config
        plot_folder = "plot/cup/top_3_comparison"
        create_dirs([plot_folder])
        plot_cup_curve(hist, 'train_mee', 'val_mee', f"Config {i+1} Train vs Val MEE", f"{plot_folder}/config_{i+1}_curve.png")
        
        # Check for best
        if val_mee < best_mee:
            best_mee = val_mee
            best_config = cfg
            best_hist = hist # Capture the best history

    print(f"\nBest Config Found: {best_config} (Val MEE: {best_mee:.4f})")
    
    # --- NEW: Plot Best Model Specifically ---
    if best_hist is not None:
        print("Displaying best model plot...")
        plot_cup_curve(
            best_hist, 
            'train_mee', 
            'val_mee', 
            "Loss (MEE) vs Epochs for Cup best model", 
            "plot/cup/best_model_selection_curve.png"
        )
    
    # 4. Internal Test
    print("\n--- Retraining on DEV (Train + Val) for Internal Test ---")
    X_dev = np.vstack((X_train, X_val))
    Y_dev = np.vstack((Y_train, Y_val))
    
    model, stats, _, _, _ = train_mlp3_flexible(
        X_dev, Y_dev, X_dev, Y_dev,
        **best_config, epochs=4000, patience=10**9, verbose_n=500, plot=False
    )
    
    Y_pred_int = predict_mlp3(model, stats, X_test_int)
    final_int_mee = float(np.mean(np.linalg.norm(Y_test_int - Y_pred_int, axis=1)))
    print(f"Final Internal Test MEE: {final_int_mee:.4f}")
    
    param_folder = "cup_parameters"
    create_dirs([param_folder])
    save_config({"best_config": best_config, "internal_test_mee": final_int_mee}, param_folder, "cup_internal_test_metrics.json")

    # 5. Blind Test
    print("\n--- Final Retraining on ALL Data & Blind Prediction ---")
    X_all_shuffled = X 
    Y_all_shuffled = Y
    
    final_model_blind, final_stats_blind, _, _, _ = train_mlp3_flexible(
        X_all_shuffled, Y_all_shuffled, X_all_shuffled, Y_all_shuffled, 
        **best_config, epochs=4000, patience=10**9, verbose_n=500, plot=False
    )
    
    final_folder = "final_model"
    create_dirs([final_folder])
    save_cup_model_npz(f"{final_folder}/cup_best_model.npz", final_model_blind, final_stats_blind, best_config)
    save_config(best_config, final_folder, "cup_best_hyperparameters.json")
    
    Y_blind_pred = predict_mlp3(final_model_blind, final_stats_blind, X_blind)
    
    res_path = f"results/Spice-Devs_ML-CUP25-TS.csv"
    create_dirs(["results"])
    with open(res_path, 'w') as f:
        f.write("# Israel Fitsum Yohannes, Shaikh Asif Hossain\n# Spice-Devs\n# ML-CUP25 v1\n# 21 Jan 2026\n")
        for i, pred in zip(ids_blind, Y_blind_pred):
            f.write(f"{i},{pred[0]},{pred[1]},{pred[2]},{pred[3]}\n")
            
    print(f"Blind Test Results saved to {res_path}")

if __name__ == "__main__":
    """
    Main entry point of the script.
    Provides a CLI (Command Line Interface) menu allowing the user to choose 
    between running MONK tasks, the ML-CUP task, or exiting.
    """
    
    while True:
        print("\n" + "="*40); print("      ML-2025 PROJECT CLI MENU      "); print("="*40)
        print("1. Monks Tasks"); print("2. ML-CUP Task"); print("3. Exit"); print("="*40)
        choice = input("Enter choice (1-3): ").strip()
        if choice == '1':
            try: run_monk()
            except Exception as e: print(f"Error: {e}")
        elif choice == '2':
            try: run_cup()
            except Exception as e: print(f"Error: {e}")
        elif choice == '3': sys.exit()