#!/usr/bin/env python3
"""
Test: Save/Load ML Models
Tests that ML models can be saved and loaded correctly
"""

import sys

from grnexus import *
import numpy as np

def print_header(title):
    print("\n" + "=" * 70)
    print(f"  {title}")
    print("=" * 70)

def test_knn_save_load():
    print_header("TEST 1: KNN Save/Load")
    
    # Train model
    train_data = [[1.0, 1.0], [1.5, 1.2], [5.0, 5.0], [5.2, 4.8]]
    train_labels = [0.0, 0.0, 1.0, 1.0]
    
    print("\nTraining KNN...")
    knn = KNeighborsClassifier(n_neighbors=3)
    knn.fit(train_data, train_labels)
    
    # Make prediction before save
    test_point = [[1.1, 1.1]]
    pred_before = knn.predict(test_point)[0]
    print(f"Prediction before save: {int(pred_before)}")
    
    # Save model
    filepath = 'test_knn_model.lnexus'
    knn.save(filepath)
    print(f"✓ Model saved to {filepath}")
    
    # Load model
    knn_loaded = KNeighborsClassifier.load(filepath)
    print(f"✓ Model loaded from {filepath}")
    
    # Make prediction after load
    pred_after = knn_loaded.predict(test_point)[0]
    print(f"Prediction after load: {int(pred_after)}")
    
    # Verify predictions match
    assert round(pred_before) == round(pred_after), "Predictions don't match!"
    
    # Cleanup
    if os.path.exists(filepath):
        os.remove(filepath)
    
    print("✓ KNN save/load test passed")

def test_linear_regression_save_load():
    print_header("TEST 2: Linear Regression Save/Load")
    
    # Train model
    x_train = [[1.0], [2.0], [3.0], [4.0]]
    y_train = [3.0, 5.0, 7.0, 9.0]
    
    print("\nTraining Linear Regression...")
    lr = LinearRegression()
    lr.fit(x_train, y_train)
    
    # Make prediction before save
    test_x = [[5.0]]
    pred_before = lr.predict(test_x)[0]
    print(f"Prediction before save: {pred_before:.2f}")
    
    # Save model
    filepath = 'test_lr_model.lnexus'
    lr.save(filepath)
    print(f"✓ Model saved to {filepath}")
    
    # Load model
    lr_loaded = LinearRegression.load(filepath)
    print(f"✓ Model loaded from {filepath}")
    
    # Make prediction after load
    pred_after = lr_loaded.predict(test_x)[0]
    print(f"Prediction after load: {pred_after:.2f}")
    
    # Verify predictions match
    error = abs(pred_before - pred_after)
    assert error < 0.001, f"Predictions don't match! Error: {error}"
    
    # Cleanup
    if os.path.exists(filepath):
        os.remove(filepath)
    
    print("✓ Linear Regression save/load test passed")

def test_kmeans_save_load():
    print_header("TEST 3: K-Means Save/Load")
    
    # Train model
    data = [[1.0, 1.0], [1.5, 1.2], [5.0, 5.0], [5.2, 4.8]]
    
    print("\nTraining K-Means...")
    kmeans = KMeans(n_clusters=2, max_iter=100)
    kmeans.fit(data)
    
    # Get centroids before save
    centroids_before = kmeans.centroids.copy()
    print(f"Centroids before save:\n{centroids_before}")
    
    # Save model
    filepath = 'test_kmeans_model.lnexus'
    kmeans.save(filepath)
    print(f"✓ Model saved to {filepath}")
    
    # Load model
    kmeans_loaded = KMeans.load(filepath)
    print(f"✓ Model loaded from {filepath}")
    
    # Get centroids after load
    centroids_after = kmeans_loaded.centroids
    print(f"Centroids after load:\n{centroids_after}")
    
    # Verify centroids match
    error = np.abs(centroids_before - centroids_after).max()
    assert error < 0.001, f"Centroids don't match! Max error: {error}"
    
    # Cleanup
    if os.path.exists(filepath):
        os.remove(filepath)
    
    print("✓ K-Means save/load test passed")

def test_naive_bayes_save_load():
    print_header("TEST 4: Naive Bayes Save/Load")
    
    # Train model
    train_data = [[1.0, 1.0], [1.5, 1.2], [5.0, 5.0], [5.2, 4.8]]
    train_labels = [0.0, 0.0, 1.0, 1.0]
    
    print("\nTraining Naive Bayes...")
    nb = GaussianNB()
    nb.fit(train_data, train_labels)
    
    # Make prediction before save
    test_point = [[1.0, 1.0]]
    pred_before = nb.predict(test_point)[0]
    print(f"Prediction before save: {pred_before}")
    
    # Save model
    filepath = 'test_nb_model.lnexus'
    nb.save(filepath)
    print(f"✓ Model saved to {filepath}")
    
    # Load model
    nb_loaded = GaussianNB.load(filepath)
    print(f"✓ Model loaded from {filepath}")
    
    # Make prediction after load
    pred_after = nb_loaded.predict(test_point)[0]
    print(f"Prediction after load: {pred_after}")
    
    # Verify predictions match
    assert pred_before == pred_after, "Predictions don't match!"
    
    # Cleanup
    if os.path.exists(filepath):
        os.remove(filepath)
    
    print("✓ Naive Bayes save/load test passed")

def test_logistic_regression_save_load():
    print_header("TEST 5: Logistic Regression Save/Load")
    
    # Train model
    x_train = [[1.0], [1.5], [2.0], [3.5], [4.0], [4.5]]
    y_train = [0.0, 0.0, 0.0, 1.0, 1.0, 1.0]
    
    print("\nTraining Logistic Regression...")
    logistic = LogisticRegression(learning_rate=0.1, max_iter=1000)
    logistic.fit(x_train, y_train)
    
    # Make prediction before save
    test_x = [[1.5]]
    pred_before = logistic.predict(test_x)[0]
    print(f"Prediction before save: {pred_before}")
    
    # Save model
    filepath = 'test_logistic_model.lnexus'
    logistic.save(filepath)
    print(f"✓ Model saved to {filepath}")
    
    # Load model
    logistic_loaded = LogisticRegression.load(filepath)
    print(f"✓ Model loaded from {filepath}")
    
    # Make prediction after load
    pred_after = logistic_loaded.predict(test_x)[0]
    print(f"Prediction after load: {pred_after}")
    
    # Verify predictions match
    error = abs(pred_before - pred_after)
    assert error < 0.001, f"Predictions don't match! Error: {error}"
    
    # Cleanup
    if os.path.exists(filepath):
        os.remove(filepath)
    
    print("✓ Logistic Regression save/load test passed")

def main():
    print("\n" + "=" * 70)
    print("  GRNexus ML Models - Save/Load Tests (Python)")
    print("=" * 70)
    
    try:
        test_knn_save_load()
        test_linear_regression_save_load()
        test_kmeans_save_load()
        test_naive_bayes_save_load()
        test_logistic_regression_save_load()
        
        print("\n" + "=" * 70)
        print("  ✅ ALL ML SAVE/LOAD TESTS PASSED!")
        print("=" * 70 + "\n")
        
        return 0
        
    except Exception as e:
        print(f"\n❌ TEST FAILED: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())
