# BRIDGE SCRIPT: Connects JuliaMSI framework to the CNN Preprocessing pipeline.
# VERSION (V3): Relying directly on precomputed analytics and multiple mz slices for max throughtput

using Pkg
juliamsi_path = "/app/JuliaMSI"
push!(LOAD_PATH, joinpath(juliamsi_path, "src"))

# Load JuliaMSI core
include(joinpath(juliamsi_path, "src", "MSI_src.jl"))

using .MSI_src
using Statistics, Images, Interpolations, BenchmarkTools

# CONFIGURATION
data_path = "/app/data/Leaf.imzML"
mask_path = "/app/data/region_of_interest.png"
output_folder = "/app/datos_para_ai"
img_size = (256, 256)
tolerance = 0.05
mkpath(output_folder)

println("--- JULIA BATCH ENGINE ---")

function run_production_pipeline()
    # 1. LOAD DATA
    println("Step 1: Ingesting Data...")
    data = OpenMSIData(data_path)

    orig_height = data.image_dims[2]
    orig_width  = data.image_dims[1]
    if orig_height < 256 || orig_width < 256
        final_img_size = (orig_height, orig_width)
        println("  [Notice] Original image is smaller than 256x256. Keeping original size: ", final_img_size)
    else
        final_img_size = (256, 256)
        println("  [Notice] Resizing down to target size: ", final_img_size)
    end
    
    # 2. RUN FRAMEWORK ANALYTICS
    println("Step 2: Checking Dataset Precomputed Boundaries...")
    # This matches the framework's internal logic to safely discover the true boundaries.
    # It populates global_min_mz and global_max_mz by reading the binary headers properly.
    MSI_src.precompute_analytics(data)
    
    # 3. EXTRACT AVERAGE SPECTRUM 
    println("Step 3: Calculating ROI Average Spectrum...")
    mzs, avg_ints = MSI_src.get_average_spectrum(data, mask_path=mask_path)
    
    println("  True Dataset Min m/z: ", round(minimum(mzs), digits=2))
    println("  True Dataset Max m/z: ", round(maximum(mzs), digits=2))

    # 4. PEAK PICKING USING TRUE SPECTRA LIMITS
    # Identify relevant ion channels above a 2% intensity base peak threshold
    threshold = maximum(avg_ints) * 0.02
    peak_indices = findall(x -> x > threshold, avg_ints)

    min_vis_mz = data.global_min_mz[] > 0.0 ? data.global_min_mz[] : 50.0
    max_vis_mz = data.global_max_mz[] > 0.0 ? data.global_max_mz[] : 1000.0
    
    # Isolate targets that fall strictly within the standard CNN metabolic training window
    target_masses = Float64[]
    for idx in peak_indices
        mz_val = mzs[idx]
        if mz_val >= min_vis_mz && mz_val <= max_vis_mz
            push!(target_masses, mz_val)
        end
    end
    
    total_ions = length(target_masses)
    println("Found $total_ions ions within training range ($(min_vis_mz) to $(max_vis_mz) m/z) above threshold.")
    if total_ions == 0
        println("Warning: No peaks met the criteria. Check your threshold or mask alignment.")
        return
    end

    # 5. MULTI-CHANNEL BATCH EXTRACTION (Single pass over binary data)
    println("Step 4: Multi-channel Batch Extraction (Optimized I/O)...")
    slice_dict = MSI_src.get_multiple_mz_slices(data, target_masses, tolerance, mask_path=mask_path)

    # 6. THREADED PROCESSING & EXPORT
    println("Step 5: Threaded TrIQ Normalization & Safe 256x256 Zero-Padding...")
    mask_matrix = MSI_src.load_and_prepare_mask(mask_path, (data.image_dims[1], data.image_dims[2]))
    exported_count = Threads.Atomic{Int}(0)

    # Fixed target dimensions for the CNN
    target_size = (256, 256)

    Threads.@threads for mz_val in target_masses
        slice = get(slice_dict, mz_val, nothing)
        if slice === nothing || all(iszero, slice) continue end

        # 1. Framework-Native TrIQ (0.0 to 255.0 Scaling)
        slice_quant, _ = MSI_src.TrIQ(slice, 256, 0.98, mask_matrix=mask_matrix)
        
        # 2. Allocate an empty 256x256 background matrix
        padded_matrix = zeros(Float32, target_size)
        
        # 3. Compute structural centering offsets based on true resolution (161x106)
        h_orig, w_orig = size(slice_quant, 1), size(slice_quant, 2)
        offset_y = div(target_size[1] - h_orig, 2) + 1
        offset_x = div(target_size[2] - w_orig, 2) + 1
        
        # 4. Map the pristine, raw pixels directly into the center
        padded_matrix[offset_y:offset_y+h_orig-1, offset_x:offset_x+w_orig-1] .= slice_quant
        
        # 5. Convert safely to 0.0 - 1.0 float mapping for Gray image export
        normalized_gray = padded_matrix ./ 255.0
        
        # 6. Save the unblurred, structurally perfect map
        filename = "ion_$(round(mz_val, digits=2)).png"
        save(joinpath(output_folder, filename), Gray.(normalized_gray))
        
        Threads.atomic_add!(exported_count, 1)
    end
    
    println("\nSuccess: $(exported_count[]) target image maps exported to $output_folder.")
end

@time run_production_pipeline()