import time
import dask
import numpy as np
from dask.distributed import as_completed, Client, LocalCluster

def log_progress(futures, interval=100, label="Tasks"):
    """
    A simple progress logger that prints a summary line for every N completed tasks.
    Args:
        futures (list): A list of Dask futures to monitor.
        interval (int, optional): How often to print a summary line (i.e., every
                                  'interval' tasks). Defaults to 100.
        label (str, optional): A label for the task group being monitored.
                               Defaults to "Tasks".
    """
    total_tasks = len(futures)
    if total_tasks == 0:
        return
    # Use as_completed to iterate over futures as they finish
    completed_iterator = as_completed(futures)
    counter = 0
    
    for future in completed_iterator:
        counter += 1
        # Print a summary line at the specified interval or for the very last task
        if counter % interval == 0 or counter == total_tasks:
            timestamp = time.strftime('%H:%M:%S')
            percentage = (counter / total_tasks) * 100
            print(
                f"[{timestamp}] {counter} / {total_tasks} "
                f"tasks complete ({percentage:.1f}%)"
            )
            
    print(f"--- {label}: All tasks complete ---")

def _setup_dask_client(args):
    """
    Sets up and returns a Dask client and cluster.
    """
    dask.config.set({
        'distributed.worker.daemon': False,
        'distributed.nanny.environ.OMP_NUM_THREADS': '1',
        'distributed.nanny.environ.MKL_NUM_THREADS': '1',
        'distributed.nanny.environ.OPENBLAS_NUM_THREADS': '1',
        "distributed.scheduler.worker-ttl": None
    })

    # Default to None (disabled) unless explicitly requested via CLI
    dashboard_address = ':8787' if args.dashboard else None

    cluster = LocalCluster(
        n_workers=args.n_cores,
        threads_per_worker=1,
        protocol='tcp',
        dashboard_address=dashboard_address
    )
    client = Client(cluster)
    return cluster, client

def _fibonacci_sphere(samples):
    """Generate points uniformly distributed on a sphere."""
    points = []
    phi = (1 + np.sqrt(5)) / 2
    for i in range(samples):
        theta = 2 * np.pi * (i / phi - np.floor(i / phi))
        z = 1 - (2 * i) / (samples - 1)
        r = np.sqrt(1 - z * z)
        x = r * np.cos(theta)
        y = r * np.sin(theta)
        points.append((x, y, z))
    return np.array(points)

def generate_orientations(method, n_points):
    """
    Generates orientation vectors and weights for powder averaging.

    Args:
        method (str): The quadrature scheme (e.g., 'fibonacci').
        n_points (int): The number of orientation points to generate.

    Returns:
        tuple: A tuple containing:
            - np.ndarray: An array of (n_points, 3) orientation vectors.
            - np.ndarray: An array of (n_points,) corresponding weights.
    """
    if method == 'fibonacci':
        # For fibonacci sphere, all points have equal weight.
        # We generate points on the full sphere and take the top hemisphere.
        if n_points == 1:
            # Default case: single orientation along z-axis
            vectors = np.array([[0.0, 0.0, 1.0]])
            weights = np.array([1.0])
        else:
            full_sphere_points = _fibonacci_sphere(2 * n_points)
            vectors = full_sphere_points[full_sphere_points[:, 2] >= 0]
            weights = np.full(vectors.shape[0], 1.0 / vectors.shape[0])
        return vectors, weights
    else:
        raise NotImplementedError(f"Quadrature method '{method}' is not yet supported.")
