"""
Neutron scattering cross-section calculations.
"""

import numpy as np
from typing import Dict, Optional, Tuple


def sw_neutron(spec: Dict, pol: bool = False,
               uv: Tuple[np.ndarray, np.ndarray] = None,
               **kwargs) -> Dict:
    """
    Calculate neutron scattering cross-section.
    
    Parameters
    ----------
    spec : dict
        Output from SpinW.spinwave()
    pol : bool
        Calculate polarized neutron cross-sections
    uv : tuple
        Polarization direction (u, v) vectors for polarized neutron
        
    Returns
    -------
    spec : dict
        Updated spectrum with neutron cross-sections:
        - 'Sperp': Unpolarized cross-section
        - 'Mab': Polarized cross-section components (if pol=True)
    """
    omega = spec['omega']
    Sab = spec.get('Sab')
    hklA = spec.get('hklA')
    
    if Sab is None:
        return spec
    
    n_q = Sab.shape[0]
    n_modes = Sab.shape[3]
    
    # Calculate unpolarized cross-section: S_perp = S_total - (Q·S·Q)/Q²
    Sperp = np.zeros((n_q, n_modes), dtype=np.complex128)
    
    for iq in range(n_q):
        Q = hklA[iq] if hklA is not None else np.array([1, 0, 0])
        Q_norm = Q / (np.linalg.norm(Q) + 1e-10)
        
        for imode in range(n_modes):
            # Total scattering
            S_tot = np.trace(Sab[iq, :, :, imode])
            # Parallel component
            S_para = Q_norm @ Sab[iq, :, :, imode] @ Q_norm
            # Perpendicular
            Sperp[iq, imode] = S_tot - S_para
    
    spec['Sperp'] = Sperp
    
    if pol:
        # Polarized neutron cross-sections
        # See: Moon, Riste, Koehler, Phys. Rev. 181, 920 (1969)
        spec['Mab'] = _calculate_polarized(Sab, hklA, uv)
    
    return spec


def _calculate_polarized(Sab: np.ndarray, hklA: np.ndarray,
                         uv: Tuple[np.ndarray, np.ndarray] = None) -> Dict:
    """
    Calculate polarized neutron scattering cross-sections.
    
    Returns all components needed for polarization analysis:
    - NSF (non-spin-flip): M_xx, M_yy, M_zz
    - SF (spin-flip): M_xy, M_xz, M_yz
    """
    n_q, _, _, n_modes = Sab.shape
    
    # Default polarization frame
    if uv is None:
        # x || Q, z vertical
        pass
    
    Mab = {
        'Mxx': np.zeros((n_q, n_modes), dtype=np.complex128),
        'Myy': np.zeros((n_q, n_modes), dtype=np.complex128),
        'Mzz': np.zeros((n_q, n_modes), dtype=np.complex128),
        'Mxy': np.zeros((n_q, n_modes), dtype=np.complex128),
    }
    
    for iq in range(n_q):
        for imode in range(n_modes):
            S = Sab[iq, :, :, imode]
            Mab['Mxx'][iq, imode] = S[0, 0]
            Mab['Myy'][iq, imode] = S[1, 1]
            Mab['Mzz'][iq, imode] = S[2, 2]
            Mab['Mxy'][iq, imode] = S[0, 1]
    
    return Mab


