import time
import dask
import numpy as np
from dask.distributed import as_completed, Client, LocalCluster
from scipy.spatial.transform import Rotation

from .constants import MU_B_CM_T, G_E

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

    if args.n_workers:
        cluster = LocalCluster(
            n_workers=args.n_workers,
            threads_per_worker=args.n_threads,
            protocol='tcp',
            dashboard_address=dashboard_address
        )
    else:
        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.")

def _get_attr_bool(attrs, key, default=False):
    """Safely reads a boolean attribute from HDF5, handling potential string conversion."""
    val = attrs.get(key, default)
    if isinstance(val, (bytes, str)):
        if isinstance(val, bytes): val = val.decode('utf-8')
        return val.lower() == 'true'
    return bool(val)

def setup_field_orientation(args, angmom, spin, orientations=1):
    """Determines the field direction and magnitudes for the calculation.

    This function handles the g-tensor reorientation to align the calculation's
    z-axis with the principal magnetic axis of the ground state doublet. It also
    interprets user-provided field vectors and magnitudes from the command line.

    Args:
        args (argparse.Namespace): The parsed command-line arguments.
        angmom (np.ndarray): The angular momentum matrices (L_x, L_y, L_z).
        spin (np.ndarray): The spin matrices (S_x, S_y, S_z).
        orientations (int): The number of orientations being calculated.

    Returns:
        tuple: A tuple containing:
            - B_unit_input (np.ndarray): The normalised field direction vector in the input frame.
            - magnitudes_to_run (list): A list of field magnitudes (in Tesla).
            - rotation_matrix (np.ndarray): The matrix to rotate from the Principal to the Input frame.
    """
    principal_frame_z_axis = np.array([0.0, 0.0, 1.0])

    # 1. Unconditionally calculate the principal g-tensor axis for reference
    mu_op = MU_B_CM_T * (angmom + G_E * spin)
    mu_op_doublet = mu_op[:, 0:2, 0:2]
    g_tensor = 2 * np.real(np.einsum('kab,lba->kl', mu_op_doublet, mu_op_doublet))
    g_squared = g_tensor.T @ g_tensor
    eigvals, eigvecs = np.linalg.eigh(g_squared)

    principal_axis = eigvecs[:, np.argmax(eigvals)]

    # r_I2P aligns the principal axis to the z-axis (Input to Principal)
    r_I2P, _ = Rotation.align_vectors([principal_axis], [principal_frame_z_axis])
    R_I2P = r_I2P.as_matrix()
    R_P2I = R_I2P.T  # Rotates principal frame vectors back to the input frame

    # 2. Determine the rotation matrix used for the rest of the program
    if args.reorient:
        print("\nReorienting: z-axis will be aligned with the principal g-tensor axis.")
        print(f"Principal g-axis in input frame: [{principal_axis[0]:.4f}, {principal_axis[1]:.4f}, {principal_axis[2]:.4f}]")
        # We want to rotate user inputs/grids from the Principal frame to the Input frame
        rotation_matrix = R_P2I 
    else:
        rotation_matrix = np.identity(3)

    # 3. Handle the user's field direction input
    if args.field_direction:
        if orientations == 1:
            if args.reorient:
                print("Using provided --field_direction relative to the principal axis frame.")
            else:
                print("Using provided --field_direction relative to the input frame (--no_reorient).")
        field_direction_given = np.array(args.field_direction)
    else:
        field_direction_given = principal_frame_z_axis

    norm = np.linalg.norm(field_direction_given)

    if norm < 1e-9:
        if orientations == 1:
            print("Warning: Field direction vector has zero magnitude. Using z-axis.")
        B_unit_calc_frame = principal_frame_z_axis
    else:
        B_unit_calc_frame = field_direction_given / norm

    # Map the calculation frame vector to the input frame (where angmom/spin live)
    B_unit_input = rotation_matrix @ B_unit_calc_frame

    # Unconditionally map the input frame vector to the principal frame for user reference
    B_unit_principal = R_I2P @ B_unit_input

    if orientations == 1:
        print("\nNormalised field direction:")
        print(f"  principal axis frame:\t[{B_unit_principal[0]:.4f}, {B_unit_principal[1]:.4f}, {B_unit_principal[2]:.4f}]")
        print(f"  input frame:\t\t[{B_unit_input[0]:.4f}, {B_unit_input[1]:.4f}, {B_unit_input[2]:.4f}]")

    if args.field_magnitudes:
        magnitudes_to_run = args.field_magnitudes
    elif args.field_direction:
        magnitudes_to_run = [norm]
    else:
        magnitudes_to_run = [0.0002]

    return B_unit_input, magnitudes_to_run, rotation_matrix