# Auto-generated Script
# Task: Fit each pixel's spectrum with a single Lorentzian (or Drude-type) peak model on a polynomial background to extract: (1) plasmon peak center energy (eV), (2) peak FWHM/linewidth (eV), and (3) peak amplitude. Use lmfit LorentzianModel + LinearModel. Constrain center to [0.2, 0.8] eV, FWHM to [0.05, 0.6] eV. For pixels with low total signal (<5% of max), flag as background. Generate spatial maps of peak center, FWHM, and amplitude, plus histograms of each parameter. This will directly map the spatial variation of the LSPR energy across the nanocrystal array, revealing coupling effects and carrier concentration gradients.

def analyze_feature(data, axis):
    """
    Fit each pixel's spectrum with a Lorentzian peak on a linear background
    to extract plasmon peak center, FWHM, and amplitude.
    Uses raw curve_fit for speed over large pixel counts.
    """
    # Reshape to (n_pixels, n_channels)
    ny, nx, n_channels = data.shape
    spectra = data.reshape(-1, n_channels)
    n_pixels = spectra.shape[0]
    
    # Pre-allocate output maps
    peak_center = np.full(n_pixels, np.nan)
    peak_fwhm = np.full(n_pixels, np.nan)
    peak_amplitude = np.full(n_pixels, np.nan)
    background_flag = np.zeros(n_pixels, dtype=np.float64)
    fit_success = np.zeros(n_pixels, dtype=np.float64)
    
    # Determine background threshold: 5% of max total signal
    total_signal = np.sum(spectra, axis=1)
    max_total = np.max(total_signal)
    bg_threshold = 0.05 * max_total
    
    # Flag background pixels
    bg_mask = total_signal < bg_threshold
    background_flag[bg_mask] = 1.0
    
    # Define Lorentzian + linear model
    # Lorentzian: amplitude * (sigma^2) / ((x - center)^2 + sigma^2) 
    # where FWHM = 2*sigma
    # lmfit convention: LorentzianModel amplitude A gives:
    #   f(x) = (A/pi) * (sigma / ((x-center)^2 + sigma^2))
    # We'll use a direct parameterization for curve_fit:
    #   f(x) = amp * (gamma/2)^2 / ((x - x0)^2 + (gamma/2)^2) + slope*x + intercept
    # where gamma = FWHM, amp = peak height
    
    def lorentzian_linear(x, amp, x0, gamma, slope, intercept):
        half_gamma = gamma / 2.0
        return amp * (half_gamma**2) / ((x - x0)**2 + half_gamma**2) + slope * x + intercept
    
    # Smooth the axis data for initial guess estimation
    from scipy.ndimage import uniform_filter1d
    
    # Constraint bounds for curve_fit:
    # amp: [0, inf], x0: [0.2, 0.8], gamma: [0.05, 0.6], slope: [-inf, inf], intercept: [-inf, inf]
    lower_bounds = [0.0,   0.2,  0.05, -np.inf, -np.inf]
    upper_bounds = [np.inf, 0.8, 0.6,   np.inf,  np.inf]
    
    x = axis.astype(np.float64)
    
    # Restrict fitting range to where center is expected: use full axis but focus initial guess
    # Create a mask for the center search region
    center_mask = (x >= 0.2) & (x <= 0.8)
    
    for i in range(n_pixels):
        # Skip background pixels
        if bg_mask[i]:
            peak_center[i] = np.nan
            peak_fwhm[i] = np.nan
            peak_amplitude[i] = np.nan
            fit_success[i] = 0.0
            continue
        
        y = spectra[i].astype(np.float64)
        
        # Smooth for initial guess
        y_smooth = uniform_filter1d(y, size=15)
        
        # Estimate linear background from edges
        n_edge = 30
        left_val = np.mean(y_smooth[:n_edge])
        right_val = np.mean(y_smooth[-n_edge:])
        slope_init = (right_val - left_val) / (x[-1] - x[0])
        intercept_init = left_val - slope_init * x[0]
        
        # Subtract estimated background to find peak
        y_bg_sub = y_smooth - (slope_init * x + intercept_init)
        
        # Find peak in the constrained center region
        y_in_region = y_bg_sub.copy()
        y_in_region[~center_mask] = -np.inf
        
        idx_peak = np.argmax(y_in_region)
        x0_init = x[idx_peak]
        amp_init = max(y_bg_sub[idx_peak], 1e-6)
        
        # Estimate FWHM from half-max width
        half_max = amp_init / 2.0
        above_half = np.where((y_bg_sub >= half_max) & center_mask)[0]
        if len(above_half) > 1:
            gamma_init = x[above_half[-1]] - x[above_half[0]]
            gamma_init = np.clip(gamma_init, 0.05, 0.6)
        else:
            gamma_init = 0.15
        
        p0 = [amp_init, x0_init, gamma_init, slope_init, intercept_init]
        
        try:
            popt, pcov = curve_fit(
                lorentzian_linear, x, y, p0=p0,
                bounds=(lower_bounds, upper_bounds),
                maxfev=3000
            )
            
            amp_fit, x0_fit, gamma_fit, slope_fit, intercept_fit = popt
            
            # Quality check: compute residuals
            y_fit = lorentzian_linear(x, *popt)
            residual = np.sum((y - y_fit)**2)
            total_var = np.sum((y - np.mean(y))**2)
            
            if total_var > 0:
                r_squared = 1.0 - residual / total_var
            else:
                r_squared = 0.0
            
            # Check parameter uncertainties from covariance
            perr = np.sqrt(np.diag(pcov)) if np.all(np.isfinite(pcov)) else np.full(5, np.inf)
            
            # Apply quality filters
            # 1. R-squared should be reasonable
            # 2. Peak amplitude should be meaningful (at least 1% of signal range)
            # 3. Center uncertainty should be reasonable
            signal_range = np.max(y) - np.min(y)
            
            good_fit = (
                r_squared > 0.3 and
                amp_fit > 0.01 * signal_range and
                np.isfinite(perr[1]) and perr[1] < 0.2 and  # center uncertainty
                gamma_fit > 0.05 + 1e-4 and  # not pegged at lower bound
                gamma_fit < 0.6 - 1e-4  # not pegged at upper bound
            )
            
            if good_fit:
                peak_center[i] = x0_fit
                peak_fwhm[i] = gamma_fit
                peak_amplitude[i] = amp_fit
                fit_success[i] = 1.0
            else:
                peak_center[i] = np.nan
                peak_fwhm[i] = np.nan
                peak_amplitude[i] = np.nan
                fit_success[i] = 0.0
                
        except Exception:
            peak_center[i] = np.nan
            peak_fwhm[i] = np.nan
            peak_amplitude[i] = np.nan
            fit_success[i] = 0.0
    
    # Reshape all maps back to spatial dimensions
    center_map = peak_center.reshape(ny, nx)
    fwhm_map = peak_fwhm.reshape(ny, nx)
    amplitude_map = peak_amplitude.reshape(ny, nx)
    bg_flag_map = background_flag.reshape(ny, nx)
    fit_success_map = fit_success.reshape(ny, nx)
    
    return {
        "maps": {
            "Peak_Center_Energy": center_map,
            "Peak_FWHM": fwhm_map,
            "Peak_Amplitude": amplitude_map,
            "Background_Flag": bg_flag_map,
            "Fit_Success": fit_success_map
        },
        "units": {
            "Peak_Center_Energy": "eV",
            "Peak_FWHM": "eV",
            "Peak_Amplitude": "a.u.",
            "Background_Flag": "binary",
            "Fit_Success": "binary"
        },
        "description": (
            "Lorentzian peak fitting on linear background for LSPR characterization. "
            "Peak_Center_Energy maps the plasmon resonance energy (constrained 0.2-0.8 eV), "
            "revealing carrier concentration gradients and coupling effects across the nanocrystal array. "
            "Peak_FWHM maps the linewidth (constrained 0.05-0.6 eV), indicating damping mechanisms "
            "and size/shape heterogeneity. Peak_Amplitude maps the oscillator strength. "
            "Background_Flag identifies low-signal pixels (<5% of max total signal). "
            "Fit_Success indicates pixels where quality criteria (R²>0.3, meaningful amplitude, "
            "reasonable parameter uncertainties, non-pegged bounds) were met."
        )
    }
