# Auto-generated Script
# Task: Fit each pixel spectrum with a Lorentzian (or Voigt) peak model to extract the localized surface plasmon resonance (LSPR) parameters. For each pixel: (1) Fit a single Lorentzian peak (center, amplitude, FWHM) plus a linear background over the 0.2–1.13 eV range. The plasmon peak is expected near 0.35–0.50 eV. (2) Generate spatial maps of: peak center energy (eV), FWHM/linewidth (eV, related to plasmon damping), and integrated peak intensity. (3) If residuals from a single-Lorentzian fit show systematic structure above 0.6 eV, attempt a two-Lorentzian fit to check for a secondary coupled plasmon mode. Use lmfit LorentzianModel + LinearModel for robust fitting with parameter bounds (center: 0.25–0.60 eV, sigma: 0.02–0.3 eV). This will reveal spatial variations in carrier concentration and inter-particle coupling across the nanocrystal array.

def analyze_feature(data, axis):
    ny, nx, ne = data.shape
    n_pixels = ny * nx
    spectra = data.reshape(n_pixels, ne).astype(np.float64)
    x = axis.astype(np.float64)
    
    # Smooth spectra along spectral axis to aid convergence
    smoothed = gaussian_filter(spectra, sigma=(0, 2))
    
    # Lorentzian + linear background (lmfit convention: sigma=HWHM, FWHM=2*sigma)
    def lor_lin(x, amplitude, center, sigma, slope, intercept):
        return (amplitude / np.pi) * (sigma / ((x - center)**2 + sigma**2)) + slope * x + intercept
    
    # Two Lorentzians + linear background
    def two_lor_lin(x, a1, c1, s1, a2, c2, s2, slope, intercept):
        lor1 = (a1 / np.pi) * (s1 / ((x - c1)**2 + s1**2))
        lor2 = (a2 / np.pi) * (s2 / ((x - c2)**2 + s2**2))
        return lor1 + lor2 + slope * x + intercept
    
    # Initialize output arrays
    peak_center = np.full(n_pixels, np.nan)
    peak_fwhm = np.full(n_pixels, np.nan)
    peak_intensity = np.full(n_pixels, np.nan)
    peak2_center = np.full(n_pixels, np.nan)
    peak2_fwhm = np.full(n_pixels, np.nan)
    peak2_intensity = np.full(n_pixels, np.nan)
    has_second_peak = np.zeros(n_pixels)
    fit_rsq = np.full(n_pixels, np.nan)
    
    # Precompute masks for peak search regions
    mask_peak = (x >= 0.30) & (x <= 0.55)
    mask_high = x >= 0.6
    x_high = x[mask_high]
    x_peak_vals = x[mask_peak]
    
    n_edge = min(50, ne // 10)
    
    for i in range(n_pixels):
        y = smoothed[i]
        
        # Skip very low signal pixels
        y_range = np.max(y) - np.min(y)
        if y_range < 1e-10:
            continue
        
        # Estimate linear background from spectral edges
        slope_init = (np.mean(y[-n_edge:]) - np.mean(y[:n_edge])) / (x[-1] - x[0])
        intercept_init = np.mean(y[:n_edge]) - slope_init * x[0]
        
        # Subtract background to find initial peak location
        bg = slope_init * x + intercept_init
        y_sub = y - bg
        
        if np.any(mask_peak):
            peak_vals = y_sub[mask_peak]
            idx_max = np.argmax(peak_vals)
            center_init = x_peak_vals[idx_max]
            peak_height = max(peak_vals[idx_max], 1e-6)
        else:
            center_init = 0.40
            peak_height = max(np.max(y_sub), 1e-6)
        
        sigma_init = 0.08  # HWHM initial guess
        amplitude_init = max(peak_height * np.pi * sigma_init, 1e-6)
        
        p0 = [amplitude_init, center_init, sigma_init, slope_init, intercept_init]
        lower = [0, 0.25, 0.02, -np.inf, -np.inf]
        upper = [np.inf, 0.60, 0.30, np.inf, np.inf]
        
        try:
            popt, _ = curve_fit(lor_lin, x, y, p0=p0, bounds=(lower, upper), maxfev=3000)
            amp, cen, sig, slp, intcpt = popt
            
            peak_center[i] = cen
            peak_fwhm[i] = 2.0 * sig  # FWHM = 2 * HWHM
            peak_intensity[i] = amp   # integrated area under Lorentzian
            
            y_fit = lor_lin(x, *popt)
            residuals = y - y_fit
            
            ss_res = np.sum(residuals**2)
            ss_tot = np.sum((y - np.mean(y))**2)
            if ss_tot > 0:
                fit_rsq[i] = 1.0 - ss_res / ss_tot
            
            # Check for systematic residual structure above 0.6 eV
            res_high = residuals[mask_high]
            if len(res_high) > 10:
                res_smooth = gaussian_filter(res_high, sigma=3)
                noise_level = np.std(residuals)
                
                if noise_level > 0 and np.max(res_smooth) > 3.0 * noise_level:
                    idx_max2 = np.argmax(res_smooth)
                    c2_init = x_high[idx_max2]
                    h2 = res_smooth[idx_max2]
                    s2_init = 0.08
                    a2_init = max(h2 * np.pi * s2_init, 1e-6)
                    
                    p0_2 = [amp, cen, sig, a2_init, c2_init, s2_init, slp, intcpt]
                    lower_2 = [0, 0.25, 0.02, 0, 0.55, 0.02, -np.inf, -np.inf]
                    upper_2 = [np.inf, 0.60, 0.30, np.inf, 1.15, 0.30, np.inf, np.inf]
                    
                    try:
                        popt2, _ = curve_fit(two_lor_lin, x, y, p0=p0_2,
                                             bounds=(lower_2, upper_2), maxfev=5000)
                        a1f, c1f, s1f, a2f, c2f, s2f, sl2f, in2f = popt2
                        
                        y_fit2 = two_lor_lin(x, *popt2)
                        rss2 = np.sum((y - y_fit2)**2)
                        
                        # Accept if >25% RSS improvement and secondary peak is significant
                        if rss2 < ss_res * 0.75 and a2f > amp * 0.05:
                            peak_center[i] = c1f
                            peak_fwhm[i] = 2.0 * s1f
                            peak_intensity[i] = a1f
                            peak2_center[i] = c2f
                            peak2_fwhm[i] = 2.0 * s2f
                            peak2_intensity[i] = a2f
                            has_second_peak[i] = 1.0
                            
                            if ss_tot > 0:
                                fit_rsq[i] = 1.0 - rss2 / ss_tot
                    except Exception:
                        pass
        except Exception:
            pass
    
    # Reshape all arrays to spatial maps
    return {
        "maps": {
            "LSPR_Peak_Center": peak_center.reshape(ny, nx),
            "LSPR_FWHM": peak_fwhm.reshape(ny, nx),
            "LSPR_Integrated_Intensity": peak_intensity.reshape(ny, nx),
            "Secondary_Peak_Center": peak2_center.reshape(ny, nx),
            "Secondary_Peak_FWHM": peak2_fwhm.reshape(ny, nx),
            "Secondary_Peak_Intensity": peak2_intensity.reshape(ny, nx),
            "Two_Peak_Flag": has_second_peak.reshape(ny, nx),
            "Fit_R_Squared": fit_rsq.reshape(ny, nx)
        },
        "units": {
            "LSPR_Peak_Center": "eV",
            "LSPR_FWHM": "eV",
            "LSPR_Integrated_Intensity": "a.u.",
            "Secondary_Peak_Center": "eV",
            "Secondary_Peak_FWHM": "eV",
            "Secondary_Peak_Intensity": "a.u.",
            "Two_Peak_Flag": "binary",
            "Fit_R_Squared": "dimensionless"
        },
        "description": "Pixel-wise Lorentzian+linear background fitting of LSPR in low-loss EELS (0.2-1.13 eV). "
                       "Peak center energy maps carrier concentration variations in FT:IO nanocrystals (Drude model: "
                       "higher energy = higher carrier density). FWHM maps plasmon damping rates (scattering, "
                       "size effects, inter-band transitions). Integrated intensity reflects oscillator strength. "
                       "Secondary peak detection above 0.6 eV identifies inter-particle coupled plasmon modes "
                       "arising from electromagnetic coupling in the nanocrystal array."
    }
