import numpy as np
from scipy.spatial.transform import Rotation

from .constants import MU_B_CM_T, G_E

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 (np.ndarray): The normalised field direction vector for a single calculation.
            - magnitudes_to_run (list): A list of field magnitudes (in Tesla).
            - rotation_matrix (np.ndarray): The matrix to rotate from the input to the principal axis frame.
    """
    rotation_matrix = np.identity(3)
    principal_frame_z_axis = np.array([0.0, 0.0, 1.0])

    if  args.reorient:
        print("\nReorienting: z-axis will be aligned with the principal g-tensor axis.")
        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)]
        print(f"Principal g-axis in input frame: [{principal_axis[0]:.4f}, {principal_axis[1]:.4f}, {principal_axis[2]:.4f}]")

        r, _ = Rotation.align_vectors([principal_axis], [principal_frame_z_axis])
        rotation_matrix = r.as_matrix()

    if args.field_direction:
        if orientations == 1:
            print("Using user-provided --field_direction relative to the principal axis frame.")
        field_direction_principal_frame = np.array(args.field_direction)
    else:
        field_direction_principal_frame = principal_frame_z_axis

    field_direction_input_frame = rotation_matrix @ field_direction_principal_frame
    norm = np.linalg.norm(field_direction_principal_frame)

    if norm < 1e-9:
        if orientations == 1:
            print("Warning: Field direction vector has zero magnitude. Using principal frame z-axis.")
        B_unit = rotation_matrix.T @ principal_frame_z_axis
        B_unit_principal = principal_frame_z_axis
    else:
        B_unit = field_direction_input_frame / np.linalg.norm(field_direction_input_frame)
        B_unit_principal = field_direction_principal_frame / norm

    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[0]:.4f}, {B_unit[1]:.4f}, {B_unit[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, magnitudes_to_run, rotation_matrix
