"""
Energy grid convolution for spin wave spectra.

Convolves delta function spin wave dispersions with energy resolution
functions to produce intensity maps suitable for comparison with
inelastic neutron scattering data.
"""

import numpy as np
from typing import Dict, Union, Optional, List


def sw_egrid(spec: Dict, Evect: np.ndarray = None, 
             component: str = 'Sperp',
             dE: float = 0.1,
             T: float = 0,
             imagChk: bool = True,
             **kwargs) -> Dict:
    """
    Convolve spin wave spectrum with energy resolution.
    
    Parameters
    ----------
    spec : dict
        Output from SpinW.spinwave()
    Evect : array
        Energy grid for output, shape (n_E,)
        Default: 0 to 1.2*max(omega) with 100 points
    component : str
        Correlation function component to calculate:
        - 'Sperp': Neutron scattering cross-section (S_perp)
        - 'Sxx': S^xx correlation
        - 'Syy': S^yy correlation
        - 'Szz': S^zz correlation
        - 'Sxx+Syy': Sum of in-plane components
        - 'Sxx+Syy+Szz': Total scattering
    dE : float
        Energy resolution (Gaussian FWHM) in meV
    T : float
        Temperature in Kelvin (for Bose factor)
    imagChk : bool
        If True, check for imaginary frequencies (unstable modes)
        
    Returns
    -------
    spec : dict
        Updated spectrum dictionary with:
        - 'Evect': Energy grid
        - 'swConv': Convolved intensity, shape (n_E, n_q)
        - 'swInt': Integrated intensity, shape (n_q, n_modes)
    """
    omega = spec['omega']  # (n_q, n_modes)
    Sab = spec.get('Sab')   # (n_q, 3, 3, n_modes)
    hkl = spec['hkl']
    
    n_q = omega.shape[0]
    n_modes = omega.shape[1]
    
    # Set up energy grid
    if Evect is None:
        omega_max = np.max(np.real(omega))
        if omega_max < 1e-10:
            omega_max = 1.0
        Evect = np.linspace(0, 1.2 * omega_max, 100)
    
    n_E = len(Evect)
    
    # Calculate the requested component
    swInt = _calculate_component(Sab, component, hkl, spec.get('hklA'))
    
    # Check for imaginary frequencies
    if imagChk:
        omega_imag = np.imag(omega)
        if np.any(omega_imag > 1e-6):
            import warnings
            warnings.warn(
                f"Found {np.sum(omega_imag > 1e-6)} modes with imaginary frequencies. "
                "The magnetic structure may not be the ground state."
            )
    
    # Use real part of omega
    omega_real = np.real(omega)
    
    # Gaussian convolution
    sigma = dE / (2 * np.sqrt(2 * np.log(2)))  # FWHM to sigma
    
    swConv = np.zeros((n_E, n_q))
    
    for iq in range(n_q):
        for imode in range(n_modes):
            E_mode = omega_real[iq, imode]
            I_mode = np.real(swInt[iq, imode])
            
            if E_mode > 0 and I_mode > 0:
                # Gaussian profile
                gauss = np.exp(-0.5 * ((Evect - E_mode) / sigma)**2)
                gauss = gauss / (sigma * np.sqrt(2 * np.pi))
                
                # Add Bose factor if T > 0
                if T > 0:
                    from ..core.constants import CONSTANTS
                    n_bose = 1.0 / (np.exp(E_mode / (CONSTANTS.kB * T)) - 1) + 1
                    I_mode *= n_bose
                
                swConv[:, iq] += I_mode * gauss
    
    # Update spec dictionary
    spec['Evect'] = Evect
    spec['swConv'] = swConv
    spec['swInt'] = swInt
    spec['component'] = component
    
    return spec


def _calculate_component(Sab: np.ndarray, component: str,
                         hkl: np.ndarray, hklA: np.ndarray = None) -> np.ndarray:
    """Calculate the requested spin-spin correlation component."""
    if Sab is None:
        # Return placeholder if correlation not calculated
        n_q = len(hkl)
        return np.ones((n_q, 1))
    
    n_q = Sab.shape[0]
    n_modes = Sab.shape[3]
    
    component = component.lower()
    
    if component == 'sperp':
        # Perpendicular component for neutron scattering
        # S_perp = S_total - (Q · S · Q) / Q^2
        if hklA is None:
            # Fallback to sum
            swInt = np.sum(np.diagonal(Sab, axis1=1, axis2=2), axis=1)
        else:
            Q_norm = hklA / (np.linalg.norm(hklA, axis=1, keepdims=True) + 1e-10)
            
            swInt = np.zeros((n_q, n_modes), dtype=np.complex128)
            for iq in range(n_q):
                Q_hat = Q_norm[iq]
                for imode in range(n_modes):
                    S_tot = np.trace(Sab[iq, :, :, imode])
                    # Q·S·Q
                    S_para = Q_hat @ Sab[iq, :, :, imode] @ Q_hat
                    swInt[iq, imode] = S_tot - S_para
    
    elif component == 'sxx':
        swInt = Sab[:, 0, 0, :]
    
    elif component == 'syy':
        swInt = Sab[:, 1, 1, :]
    
    elif component == 'szz':
        swInt = Sab[:, 2, 2, :]
    
    elif component == 'sxx+syy':
        swInt = Sab[:, 0, 0, :] + Sab[:, 1, 1, :]
    
    elif component == 'sxx+syy+szz':
        swInt = np.sum(np.diagonal(Sab, axis1=1, axis2=2), axis=1)
    
    else:
        raise ValueError(f"Unknown component: {component}")
    
    return swInt


def sw_instrument(spec: Dict, dE: float = None, dQ: float = None,
                  ki: float = None, **kwargs) -> Dict:
    """
    Apply instrumental resolution to spectrum.
    
    Parameters
    ----------
    spec : dict
        Output from sw_egrid()
    dE : float
        Energy resolution FWHM in meV
    dQ : float
        Q resolution FWHM in Å^-1
    ki : float
        Incident neutron wavevector in Å^-1
        
    Returns
    -------
    spec : dict
        Updated spectrum with instrumental effects applied
    """
    # Apply energy resolution if not already done
    if dE is not None and 'swConv' in spec:
        # Re-convolve with new resolution
        spec = sw_egrid(spec, spec.get('Evect'), dE=dE, **kwargs)
    
    # Q resolution broadening would go here
    # This is a simplified placeholder
    
    return spec
