import numpy as np
from numba import jit, prange
from .constants import MU_B_CM_T, G_E

@jit(nopython=True, cache=True, nogil=True)
def phase(z):
    return np.arctan2(z.imag, z.real)

@jit(nopython=True, cache=True, nogil=True)
def tr_symmetrise(op, is_even):

    n = op.shape[0]

    for j in range(0, n, 2):
        for k in range(0, n, 2):
            if j == k:

                if is_even:
                    op[j, j] = 0.5 * (op[j, j].real + op[j+1, j+1].real)
                    op[j+1, j+1] = op[j, j]
                    op[j, j+1] = 0.
                    op[j+1, j] = 0.
                else:
                    op[j, j] = 0.5 * (op[j, j].real - op[j+1, j+1].real)
                    op[j+1, j+1] = -op[j, j]
                    op[j, j+1] = 0.5 * (op[j, j+1] + op[j+1, j].conjugate())
                    op[j+1, j] = op[j, j+1].conjugate()
            else:
                del_phase = (
                        (phase(op[j, k]) - phase(op[k+1, j+1])) -
                        (phase(op[j, k+1]) - phase(op[k, j+1]) - np.pi)
                        ) % (2 * np.pi)

                if del_phase > np.pi:
                    del_phase -= 2 * np.pi
                if del_phase < -np.pi:
                    del_phase += 2 * np.pi

                avg_mag = 0.5 * (abs(op[j, k]) + abs(op[k+1, j+1]))
                op[j, k] = avg_mag * np.exp(1j * phase(op[j, k]) - 0.25 * del_phase)
                op[k+1, j+1] = avg_mag * np.exp(1j * phase(op[k+1, j+1]) + 0.25 * del_phase)

                avg_mag = 0.5 * (abs(op[j, k+1]) + abs(op[k, j+1]))
                op[j, k+1] = avg_mag * np.exp(1j * phase(op[j, k+1]) + 0.25 * del_phase)
                op[k, j+1] = avg_mag * np.exp(1j * phase(op[k, j+1]) - 0.25 * del_phase)

    return op

def symmetrise_couplings(couplings, is_kramers):
    """
    Enforces physical symmetries on the spin-phonon coupling matrices V.

    This function enforces Hermiticity. For Kramers ions, it also enforces time-reversal 
    symmetry using a JIT-compiled helper function for performance. V is time-reversal even.

    Args:
        couplings (dict): A dictionary of {mode_index: coupling_matrix}.
        is_kramers (bool): True if the system is a Kramers ion.

    Returns:
        dict: The symmetrised coupling matrices.
    """
    symmetrised_couplings = {}
    for key, V in couplings.items():
        # Enforce Hermiticity for all systems.
        V_herm = 0.5 * (V + V.conj().T)
        if is_kramers:
            # Enforce even time-reversal symmetry for Kramers ions.
            symmetrised_couplings[key] = tr_symmetrise(V_herm, True)
        else:
            symmetrised_couplings[key] = V_herm
    return symmetrised_couplings


@jit(nopython=True, cache=True, nogil=True)
def _make_hermitian_jit(a):
    """
    Enforces Hermiticity A = 0.5 * (A + A^H) efficiently.
    """
    n = a.shape[0]
    for i in prange(n):
        # Diagonal must be real
        a[i, i] = 0.5 * (a[i, i] + a[i, i].conjugate())
        # Symmetrise off-diagonals
        for j in range(i + 1, n):
            val = 0.5 * (a[i, j] + a[j, i].conjugate())
            a[i, j] = val
            a[j, i] = val.conjugate()
    return a

def symmetrise_op(op, *, is_kramers, is_tr_even):
    """
    Enforces physical symmetries on a single operator.

    This function enforces Hermiticity. For Kramers ions, it also enforces time-reversal 
    symmetry.

    Args:
        op (np.ndarray): An operator matrix.
        is_kramers (bool): True if the system is a Kramers ion.
        is_tr_even (bool): True if the operator is time-reversal even.

    Returns:
        np.ndarray: The symmetrised op.
    """
    # Enforce Hermiticity.
    op_herm = _make_hermitian_jit(op.copy())
    # op_herm = 0.5 * (op + op.conj().T)
    if is_kramers:
        # Enforce even time-reversal symmetry for Kramers ions.
        op_symm = tr_symmetrise(op_herm, is_tr_even)
    else:
        op_symm = op_herm
    return op_symm

def apply_zeeman_splitting(energies, angmom, spin, B_vec, couplings=None):
    """
    Applies a Zeeman field to the electronic Hamiltonian.

    This function diagonalises the total Hamiltonian (H_0 + H_Zeeman) to
    obtain the new energy levels and state properties. If the 'couplings'
    argument is provided, it also transforms the coupling matrices into the
    new Zeeman eigenbasis.

    Args:
        energies (np.ndarray): Zero-field electronic energies.
        angmom (np.ndarray): Angular momentum operator matrices (L).
        spin (np.ndarray): Spin operator matrices (S).
        B_vec (np.ndarray): The magnetic field vector in Tesla.
        couplings (dict, optional): The symmetrised spin-phonon coupling
                                    matrices. If None, this step is skipped.

    Returns:
        tuple: A tuple containing:
            - np.ndarray: New Zeeman-split energies.
            - dict or None: Coupling matrices in the new Zeeman basis, or None.
            - np.ndarray: Magnetic moment of each state along the field axis.
            - np.ndarray: Expectation value of <mJ> for each state.
            - np.ndarray: The unitary transformation matrix, U.
    """
    # Define the total angular momentum and magnetic moment operators.
    J_total_op = angmom + spin
    mu_op = MU_B_CM_T * (angmom + G_E * spin)

    # Build and diagonalise the Hamiltonian to get new energies and eigenvectors.
    H_zee = -np.einsum('i,ijk->jk', B_vec, mu_op)
    H_total = np.diag(energies) + H_zee
    new_energies, U = np.linalg.eigh(H_total)
    U_dagger = U.conj().T

    # Conditionally transform coupling matrices if they are provided.
    new_couplings = None
    if couplings is not None:
        new_couplings = {key: U_dagger @ V_old @ U for key, V_old in couplings.items()}

    # Calculate state properties (<mJ> and magnetic moment).
    B_norm = np.linalg.norm(B_vec)
    if B_norm > 1e-9:
        B_unit = B_vec / B_norm
        
        # Calculate operators (O(N^2) if diagonal/sparse, O(N^3) if dense)
        J_parallel_op = np.einsum('i,ijk->jk', B_unit, J_total_op)
        mu_parallel = np.einsum('i,ijk->jk', B_unit, mu_op)

        # Compute diagonal elements only: <i|Op|i>
        # Avoids full U^H @ Op @ U multiplication. 
        
        J_temp = J_parallel_op @ U
        state_mJ = np.real(np.sum(U.conj() * J_temp, axis=0))
        
        mu_temp = mu_parallel @ U
        state_mu = np.real(np.sum(U.conj() * mu_temp, axis=0)) / MU_B_CM_T
        
    else:
        # Zero field case
        mu_z_temp = mu_op[2] @ U
        state_mu = np.real(np.sum(U.conj() * mu_z_temp, axis=0)) / MU_B_CM_T
        
        Jz_temp = J_total_op[2] @ U
        state_mJ = np.real(np.sum(U.conj() * Jz_temp, axis=0))
    return new_energies, new_couplings, state_mu, state_mJ, U