def sw_formfactor(ion: str, Q: np.ndarray) -> np.ndarray:
    """
    Calculate magnetic form factor.
    
    Uses the analytical approximation:
    f(Q) = sum_i A_i * exp(-a_i * s^2) + c
    where s = Q / (4*pi)
    
    Parameters
    ----------
    ion : str
        Ion label (e.g., 'Fe3+', 'Cu2+', 'Mn2+')
    Q : np.ndarray
        Q values in Å^-1, shape (n_q,) or (n_q, 3)
        
    Returns
    -------
    ff : np.ndarray
        Form factor values
    """
    # Form factor coefficients from International Tables
    # Format: [A, a, B, b, C, c, D]
    formfactor_table = {
        'Fe0': [0.0706, 35.008, 0.3589, 15.358, 0.5819, 5.561, -0.0114],
        'Fe2': [0.0263, 34.960, 0.3668, 15.943, 0.6188, 5.594, -0.0119],
        'Fe3': [0.3972, 13.244, 0.6295, 4.903, -0.0314, 0.350, 0.0044],
        'Fe3+': [0.3972, 13.244, 0.6295, 4.903, -0.0314, 0.350, 0.0044],
        'Cu0': [0.0909, 34.984, 0.4088, 11.443, 0.5128, 3.825, -0.0124],
        'Cu2': [0.0232, 34.969, 0.4023, 11.564, 0.5882, 3.843, -0.0137],
        'Cu2+': [0.0232, 34.969, 0.4023, 11.564, 0.5882, 3.843, -0.0137],
        'Mn2': [0.4220, 17.684, 0.5948, 6.005, 0.0043, -0.609, -0.0219],
        'Mn2+': [0.4220, 17.684, 0.5948, 6.005, 0.0043, -0.609, -0.0219],
        'Co2': [0.4332, 14.355, 0.5857, 4.608, -0.0382, 0.134, 0.0179],
        'Co2+': [0.4332, 14.355, 0.5857, 4.608, -0.0382, 0.134, 0.0179],
        'Ni2': [0.0163, 35.883, 0.3916, 13.223, 0.6052, 4.339, -0.0133],
        'Ni2+': [0.0163, 35.883, 0.3916, 13.223, 0.6052, 4.339, -0.0133],
        'Cr3': [0.2840, 18.044, 0.7181, 7.047, 0.0053, -0.099, -0.0078],
        'Cr3+': [0.2840, 18.044, 0.7181, 7.047, 0.0053, -0.099, -0.0078],
    }
    
    # Clean ion label
    ion_clean = ion.replace('+', '').replace('-', '')
    
    if ion_clean not in formfactor_table:
        # Try with charge
        if ion not in formfactor_table:
            # Default to constant
            return np.ones_like(Q if Q.ndim == 1 else Q[:, 0])
    
    coeffs = formfactor_table.get(ion, formfactor_table.get(ion_clean))
    
    # Calculate |Q|
    if Q.ndim == 2:
        Q_mag = np.linalg.norm(Q, axis=1)
    else:
        Q_mag = np.abs(Q)
    
    # s = Q / (4*pi)
    s = Q_mag / (4 * np.pi)
    s2 = s**2
    
    # Calculate form factor
    A, a, B, b, C, c, D = coeffs
    ff = A * np.exp(-a * s2) + B * np.exp(-b * s2) + C * np.exp(-c * s2) + D
    
    return ff


def sw_sqw(spec: Dict, Evect: np.ndarray = None, 
           qvect: np.ndarray = None,
           T: float = 0, **kwargs) -> Dict:
    """
    Calculate S(Q, omega) dynamical structure factor.
    
    Parameters
    ----------
    spec : dict
        Output from sw_egrid()
    Evect : np.ndarray
        Energy transfer values
    qvect : np.ndarray
        Q-values (if different from spec)
    T : float
        Temperature in Kelvin
        
    Returns
    -------
    sqw : dict
        S(Q, omega) and related quantities
    """
    from ..core.constants import CONSTANTS
    
    if 'swConv' not in spec:
        raise ValueError("Run sw_egrid first")
    
    swConv = spec['swConv']
    Evect = spec['Evect']
    
    # Apply detailed balance if T > 0
    if T > 0:
        for iE, E in enumerate(Evect):
            if E > 0:
                n_bose = 1.0 / (np.exp(E / (CONSTANTS.kB * T)) - 1) + 1
                swConv[iE, :] *= n_bose
    
    return {
        'SQW': swConv,
        'Evect': Evect,
        'hkl': spec['hkl'],
    }
