#!/usr/bin/env python3
"""
Test: Practical ML Applications
Real-world use cases: Iris Classification and Customer Segmentation
"""

import sys

from grnexus import *

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

def test_iris_classification():
    print_header("PRACTICAL TEST 1: Iris Flower Classification")
    
    # Iris dataset (simplified)
    # Features: [sepal_length, sepal_width, petal_length, petal_width]
    # Classes: 0=Setosa, 1=Versicolor, 2=Virginica
    
    train_data = [
        # Setosa
        [5.1, 3.5, 1.4, 0.2], [4.9, 3.0, 1.4, 0.2], [4.7, 3.2, 1.3, 0.2],
        [4.6, 3.1, 1.5, 0.2], [5.0, 3.6, 1.4, 0.2],
        # Versicolor
        [7.0, 3.2, 4.7, 1.4], [6.4, 3.2, 4.5, 1.5], [6.9, 3.1, 4.9, 1.5],
        [5.5, 2.3, 4.0, 1.3], [6.5, 2.8, 4.6, 1.5],
        # Virginica
        [6.3, 3.3, 6.0, 2.5], [5.8, 2.7, 5.1, 1.9], [7.1, 3.0, 5.9, 2.1],
        [6.3, 2.9, 5.6, 1.8], [6.5, 3.0, 5.8, 2.2]
    ]
    
    train_labels = [
        0.0, 0.0, 0.0, 0.0, 0.0,
        1.0, 1.0, 1.0, 1.0, 1.0,
        2.0, 2.0, 2.0, 2.0, 2.0
    ]
    
    print("\nTraining Iris classifier with Naive Bayes...")
    classifier = GaussianNB()
    classifier.fit(train_data, train_labels)
    
    print(f"Model trained: {classifier}")
    
    # Save model
    model_file = 'iris_classifier.lnexus'
    classifier.save(model_file)
    print(f"\n✓ Model saved to {model_file}")
    
    # Test predictions
    test_samples = [
        {'data': [5.0, 3.5, 1.3, 0.3], 'expected': 'Setosa'},
        {'data': [6.5, 3.0, 4.5, 1.5], 'expected': 'Versicolor'},
        {'data': [6.5, 3.0, 5.5, 2.0], 'expected': 'Virginica'}
    ]
    
    species = ['Setosa', 'Versicolor', 'Virginica']
    
    print("\nTesting predictions:")
    for sample in test_samples:
        pred = classifier.predict([sample['data']])[0]
        predicted_species = species[pred]
        status = '✓' if predicted_species == sample['expected'] else '✗'
        print(f"  {status} Features: {sample['data'][0]:.1f}, {sample['data'][1]:.1f}... → {predicted_species} (expected: {sample['expected']})")
    
    # Cleanup
    if os.path.exists(model_file):
        os.remove(model_file)
    
    print("\n✓ Iris classification test completed!")

def test_customer_segmentation():
    print_header("PRACTICAL TEST 2: Customer Segmentation with K-Means")
    
    # Customer data: [age, annual_income_k, spending_score]
    customers = [
        # Low income, low spending
        [25, 15, 39], [28, 15, 81], [23, 16, 6], [27, 16, 77],
        # Medium income, medium spending
        [35, 40, 50], [38, 42, 52], [40, 43, 48], [42, 44, 55],
        # High income, high spending
        [45, 70, 88], [48, 72, 90], [50, 75, 85], [52, 78, 92],
        # High income, low spending
        [55, 80, 20], [58, 82, 18], [60, 85, 15], [62, 88, 22]
    ]
    
    print("\nSegmenting customers into 4 groups using K-Means...")
    kmeans = KMeans(n_clusters=4, max_iter=100)
    kmeans.fit(customers)
    
    print(f"Model trained: {kmeans}")
    
    # Save model
    model_file = 'customer_segments.lnexus'
    kmeans.save(model_file)
    print(f"\n✓ Model saved to {model_file}")
    
    # Analyze segments
    labels = kmeans.predict(customers)
    centroids = kmeans.centroids
    
    print("\nSegment Analysis:")
    for segment in range(4):
        centroid = centroids[segment]
        count = sum(1 for l in labels if l == segment)
        print(f"\n  Segment {segment + 1}: {count} customers")
        print(f"    Avg Age: {centroid[0]:.1f} years")
        print(f"    Avg Income: ${centroid[1]:.1f}k")
        print(f"    Avg Spending Score: {centroid[2]:.1f}/100")
    
    # Predict segment for new customer
    new_customer = [[30, 50, 60]]
    segment = kmeans.predict(new_customer)[0]
    print(f"\n  New customer [30 years, $50k, score 60] → Segment {segment + 1}")
    
    # Cleanup
    if os.path.exists(model_file):
        os.remove(model_file)
    
    print("\n✓ Customer segmentation test completed!")

def test_house_price_prediction():
    print_header("PRACTICAL TEST 3: House Price Prediction")
    
    # Simplified house data: [size_sqft / 1000] (normalized)
    # Price in hundreds of thousands
    houses = [
        [1.0], [1.2], [1.5], [1.8], [2.0], [2.2], [2.5], [2.8]
    ]
    
    prices = [1.5, 1.8, 2.2, 2.8, 3.2, 3.6, 4.2, 4.8]
    
    print("\nTraining house price predictor with Linear Regression...")
    model = LinearRegression()
    model.fit(houses, prices)
    
    print(f"Model trained: {model}")
    
    # Save model
    model_file = 'house_price_model.lnexus'
    model.save(model_file)
    print(f"\n✓ Model saved to {model_file}")
    
    # Test predictions
    test_houses = [
        {'data': [1.1], 'desc': '1100 sqft'},
        {'data': [1.6], 'desc': '1600 sqft'},
        {'data': [2.3], 'desc': '2300 sqft'}
    ]
    
    print("\nPrice predictions:")
    for house in test_houses:
        pred_price = model.predict([house['data']])[0]
        print(f"  {house['desc']} → ${pred_price * 100:.0f}k")
    
    # Cleanup
    if os.path.exists(model_file):
        os.remove(model_file)
    
    print("\n✓ House price prediction test completed!")

def main():
    print("\n" + "=" * 70)
    print("  GRNexus ML - Practical Applications (Python)")
    print("=" * 70)
    
    try:
        test_iris_classification()
        test_customer_segmentation()
        test_house_price_prediction()
        
        print("\n" + "=" * 70)
        print("  ✅ ALL PRACTICAL TESTS COMPLETED!")
        print("  These examples demonstrate real-world ML applications")
        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())
