# Auto-generated Script
# Task: Fit each pixel's EELS spectrum with a Lorentzian (or Voigt) model for the primary LSPR peak in the 0.25–0.6 eV range to extract peak center energy, FWHM (linewidth/damping), and peak amplitude. Use lmfit LorentzianModel with reasonable bounds (center: 0.3–0.55 eV, FWHM: 0.05–0.4 eV, amplitude > 0). Add a linear background to handle the tail of the zero-loss peak. For pixels where the residual after the first peak fit exceeds a threshold, attempt a second Lorentzian near 0.6–0.85 eV to capture possible coupled or higher-order plasmon modes. Output spatial maps of: (1) primary peak center energy, (2) primary peak FWHM, (3) primary peak amplitude, (4) presence/absence and center energy of secondary mode. This will directly map LSPR mode energies and coupling across the self-assembled nanocrystal array.

def analyze_feature(data, axis):
    ny, nx, ne = data.shape
    n_pixels = ny * nx

    # Smooth along energy axis to stabilize fitting
    smoothed_3d = gaussian_filter(data, sigma=(0, 0, 2))
    spectra = smoothed_3d.reshape(n_pixels, ne)

    # Define fit range
    fit_mask = (axis >= 0.22) & (axis <= 0.92)
    x_fit = axis[fit_mask]
    n_fit = len(x_fit)

    # Secondary peak residual check region
    sec_mask = (x_fit >= 0.6) & (x_fit <= 0.85)
    prim_mask_in_fit = (x_fit >= 0.30) & (x_fit <= 0.55)

    # Height-parameterized Lorentzian: height at center
    def lorentzian(x, height, center, fwhm):
        hwhm = fwhm / 2.0
        return height * hwhm**2 / ((x - center)**2 + hwhm**2)

    def model_1peak(x, h, c, w, slope, intercept):
        return lorentzian(x, h, c, w) + slope * x + intercept

    def model_2peak(x, h1, c1, w1, h2, c2, w2, slope, intercept):
        return (lorentzian(x, h1, c1, w1) + lorentzian(x, h2, c2, w2)
                + slope * x + intercept)

    # Output arrays
    pc = np.full(n_pixels, np.nan)   # primary center
    pw = np.full(n_pixels, np.nan)   # primary FWHM
    pa = np.full(n_pixels, np.nan)   # primary amplitude (height)
    sp = np.zeros(n_pixels)          # secondary present
    sc = np.full(n_pixels, np.nan)   # secondary center

    # Bounds for single peak fit
    lb1 = np.array([0.0,    0.30, 0.05, -np.inf, -np.inf])
    ub1 = np.array([np.inf, 0.55, 0.40,  np.inf,  np.inf])

    # Bounds for two-peak fit
    lb2 = np.array([0.0,    0.30, 0.05, 0.0,    0.58, 0.03, -np.inf, -np.inf])
    ub2 = np.array([np.inf, 0.55, 0.40, np.inf, 0.88, 0.40,  np.inf,  np.inf])

    # Edge indices for background estimation
    n_edge = max(5, n_fit // 20)

    for i in range(n_pixels):
        y = spectra[i, fit_mask].copy()

        # Skip empty/bad pixels
        y_max = np.max(y)
        if y_max <= 0 or np.any(np.isnan(y)):
            continue

        # --- Initial background estimate from edges ---
        y_left = np.mean(y[:n_edge])
        y_right = np.mean(y[-n_edge:])
        x_left = np.mean(x_fit[:n_edge])
        x_right = np.mean(x_fit[-n_edge:])
        slope0 = (y_right - y_left) / (x_right - x_left + 1e-12)
        intercept0 = y_left - slope0 * x_left

        # Subtract rough background
        y_nobg = y - (slope0 * x_fit + intercept0)

        # --- Find primary peak initial guess ---
        if not np.any(prim_mask_in_fit):
            continue
        y_prim = y_nobg[prim_mask_in_fit]
        x_prim = x_fit[prim_mask_in_fit]
        idx_max = np.argmax(y_prim)
        c0 = x_prim[idx_max]
        h0 = max(y_prim[idx_max], np.std(y) * 0.05, 1e-8)

        p0_single = [h0, c0, 0.15, slope0, intercept0]

        # --- Single peak fit ---
        try:
            popt, _ = curve_fit(model_1peak, x_fit, y, p0=p0_single,
                                bounds=(lb1, ub1), maxfev=4000,
                                ftol=1e-6, xtol=1e-6)

            pc[i] = popt[1]
            pw[i] = popt[2]
            pa[i] = popt[0]

            # --- Check residual for secondary peak ---
            if np.any(sec_mask):
                residual = y - model_1peak(x_fit, *popt)
                res_sec = residual[sec_mask]
                x_sec = x_fit[sec_mask]
                max_res = np.max(res_sec)

                # Threshold: 8% of primary peak height
                threshold = 0.08 * popt[0]

                if max_res > threshold and max_res > 0:
                    idx_sec = np.argmax(res_sec)
                    c2_0 = np.clip(x_sec[idx_sec], 0.59, 0.87)
                    h2_0 = max(res_sec[idx_sec], 1e-8)

                    p0_double = [popt[0], popt[1], popt[2],
                                 h2_0, c2_0, 0.12,
                                 popt[3], popt[4]]

                    try:
                        popt2, _ = curve_fit(model_2peak, x_fit, y,
                                             p0=p0_double,
                                             bounds=(lb2, ub2),
                                             maxfev=6000,
                                             ftol=1e-6, xtol=1e-6)

                        # Validate: secondary must be > 5% of primary height
                        # and two-peak fit must reduce chi-squared
                        res1 = np.sum((y - model_1peak(x_fit, *popt))**2)
                        res2 = np.sum((y - model_2peak(x_fit, *popt2))**2)

                        if popt2[3] > 0.05 * popt2[0] and res2 < 0.90 * res1:
                            pc[i] = popt2[1]
                            pw[i] = popt2[2]
                            pa[i] = popt2[0]
                            sp[i] = 1.0
                            sc[i] = popt2[4]
                    except (RuntimeError, ValueError):
                        pass

        except (RuntimeError, ValueError):
            continue

    # Reshape to spatial maps
    primary_center_map = pc.reshape(ny, nx)
    primary_fwhm_map = pw.reshape(ny, nx)
    primary_amplitude_map = pa.reshape(ny, nx)
    secondary_present_map = sp.reshape(ny, nx)
    secondary_center_map = sc.reshape(ny, nx)

    return {
        "maps": {
            "LSPR_Peak_Center": primary_center_map,
            "LSPR_Peak_FWHM": primary_fwhm_map,
            "LSPR_Peak_Amplitude": primary_amplitude_map,
            "Secondary_Mode_Present": secondary_present_map,
            "Secondary_Mode_Center": secondary_center_map
        },
        "units": {
            "LSPR_Peak_Center": "eV",
            "LSPR_Peak_FWHM": "eV",
            "LSPR_Peak_Amplitude": "a.u.",
            "Secondary_Mode_Present": "binary",
            "Secondary_Mode_Center": "eV"
        },
        "description": (
            "Pixel-wise Lorentzian + linear background fitting of low-loss EELS spectra "
            "to map LSPR modes in self-assembled FT:IO nanocrystal arrays. "
            "Primary peak center (0.3-0.55 eV) maps the dominant plasmon resonance energy; "
            "FWHM maps the mode linewidth/damping; amplitude maps oscillator strength/intensity. "
            "Secondary mode detection (0.58-0.88 eV) from fit residuals identifies "
            "coupled or higher-order plasmon modes arising from inter-particle electromagnetic "
            "coupling in the nanocrystal superlattice. Two-peak fits are validated by "
            "requiring >10% chi-squared improvement and secondary height >5% of primary."
        )
    }