using MSI_src
using Statistics, Images, Interpolations

# 1. Configuration
data_path = "Leaf.imzML"
mask_path = "region_of_interest.png"
output_folder = "datos_para_ai"
img_size = (224, 224) # Standard size for CNN
mkpath(output_folder)

# 2. Data loading
data = OpenMSIData(data_path)

# 3. Get important peaks (you can use your previous logic)
mzs, avg_ints = MSI_src.get_average_spectrum(data, mask_path=mask_path)
threshold = maximum(avg_ints) * 0.02
peak_indices = findall(x -> x > threshold, avg_ints)

println("Exporting $(length(peak_indices)) ions...")

for idx in peak_indices
    mz_val = mzs[idx]
    if mz_val < 50.0 || mz_val > 700.0 continue end

    # A. Extract slice (raw, no plots)
    raw_slice = MSI_src.get_mz_slice(data, mz_val, 0.05, mask_path=mask_path)
    
    # B. TrIQ Normalization (as in your script)
    vals = filter(x -> x > 0 && isfinite(x), raw_slice)
    if isempty(vals) continue end
    q = quantile(vals, 0.98)
    q = q == 0 ? maximum(vals) : q
    enhanced = clamp.(raw_slice ./ q, 0.0, 1.0)

    # C. Rescale to 224x224 (Crucial for CNN)
    # We use imresize from Images.jl
    img_resized = imresize(enhanced, img_size)

    # D. Save as raw PNG (Grayscale, no axes)
    filename = "ion_$(round(mz_val, digits=2)).png"
    save(joinpath(output_folder, filename), Gray.(img_resized))
end