# Auto-generated Script
# Task: Fit each pixel's spectrum with a single Lorentzian or Drude-Lorentz model (using lmfit LorentzianModel + linear background) in the 0.2–0.8 eV range to extract three maps: (1) plasmon peak center energy, (2) peak FWHM/linewidth, and (3) peak amplitude. Initialize center at 0.4 eV, sigma at 0.05 eV, constrain center to [0.2, 0.7] eV and sigma to [0.01, 0.3] eV. For pixels with low total intensity (bottom 5%), skip fitting. This will directly map the spatial variation of the LSPR energy across the nanocrystal array, revealing coupling effects at interfaces and mode confinement at edges/corners.

def analyze_feature(data, axis):
    # Reshape data to (pixels, energy)
    ny, nx, ne = data.shape
    spectra = data.reshape(-1, ne)
    n_pixels = spectra.shape[0]
    
    # Define energy range mask for fitting: 0.2 to 0.8 eV
    mask = (axis >= 0.2) & (axis <= 0.8)
    x_fit = axis[mask]
    
    # Initialize output maps
    center_map = np.full(n_pixels, np.nan)
    fwhm_map = np.full(n_pixels, np.nan)
    amplitude_map = np.full(n_pixels, np.nan)
    
    # Compute total intensity per pixel for thresholding
    total_intensity = np.sum(spectra[:, mask], axis=1)
    threshold = np.percentile(total_intensity, 5)
    
    # Define Lorentzian + linear background model function for curve_fit
    # Lorentzian: amplitude * (sigma**2) / ((x - center)**2 + sigma**2) 
    # This is the Cauchy/Lorentzian form where FWHM = 2*sigma
    # Using the form: A / (1 + ((x - x0)/gamma)**2) + slope*x + intercept
    # where gamma = HWHM, FWHM = 2*gamma
    # Peak height = A, amplitude in area sense = A * pi * gamma
    def lorentzian_linear(x, A, center, gamma, slope, intercept):
        return A / (1.0 + ((x - center) / gamma)**2) + slope * x + intercept
    
    # Smooth spectra slightly for better initial guesses and convergence
    # Use gaussian filter along energy axis
    smoothed_spectra = np.zeros_like(spectra)
    for i in range(n_pixels):
        smoothed_spectra[i] = gaussian_filter(spectra[i].astype(np.float64), sigma=3)
    
    # Bounds for curve_fit: [A, center, gamma, slope, intercept]
    # A > 0, center in [0.2, 0.7], gamma (HWHM) in [0.005, 0.15] (so sigma=HWHM, FWHM=2*gamma in [0.01, 0.3])
    lower_bounds = [0.0,   0.2,  0.005, -np.inf, -np.inf]
    upper_bounds = [np.inf, 0.7, 0.15,   np.inf,  np.inf]
    
    for i in range(n_pixels):
        # Skip low intensity pixels
        if total_intensity[i] < threshold:
            continue
        
        y_fit = spectra[i, mask].astype(np.float64)
        y_smooth = smoothed_spectra[i, mask]
        
        # Skip if all zeros or negative
        if np.all(y_fit <= 0) or np.max(y_fit) == 0:
            continue
        
        # Initial guesses from smoothed data
        # Estimate linear background from endpoints
        bg_slope = (y_smooth[-1] - y_smooth[0]) / (x_fit[-1] - x_fit[0]) if (x_fit[-1] - x_fit[0]) != 0 else 0
        bg_intercept = y_smooth[0] - bg_slope * x_fit[0]
        
        # Subtract background to find peak
        y_nobg = y_smooth - (bg_slope * x_fit + bg_intercept)
        
        # Find peak position from residual
        peak_idx = np.argmax(y_nobg)
        init_center = x_fit[peak_idx]
        # Clamp initial center
        init_center = np.clip(init_center, 0.2, 0.7)
        init_A = max(y_nobg[peak_idx], 1.0)
        init_gamma = 0.05  # HWHM = 0.05 eV, so FWHM = 0.1 eV, sigma ~ 0.05
        
        p0 = [init_A, init_center, init_gamma, bg_slope, bg_intercept]
        
        try:
            popt, pcov = curve_fit(
                lorentzian_linear, x_fit, y_fit,
                p0=p0,
                bounds=(lower_bounds, upper_bounds),
                maxfev=5000
            )
            A_fit, center_fit, gamma_fit, slope_fit, intercept_fit = popt
            
            # Quality check: fitted center should be reasonable
            # and amplitude should be positive and meaningful
            if A_fit > 0 and 0.2 <= center_fit <= 0.7:
                center_map[i] = center_fit
                fwhm_map[i] = 2.0 * gamma_fit  # FWHM = 2 * HWHM
                amplitude_map[i] = A_fit  # Peak height
        except (RuntimeError, ValueError):
            # Fitting failed, leave as NaN
            continue
    
    # Reshape maps back to spatial dimensions
    center_map = center_map.reshape(ny, nx)
    fwhm_map = fwhm_map.reshape(ny, nx)
    amplitude_map = amplitude_map.reshape(ny, nx)
    
    return {
        "maps": {
            "Plasmon_Peak_Center": center_map,
            "Plasmon_Peak_FWHM": fwhm_map,
            "Plasmon_Peak_Amplitude": amplitude_map
        },
        "units": {
            "Plasmon_Peak_Center": "eV",
            "Plasmon_Peak_FWHM": "eV",
            "Plasmon_Peak_Amplitude": "a.u."
        },
        "description": (
            "Single Lorentzian + linear background fit in 0.2-0.8 eV range to extract "
            "LSPR characteristics. The plasmon peak center map reveals spatial variation "
            "of the localized surface plasmon resonance energy across the FT:IO nanocrystal "
            "array, highlighting coupling effects at interfaces and mode confinement at "
            "edges/corners. FWHM maps indicate damping/lifetime variations, while amplitude "
            "maps show the strength of plasmonic excitation. Bottom 5% intensity pixels "
            "are excluded to avoid fitting noise in vacuum/background regions."
        )
    }
