"""
Heavy tests for machine learning training tasks using distributed compute
"""
import pytest
import time
import threading
from distributed_compute import Coordinator, Worker


def train_simple_model(data):
    """Simple training simulation"""
    X, y = data
    # Simulate training time
    result = sum(y) / len(y) if y else 0
    return result


@pytest.mark.timeout(60)
def test_distributed_training():
    """Test distributed model training"""
    import random
    random.seed(42)
    
    # Create training batches
    batches = []
    for i in range(10):
        X = [[random.uniform(-10, 10) for _ in range(5)] for _ in range(100)]
        y = [sum(x[:2]) + random.uniform(-0.5, 0.5) for x in X]
        batches.append((X, y))
    
    coordinator = Coordinator(port=5563, verbose=False)
    coordinator.start_server()
    time.sleep(0.5)
    
    workers = []
    for i in range(3):
        worker = Worker('localhost', 5563, name=f'worker_{i}')
        workers.append(worker)
        thread = threading.Thread(target=worker.start, daemon=True)
        thread.start()
    
    time.sleep(2)
    
    try:
        results = coordinator.map(train_simple_model, batches, timeout=50)
        assert len(results) == len(batches)
        assert all(isinstance(r, (int, float)) for r in results)
    finally:
        coordinator.stop_server()
        for worker in workers:
            worker.stop()
