import numpy as np
from numba import jit
from scipy.special import erf

# --- Original Lineshape Functions ---

def get_lineshape_function(lineshape, fwhm):
    """
    Returns a callable function for a given phonon lineshape.

    The Gaussian and Lorentzian functions are defined to be zero for
    non-positive energies, reflecting the physical nature of phonons.

    Args:
        lineshape (str): The name of the lineshape ('gaussian',
                         'lorentzian', 'antilorentzian').
        fwhm (float): The full-width at half-maximum of the lineshape.

    Returns:
        function: A function that takes energy (x) and a center (mu) and
                  returns the lineshape value.
    """
    if fwhm <= 0:
        return lambda x, mu: np.isclose(x, mu).astype(float)

    if lineshape == 'gaussian':
        sigma = fwhm / (2 * np.sqrt(2 * np.log(2)))
        norm = 1 / (sigma * np.sqrt(2 * np.pi))

        def gaussian(x, mu):
            res = np.zeros_like(x, dtype=float)
            mask = x > 0
            if np.any(mask):
                res[mask] = norm * np.exp(-(x[mask] - mu)**2 / (2 * sigma**2))
            return res
        return gaussian

    elif lineshape == 'lorentzian':
        hwhm = fwhm / 2.0

        def lorentzian(x, mu):
            res = np.zeros_like(x, dtype=float)
            mask = x > 0
            if np.any(mask):
                res[mask] = (hwhm / np.pi) / ((x[mask] - mu)**2 + hwhm**2)
            return res
        return lorentzian

    elif lineshape == 'antilorentzian':
        hwhm = fwhm / 2.0

        def antilorentzian(x, mu):
            res = np.zeros_like(x, dtype=float)
            mask = x > 0
            if np.any(mask):
                norm_factor = 2 * np.arctan(2 * mu / fwhm) if fwhm > 0 else np.pi
                l_minus = (hwhm / np.pi) / ((x[mask] - mu)**2 + hwhm**2)
                l_plus = (hwhm / np.pi) / ((x[mask] + mu)**2 + hwhm**2)
                if norm_factor > 1e-9:
                    res[mask] = (np.pi * (l_minus - l_plus) / norm_factor)
            return res
        return antilorentzian

    raise ValueError(f"Unknown lineshape: {lineshape}")

def get_integration_width(lineshape, fwhm, scale, erange=[0,np.inf]):
    """
    Determines the energy window (half-width) around a phonon's central energy for
    integration, based on a scaling factor.
    """
    # Return full range if scale is not positive
    if scale <= 0:
        return erange[1] - erange[0]

    # --- Gaussian lineshape ---
    if lineshape == 'gaussian':
        # Convert FWHM to sigma, then return the half-width (scale * sigma)
        sigma = fwhm / (2 * np.sqrt(2 * np.log(2)))
        return scale * sigma

    # --- Lorentzian or Antilorentzian lineshape ---
    elif lineshape in ['lorentzian', 'antilorentzian']:
        # Calculate the fractional area 'p' under a Gaussian for 'scale' standard deviations.
        p = erf(scale / np.sqrt(2))
        
        # Calculate the corresponding full integration width in units of FWHM for a Lorentzian.
        full_width_in_fwhm = np.tan(p * np.pi / 2)
        
        # Return the half-width.
        return (full_width_in_fwhm / 2) * fwhm
        
    # --- Fallback ---
    return 5 * fwhm
# --- JIT-compiled Lineshape Functions for Numba ---

@jit(nopython=True, cache=True, nogil=True)
def lorentzian_jit(x, center, hwhm):
    """JIT-compiled Lorentzian part for the antilorentzian (normalised)."""
    return (hwhm / np.pi) / ((x - center)**2 + hwhm**2)

@jit(nopython=True, cache=True, nogil=True)
def antilorentzian_jit(x, center, hwhm):
    """
    JIT-compiled anti-Lorentzian, faithfully reproducing the original's shape.
    This function remains normalised as the shape depends on the normalisation factor.
    """
    res = np.zeros_like(x)
    fwhm = 2.0 * hwhm
    
    # Loop to handle the x > 0 mask, which is more Numba-friendly
    for i in range(len(x)):
        if x[i] > 0:
            if fwhm > 0:
                norm_factor = 2.0 * np.arctan(2.0 * center / fwhm)
            else:
                norm_factor = np.pi
            
            l_minus = lorentzian_jit(x[i], center, hwhm)
            l_plus = lorentzian_jit(x[i], -center, hwhm)
            
            if norm_factor > 1e-9: # Avoid division by zero
                res[i] = (np.pi * (l_minus - l_plus) / norm_factor)
    return res

@jit(nopython=True, cache=True, nogil=True)
def gaussian_jit(x, center, hwhm):
    """JIT-compiled Gaussian lineshape (unnormalised)."""
    sigma = hwhm / np.sqrt(2.0 * np.log(2.0))
    return np.exp(-0.5 * ((x - center) / sigma)**2)

@jit(nopython=True, cache=True, nogil=True)
def dispatch_lineshape_jit(x, center, hwhm, lineshape_code):
    """Calls the correct JIT-compiled lineshape based on an integer code."""
    if lineshape_code == 0: # antilorentzian
        return antilorentzian_jit(x, center, hwhm)
    elif lineshape_code == 1: # lorentzian
        # For direct Lorentzian calculations, we can use an unnormalised form
        return hwhm**2 / ((x - center)**2 + hwhm**2)
    elif lineshape_code == 2: # gaussian
        return gaussian_jit(x, center, hwhm)
    return np.zeros_like(x) # Should not be reached
