"""Nuclear spin-rotation Hamiltonian operators for B state in coupled Omega basis.

This module implements the nuclear spin-rotation (NSR) interaction for the thallium
nucleus in the B electronic state. The NSR interaction couples nuclear spin angular
momentum to molecular rotational angular momentum.

The nuclear spin-rotation effect arises from the magnetic field generated by the
rotating charged nuclei acting on the nuclear magnetic moment.
"""

from __future__ import annotations

from functools import lru_cache

import numpy as np

from centrex_tlf.constants import BConstants
from centrex_tlf.states import CoupledBasisState, CoupledState

from ..wigner import sixj_f


@lru_cache(maxsize=int(1e6))
def H_c_Tl(psi: CoupledBasisState, constants: BConstants) -> CoupledState:
    """Nuclear spin-rotation operator for Tl nucleus in B state (coupled Omega basis).

    Calculates the nuclear spin-rotation (NSR) interaction for the thallium nucleus,
    represented by the scalar coupling constant c(Tl). This diagonal operator couples
    the Tl nuclear spin I₁ to the molecular rotational angular momentum J.

    The operator is diagonal in all quantum numbers (J, I₁, I₂, F₁, F, mF, Ω), making
    it a first-order energy correction to the hyperfine structure. The matrix element
    is proportional to I₁·J.

    Args:
        psi: Input coupled basis state |J,F₁,F,mF,I₁,I₂,Ω⟩ in B electronic state.
        constants: Molecular constants for TlF B state, requires c_Tl parameter
            (Tl nuclear spin-rotation constant).

    Returns:
        CoupledState with a single term containing the NSR energy shift and the
        same basis state as input. The amplitude includes a Wigner 6-j symbol
        for angular momentum recoupling and appropriate normalization factors.

    Notes:
        - The operator is completely diagonal (no state mixing)
        - Selection rules: ΔJ = 0; ΔI₁ = 0; ΔF₁ = 0; ΔF = 0; ΔmF = 0; ΔΩ = 0
        - For TlF B state: c(Tl) ≈ -7.83 MHz
        - The negative sign indicates antiparallel coupling is favored
        - NSR constants are typically much smaller than hyperfine constants
        - The interaction scales with J(J+1) and I₁(I₁+1)

    References:
        Ramsey, N. F. (1953). "Nuclear Moments." Wiley, New York.
        Brown, J. M., & Carrington, A. (2003). "Rotational Spectroscopy of Diatomic
        Molecules." Cambridge University Press. Section 8.6.

    Example:
        >>> state = CoupledBasisState(F=1, mF=0, F1=0.5, J=1, I1=0.5, I2=0.5, Omega=1)
        >>> constants = BConstants(c_Tl=-7830000.0)
        >>> result = H_c_Tl(state, constants)
        >>> # Returns the same state with NSR energy shift
    """
    # Find the quantum numbers of the input state
    Jp = psi.J
    I1p = psi.I1
    I2p = psi.I2
    F1p = psi.F1
    Fp = psi.F
    mFp = psi.mF
    Omegap = psi.Omega

    # J, I1, I2, F and mF are the same for both states
    J = Jp
    I1 = I1p
    I2 = I2p
    F = Fp
    F1 = F1p
    mF = mFp

    # Omega also doesn't change
    Omega = Omegap

    # Calculate matrix element
    amp = (
        constants.c_Tl
        * (-1) ** (I1 + F1 + J)
        * sixj_f(Jp, I1, F1, I1, J, 1)
        * np.sqrt(I1 * (I1 + 1) * (2 * I1 + 1) * J * (J + 1) * (2 * J + 1))
    )

    basis_state = CoupledBasisState(
        F,
        mF,
        F1,
        J,
        I1,
        I2,
        Omega=Omega,
        electronic_state=psi.electronic_state,
        P=psi.P,
    )

    return CoupledState([(amp, basis_state)])
