"""
Heavy tests for scientific computing tasks using distributed compute
"""
import pytest
import math
import time
import threading
from distributed_compute import Coordinator, Worker


def monte_carlo_pi(num_samples):
    """Estimate pi using Monte Carlo method"""
    import random
    random.seed()
    
    inside_circle = 0
    for _ in range(num_samples):
        x = random.uniform(-1, 1)
        y = random.uniform(-1, 1)
        if x*x + y*y <= 1:
            inside_circle += 1
    
    return inside_circle


def compute_primes(n):
    """Count primes up to n"""
    if n < 2:
        return 0
    
    primes = []
    for num in range(2, min(n, 10000)):  # Limit for performance
        is_prime = True
        for p in primes:
            if p * p > num:
                break
            if num % p == 0:
                is_prime = False
                break
        if is_prime:
            primes.append(num)
    
    return len(primes)


def numerical_integration(params):
    """Perform numerical integration"""
    start, end, num_points = params
    
    def f(x):
        return x**2 + 2*x + 1
    
    dx = (end - start) / num_points
    integral = 0.5 * (f(start) + f(end))
    
    for i in range(1, num_points):
        x = start + i * dx
        integral += f(x)
    
    integral *= dx
    return integral


@pytest.mark.timeout(60)
def test_monte_carlo_pi_estimation():
    """Test distributed Monte Carlo pi estimation"""
    samples_per_task = 100000
    num_tasks = 10
    
    tasks_data = [samples_per_task for _ in range(num_tasks)]
    
    coordinator = Coordinator(port=5566, verbose=False)
    coordinator.start_server()
    time.sleep(0.5)
    
    workers = []
    for i in range(3):
        worker = Worker('localhost', 5566, name=f'worker_{i}')
        workers.append(worker)
        thread = threading.Thread(target=worker.start, daemon=True)
        thread.start()
    
    time.sleep(2)
    
    try:
        results = coordinator.map(monte_carlo_pi, tasks_data, timeout=50)
        
        # Estimate pi
        total_inside = sum(results)
        total_samples = samples_per_task * num_tasks
        pi_estimate = 4 * total_inside / total_samples
        
        # Pi should be approximately 3.14159
        assert 3.0 < pi_estimate < 3.3
    finally:
        coordinator.stop_server()
        for worker in workers:
            worker.stop()


@pytest.mark.timeout(60)
def test_prime_counting():
    """Test distributed prime counting"""
    numbers = [1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000, 10000]
    
    coordinator = Coordinator(port=5567, verbose=False)
    coordinator.start_server()
    time.sleep(0.5)
    
    workers = []
    for i in range(4):
        worker = Worker('localhost', 5567, name=f'worker_{i}')
        workers.append(worker)
        thread = threading.Thread(target=worker.start, daemon=True)
        thread.start()
    
    time.sleep(2)
    
    try:
        results = coordinator.map(compute_primes, numbers, timeout=50)
        
        assert len(results) == len(numbers)
        assert all(isinstance(r, int) for r in results)
        assert all(r > 0 for r in results)
    finally:
        coordinator.stop_server()
        for worker in workers:
            worker.stop()


@pytest.mark.timeout(60)
def test_numerical_integration():
    """Test distributed numerical integration"""
    # Integrate x^2 + 2x + 1 over different intervals
    intervals = [
        (0, 1, 1000),
        (1, 2, 1000),
        (2, 3, 1000),
        (3, 4, 1000),
        (4, 5, 1000),
        (0, 5, 1000)
    ]
    
    coordinator = Coordinator(port=5568, verbose=False)
    coordinator.start_server()
    time.sleep(0.5)
    
    workers = []
    for i in range(3):
        worker = Worker('localhost', 5568, name=f'worker_{i}')
        workers.append(worker)
        thread = threading.Thread(target=worker.start, daemon=True)
        thread.start()
    
    time.sleep(2)
    
    try:
        results = coordinator.map(numerical_integration, intervals, timeout=50)
        
        assert len(results) == len(intervals)
        assert all(isinstance(r, (int, float)) for r in results)
        
        # Integral of x^2 + 2x + 1 from 0 to 1 should be approximately 2.333...
        assert abs(results[0] - 7/3) < 0.1
    finally:
        coordinator.stop_server()
        for worker in workers:
            worker.stop()
