#!/usr/bin/env python3
"""
Test: Machine Learning Integration with GRNexus Framework
Tests that ML algorithms work correctly when imported through main framework
"""

import sys

from grnexus import *

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

def test_knn_integration():
    print_header("TEST 1: KNN Integration with Framework")
    
    # Training data
    train_data = [
        [1.0, 1.0], [1.5, 1.2], [1.2, 1.5], [0.8, 1.1],
        [5.0, 5.0], [5.2, 4.8], [4.8, 5.2], [5.1, 5.1]
    ]
    train_labels = [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0]
    
    print("\nTraining KNN classifier...")
    knn = KNeighborsClassifier(n_neighbors=3)
    knn.fit(train_data, train_labels)
    print("✓ KNN trained successfully")
    
    # Test predictions
    test_points = [[1.1, 1.1], [5.0, 4.9]]
    predictions = knn.predict(test_points)
    
    print("\nPredictions:")
    for i, point in enumerate(test_points):
        print(f"  Point {point}: class {int(predictions[i])}")
    
    # Verify accuracy
    assert round(predictions[0]) == 0, "KNN prediction 1 incorrect"
    assert round(predictions[1]) == 1, "KNN prediction 2 incorrect"
    
    print("✓ KNN integration test passed")

def test_kmeans_integration():
    print_header("TEST 2: K-Means Integration with Framework")
    
    # Data with 2 clear clusters
    data = [
        [1.0, 1.0], [1.5, 1.2], [1.2, 1.5],
        [5.0, 5.0], [5.2, 4.8], [4.8, 5.2]
    ]
    
    print("\nTraining K-Means...")
    kmeans = KMeans(n_clusters=2, max_iter=100)
    kmeans.fit(data)
    print("✓ K-Means trained successfully")
    
    # Get cluster assignments
    labels = kmeans.predict(data)
    print(f"\nCluster assignments: {labels}")
    
    # Verify that similar points are in same cluster
    cluster_0 = labels[0]
    same_cluster = all(l == cluster_0 for l in labels[0:3])
    assert same_cluster, "K-Means clustering incorrect"
    
    print("✓ K-Means integration test passed")

def test_linear_regression_integration():
    print_header("TEST 3: Linear Regression Integration with Framework")
    
    # Simple linear data: y = 2x + 1
    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)
    print("✓ Linear Regression trained successfully")
    
    # Test prediction
    x_test = [[5.0]]
    prediction = lr.predict(x_test)[0]
    expected = 11.0
    
    print(f"\nPrediction for x=5: {prediction:.2f} (expected ~{expected})")
    
    # Verify prediction is close
    error = abs(prediction - expected)
    assert error < 1.0, f"Linear Regression prediction error too large: {error}"
    
    print("✓ Linear Regression integration test passed")

def test_logistic_regression_integration():
    print_header("TEST 4: Logistic Regression Integration with Framework")
    
    # Binary classification data
    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)
    print("✓ Logistic Regression trained successfully")
    
    # Test predictions
    x_test = [[1.5], [4.0]]
    predictions = logistic.predict(x_test)
    
    print("\nPredictions:")
    print(f"  x=1.5: class {predictions[0]} (expected 0)")
    print(f"  x=4.0: class {predictions[1]} (expected 1)")
    
    print("✓ Logistic Regression integration test passed")

def test_naive_bayes_integration():
    print_header("TEST 5: Naive Bayes Integration with Framework")
    
    # Training data
    train_data = [
        [1.0, 1.0], [1.5, 1.2], [1.2, 1.5],
        [5.0, 5.0], [5.2, 4.8], [4.8, 5.2]
    ]
    train_labels = [0.0, 0.0, 0.0, 1.0, 1.0, 1.0]
    
    print("\nTraining Gaussian Naive Bayes...")
    nb = GaussianNB()
    nb.fit(train_data, train_labels)
    print("✓ Naive Bayes trained successfully")
    
    # Test predictions
    test_points = [[1.0, 1.0], [5.0, 5.0]]
    predictions = nb.predict(test_points)
    
    print("\nPredictions:")
    print(f"  Point [1.0, 1.0]: class {predictions[0]} (expected 0)")
    print(f"  Point [5.0, 5.0]: class {predictions[1]} (expected 1)")
    
    # Verify predictions
    assert predictions[0] == 0, "Naive Bayes prediction 1 incorrect"
    assert predictions[1] == 1, "Naive Bayes prediction 2 incorrect"
    
    print("✓ Naive Bayes integration test passed")

def test_ml_with_neural_network():
    print_header("TEST 6: ML + Neural Network Hybrid")
    
    print("\nDemonstrating ML and DL can coexist in framework...")
    
    # Use KNN for feature selection/preprocessing
    train_data = [[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]]
    train_labels = [0.0, 1.0, 0.0]
    
    knn = KNeighborsClassifier(n_neighbors=2)
    knn.fit(train_data, train_labels)
    
    # Create a simple neural network
    nn = NeuralNetwork(learning_rate=0.01)
    nn.add(DenseLayer(units=4, input_dim=2, activation='ReLU'))
    nn.add(DenseLayer(units=1, input_dim=4, activation='Sigmoid'))
    
    print("✓ Both ML and DL models created successfully")
    print("✓ Framework supports hybrid ML/DL workflows")

def main():
    print("\n" + "=" * 70)
    print("  GRNexus Framework - Machine Learning Integration Tests (Python)")
    print("=" * 70)
    
    try:
        test_knn_integration()
        test_kmeans_integration()
        test_linear_regression_integration()
        test_logistic_regression_integration()
        test_naive_bayes_integration()
        test_ml_with_neural_network()
        
        print("\n" + "=" * 70)
        print("  ✅ ALL MACHINE LEARNING INTEGRATION 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())
