#!/usr/bin/env python3
"""
Test: Basic ML Model Save/Load and Format Validation
Demonstrates saving models and format validation errors
"""

import sys

from grnexus import *

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

def test_save_and_load():
    print_header("TEST 1: Save and Load ML Model")
    
    # Create and train a simple model
    train_data = [[1.0, 2.0], [2.0, 3.0], [3.0, 4.0], [5.0, 6.0], [6.0, 7.0]]
    train_labels = [0.0, 0.0, 0.0, 1.0, 1.0]
    
    print("\nTraining KNN classifier...")
    knn = KNeighborsClassifier(n_neighbors=3)
    knn.fit(train_data, train_labels)
    
    print(f"Model info: {knn}")
    
    # Save model
    filepath = 'test_basic_model.lnexus'
    print(f"\nSaving model to {filepath}...")
    knn.save(filepath)
    
    # Load model
    print(f"Loading model from {filepath}...")
    knn_loaded = KNeighborsClassifier.load(filepath)
    print(f"Loaded model info: {knn_loaded}")
    
    # Test prediction
    test_point = [[2.5, 3.5]]
    pred_original = knn.predict(test_point)[0]
    pred_loaded = knn_loaded.predict(test_point)[0]
    
    print("\nPredictions:")
    print(f"  Original model: class {int(pred_original)}")
    print(f"  Loaded model:   class {int(pred_loaded)}")
    
    assert round(pred_original) == round(pred_loaded), "Predictions don't match!"
    
    # Cleanup
    if os.path.exists(filepath):
        os.remove(filepath)
    
    print("\n✓ Save/Load test passed!")

def test_format_validation():
    print_header("TEST 2: Format Validation - .nexus vs .lnexus")
    
    # Create a dummy .nexus file (Neural Network format)
    dummy_nexus = 'dummy_neural_network.nexus'
    with open(dummy_nexus, 'wb') as f:
        f.write(b'This is a neural network model')
    
    print("\nAttempting to load .nexus file (Neural Network format)...")
    print("Expected: Error message about invalid format")
    
    try:
        KNeighborsClassifier.load(dummy_nexus)
        print("\n❌ ERROR: Should have raised an exception!")
        raise AssertionError("Expected ValueError was not raised")
    except ValueError as e:
        print("\n✓ Correct error caught:")
        print(f"  {e}")
    
    # Cleanup
    if os.path.exists(dummy_nexus):
        os.remove(dummy_nexus)
    
    print("\n✓ Format validation test passed!")

def test_auto_extension_conversion():
    print_header("TEST 3: Auto Extension Conversion")
    
    train_data = [[1.0, 2.0], [3.0, 4.0]]
    train_labels = [0.0, 1.0]
    
    knn = KNeighborsClassifier(n_neighbors=1)
    knn.fit(train_data, train_labels)
    
    # Try to save with .nexus extension
    print("\nAttempting to save with .nexus extension...")
    knn.save('test_model.nexus')
    
    # Check that it was saved as .lnexus
    if os.path.exists('test_model.lnexus'):
        print("✓ File automatically saved as .lnexus")
        os.remove('test_model.lnexus')
    else:
        raise AssertionError("Auto conversion failed!")
    
    print("\n✓ Auto extension conversion test passed!")

def main():
    print("\n" + "=" * 70)
    print("  GRNexus ML - Basic Tests (Python)")
    print("=" * 70)
    
    try:
        test_save_and_load()
        test_format_validation()
        test_auto_extension_conversion()
        
        print("\n" + "=" * 70)
        print("  ✅ ALL BASIC 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())
