peacoqc-rs 0.2.4

//! QC plotting module for PeacoQC
//!
//! This module provides functionality to create QC plots similar to the R PeacoQC package.
//! It generates:
//! - Time vs events/second plot
//! - Signal value vs cell event plots for each QC'd channel with highlighted unstable regions

use crate::PeacoQCData;
use crate::error::{PeacoQCError, Result};
use crate::qc::peacoqc::PeacoQCResult;
use plotters::prelude::*;
use plotters::style::{BLACK, RGBAColor, RGBColor, WHITE};
use std::path::Path;

/// Default grid line color (light gray, supports reading without dominating)
const GRID_LINE_COLOR: RGBColor = RGBColor(218, 218, 218);

/// Dash length in pixels for MAD threshold lines (plot and legend)
const MAD_DASH_LEN: u32 = 14;
/// Gap in pixels between dashes for MAD threshold lines (plot and legend)
const MAD_DASH_GAP: u32 = 3;

/// Semi-transparent white for legend background (readable over plot)
const LEGEND_BG_COLOR: RGBAColor = RGBAColor(255, 255, 255, 0.85);

/// Configuration for QC plots
#[derive(Debug, Clone)]
pub struct QCPlotConfig {
    /// Output image width in pixels
    pub width: u32,

    /// Output image height in pixels
    pub height: u32,

    /// Number of columns in the plot grid
    pub n_cols: usize,

    /// Number of rows in the plot grid
    pub n_rows: usize,

    /// Color for unstable regions (RGBA)
    pub unstable_color: RGBColor,

    /// Color for good data points
    pub good_color: RGBColor,

    /// Color for bad (unstable) data points
    pub bad_color: RGBColor,

    /// Color for median line
    pub median_color: RGBColor,

    /// Color for smoothed spline line
    pub smoothed_spline_color: RGBColor,

    /// Color for MAD threshold lines
    pub mad_threshold_color: RGBColor,

    /// Show smoothed spline and MAD threshold lines (default: true)
    pub show_spline_and_mad: bool,

    /// Show bin boundaries (gray vertical lines, default: false)
    pub show_bin_boundaries: bool,

    /// Font size for axis labels (description); default 18
    pub axis_label_size: u32,

    /// Font size for tick labels; default 15 (one step down from axis labels)
    pub tick_label_size: u32,

    /// Font size for legend text; default 17
    pub legend_font_size: u32,

    /// Font size for plot title/caption; default 22
    pub caption_font_size: u32,

    /// Font family for all text (e.g. "sans-serif", "serif"); None = "sans-serif"
    pub font_family: Option<String>,

    /// Background color; None = white (light theme)
    pub background_color: Option<RGBColor>,

    /// Foreground color for text, axes, grid; None = black (for contrast on light background)
    pub foreground_color: Option<RGBColor>,

    /// Alpha for scatter points (0.0–1.0) to reduce overplotting; None = opaque
    pub scatter_alpha: Option<f32>,
}

impl Default for QCPlotConfig {
    fn default() -> Self {
        Self {
            width: 2400,
            height: 1800,
            n_cols: 4,
            n_rows: 6,
            unstable_color: RGBColor(200, 150, 255), // Light purple
            good_color: RGBColor(128, 128, 128),     // Grey
            bad_color: RGBColor(200, 50, 50),        // Red for bad events
            median_color: RGBColor(0, 0, 0),         // Black
            smoothed_spline_color: RGBColor(0, 0, 255), // Blue (distinct from red bad events)
            mad_threshold_color: RGBColor(0, 200, 80), // Bright green for MAD bounds
            show_spline_and_mad: true,               // Enabled by default
            show_bin_boundaries: false,              // Disabled by default
            axis_label_size: 20,
            tick_label_size: 17,
            legend_font_size: 17,
            caption_font_size: 22,
            font_family: None,
            background_color: None,
            foreground_color: None,
            scatter_alpha: Some(0.5), // Slight transparency to show density
        }
    }
}

/// Find the time channel name
fn find_time_channel<T: PeacoQCData>(fcs: &T) -> Option<String> {
    fcs.channel_names().into_iter().find(|name| {
        let upper = name.to_uppercase();
        upper.contains("TIME") || upper == "TIME"
    })
}

/// Calculate events per second over time
fn calculate_events_per_second<T: PeacoQCData>(
    fcs: &T,
    time_channel: &str,
    window_size: usize,
) -> Result<Vec<(f64, f64)>> {
    let time_values = fcs.get_channel_f64(time_channel)?;

    if time_values.is_empty() {
        return Err(PeacoQCError::InsufficientData { min: 1, actual: 0 });
    }

    let mut events_per_second = Vec::new();
    let mut i = 0;

    while i < time_values.len() {
        let window_end = (i + window_size).min(time_values.len());
        if window_end <= i {
            break;
        }

        let window_times: Vec<f64> = time_values[i..window_end].to_vec();
        let time_start = window_times.first().copied().unwrap_or(0.0);
        let time_end = window_times.last().copied().unwrap_or(time_start);
        let time_span = time_end - time_start;

        let rate = if time_span > 0.0 {
            (window_end - i) as f64 / time_span
        } else {
            0.0
        };

        // Use middle of window as x position
        let mid_time = (time_start + time_end) / 2.0;
        events_per_second.push((mid_time, rate));

        i = window_end;
    }

    Ok(events_per_second)
}

/// Get channel data as vector
fn get_channel_data<T: PeacoQCData>(fcs: &T, channel: &str) -> Result<Vec<f64>> {
    fcs.get_channel_f64(channel)
}

/// Calculate median value per bin for a channel
fn calculate_median_per_bin(values: &[f64], events_per_bin: usize) -> Vec<(usize, f64)> {
    let mut medians = Vec::new();
    let n_bins = (values.len() + events_per_bin - 1) / events_per_bin;

    for bin_idx in 0..n_bins {
        let start = bin_idx * events_per_bin;
        let end = ((bin_idx + 1) * events_per_bin).min(values.len());

        if start < values.len() {
            let bin_values: Vec<f64> = values[start..end].to_vec();
            if !bin_values.is_empty() {
                let mut sorted = bin_values.clone();
                sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                let median = if sorted.len() % 2 == 0 {
                    (sorted[sorted.len() / 2 - 1] + sorted[sorted.len() / 2]) / 2.0
                } else {
                    sorted[sorted.len() / 2]
                };
                medians.push((bin_idx, median));
            }
        }
    }

    medians
}

/// Calculate grid dimensions for a given number of plots
/// Returns (n_rows, n_cols) that is relatively square and can fit all plots
fn calculate_grid_dimensions(n_plots: usize) -> (usize, usize) {
    if n_plots == 0 {
        return (1, 1);
    }

    // Start with a 1x1 grid
    let mut n_rows = 1;
    let mut n_cols = 1;

    // Alternate incrementing rows and cols until we have enough cells
    let mut increment_rows = true;
    while n_rows * n_cols < n_plots {
        if increment_rows {
            n_rows += 1;
        } else {
            n_cols += 1;
        }
        increment_rows = !increment_rows;
    }
    (n_rows, n_cols)
}

/// Find unstable regions (ranges of cell indices where good_cells is false)
fn find_unstable_regions(good_cells: &[bool]) -> Vec<(usize, usize)> {
    let mut regions = Vec::new();
    let mut in_unstable = false;
    let mut start = 0;

    for (i, &is_good) in good_cells.iter().enumerate() {
        if !is_good {
            if !in_unstable {
                start = i;
                in_unstable = true;
            }
        } else {
            if in_unstable {
                regions.push((start, i));
                in_unstable = false;
            }
        }
    }

    // Handle case where unstable region extends to end
    if in_unstable {
        regions.push((start, good_cells.len()));
    }

    regions
}

/// Create QC plots and save to file
///
/// # Arguments
/// * `fcs` - FCS data implementing PeacoQCData
/// * `qc_result` - Result from PeacoQC analysis
/// * `output_path` - Path to save the plot image
/// * `config` - Plot configuration
/// * `plot_index` - When `Some(i)`, render only the i-th plot (0 = time, 1..n = channels). Uses full canvas for single plot.
pub fn create_qc_plots<T: PeacoQCData>(
    fcs: &T,
    qc_result: &PeacoQCResult,
    output_path: impl AsRef<Path>,
    config: QCPlotConfig,
    plot_index: Option<usize>,
) -> Result<()> {
    let output_path = output_path.as_ref();

    // Find time channel
    let time_channel = find_time_channel(fcs)
        .ok_or_else(|| PeacoQCError::ConfigError("Time channel not found".to_string()))?;

    // Get channels to plot (those that were QC'd)
    let channels: Vec<String> = qc_result.peaks.keys().cloned().collect();

    if channels.is_empty() {
        return Err(PeacoQCError::ConfigError("No channels to plot".to_string()));
    }

    // When plot_index is set, render only that one plot using full canvas
    let (_n_plots, n_rows, n_cols) = match plot_index {
        Some(idx) => {
            let total = 1 + channels.len();
            if idx >= total {
                return Err(PeacoQCError::ConfigError(format!(
                    "plot_index {} out of range (0..{})",
                    idx, total
                )));
            }
            (1, 1, 1)
        }
        None => {
            let n = 1 + channels.len();
            let (r, c) = calculate_grid_dimensions(n);
            (n, r, c)
        }
    };

    let bg = config.background_color.unwrap_or(WHITE);
    let fg = config.foreground_color.unwrap_or(BLACK);
    let font_family = config.font_family.as_deref().unwrap_or("sans-serif");

    // Create drawing area
    let root = BitMapBackend::new(output_path, (config.width, config.height)).into_drawing_area();
    root.fill(&bg)
        .map_err(|e| PeacoQCError::ExportError(format!("Failed to fill background: {:?}", e)))?;

    // Split root into subplot areas
    let subplot_areas = root.split_evenly((n_rows, n_cols));

    // Plot 1: Time vs events/second
    let draw_time = plot_index.map_or(true, |i| i == 0);
    if draw_time {
        let events_per_sec = calculate_events_per_second(fcs, &time_channel, 1000)?;

        if !events_per_sec.is_empty() {
            let x_range = events_per_sec
                .iter()
                .map(|(t, _)| *t)
                .fold((f64::INFINITY, f64::NEG_INFINITY), |(min, max), x| {
                    (min.min(x), max.max(x))
                });

            let y_range = events_per_sec
                .iter()
                .map(|(_, r)| *r)
                .fold((f64::INFINITY, f64::NEG_INFINITY), |(min, max), x| {
                    (min.min(x), max.max(x))
                });

            let x_range = if x_range.0 == x_range.1 {
                (x_range.0 - 1.0)..(x_range.1 + 1.0)
            } else {
                x_range.0..x_range.1
            };

            // Enforce minimum y-scale maximum of 1 event/sec; rates below that show as "low"
            let actual_y_max = y_range.1;
            let (y_range, y_max_is_low) = if actual_y_max < 1.0 {
                (0.0..1.0, true)
            } else {
                let yr = if y_range.0 == y_range.1 {
                    (y_range.0 - 1.0)..(y_range.1 + 1.0)
                } else {
                    y_range.0..y_range.1
                };
                (yr, false)
            };

            let subplot_area = &subplot_areas[0];

            // Create title with percentage removed
            let title_text = format!(
                "{:.3}% of the data was removed",
                qc_result.percentage_removed
            );

            let y_range_clone = y_range.clone();
            let mut chart = ChartBuilder::on(&subplot_area)
                .margin(12)
                .caption(title_text, (font_family, config.caption_font_size).into_font().color(&fg))
                .x_label_area_size(58)
                .y_label_area_size(82)
                .build_cartesian_2d(x_range.clone(), y_range_clone)
                .map_err(|e| {
                    PeacoQCError::ExportError(format!("Failed to build chart: {:?}", e))
                })?;

            let draw_result = if y_max_is_low {
                chart
                    .configure_mesh()
                    .axis_desc_style(
                        (font_family, config.axis_label_size)
                            .into_font()
                            .color(&fg),
                    )
                    .label_style(
                        (font_family, config.tick_label_size)
                            .into_font()
                            .color(&fg),
                    )
                    .light_line_style(GRID_LINE_COLOR.stroke_width(1))
                    .x_desc("Time")
                    .y_desc("Nº of cells per second")
                    .x_label_formatter(&|v: &f64| format!("{:>8.1}", v))
                    .y_label_formatter(&|v: &f64| {
                        if *v >= 0.99 {
                            " low".to_string()
                        } else {
                            format!("{:>6.2}", v)
                        }
                    })
                    .draw()
            } else {
                chart
                    .configure_mesh()
                    .axis_desc_style(
                        (font_family, config.axis_label_size)
                            .into_font()
                            .color(&fg),
                    )
                    .label_style(
                        (font_family, config.tick_label_size)
                            .into_font()
                            .color(&fg),
                    )
                    .light_line_style(GRID_LINE_COLOR.stroke_width(1))
                    .x_desc("Time")
                    .y_desc("Nº of cells per second")
                    .x_label_formatter(&|v: &f64| format!("{:>8.1}", v))
                    .y_label_formatter(&|v: &f64| format!("{:>6.2}", v))
                    .draw()
            };
            draw_result
                .map_err(|e| PeacoQCError::ExportError(format!("Failed to draw mesh: {:?}", e)))?;

            // Highlight unstable regions on time plot
            let unstable_regions = find_unstable_regions(&qc_result.good_cells);
            let time_values = get_channel_data(fcs, &time_channel)?;

            for (start_idx, end_idx) in unstable_regions {
                if start_idx < time_values.len() && end_idx <= time_values.len() {
                    let start_time = time_values[start_idx];
                    let end_time = time_values[(end_idx - 1).min(time_values.len() - 1)];

                    let fill_color = RGBAColor(
                        config.unstable_color.0,
                        config.unstable_color.1,
                        config.unstable_color.2,
                        0.3,
                    );
                    chart
                        .draw_series(std::iter::once(Rectangle::new(
                            [(start_time, y_range.start), (end_time, y_range.end)],
                            fill_color.filled(),
                        )))
                        .map_err(|e| {
                            PeacoQCError::ExportError(format!("Failed to draw rectangle: {:?}", e))
                        })?;
                }
            }

            // Draw events per second line
            chart
                .draw_series(LineSeries::new(
                    events_per_sec.iter().map(|(t, r)| (*t, *r)),
                    fg.stroke_width(2),
                ))
                .map_err(|e| {
                    PeacoQCError::ExportError(format!("Failed to draw line series: {:?}", e))
                })?;

            // Legend: Removed events (shaded regions)
            let x_range_size = x_range.end - x_range.start;
            let y_range_size = y_range.end - y_range.start;
            let legend_x_start = x_range.end - (x_range_size * 0.22);
            let legend_y_start = y_range.end - (y_range_size * 0.06);
            let rect_w = x_range_size * 0.025;
            let rect_h = y_range_size * 0.04;
            let text_gap = x_range_size * 0.008;
            let pad_x = x_range_size * 0.008;
            let pad_y = y_range_size * 0.008;
            // Semi-transparent white background behind legend
            let legend_bg_left = legend_x_start - pad_x;
            let legend_bg_bottom = legend_y_start - rect_h - pad_y;
            let legend_bg_right = legend_x_start + rect_w + text_gap + x_range_size * 0.12;
            let legend_bg_top = legend_y_start + pad_y;
            chart
                .draw_series(std::iter::once(Rectangle::new(
                    [
                        (legend_bg_left, legend_bg_bottom),
                        (legend_bg_right, legend_bg_top),
                    ],
                    LEGEND_BG_COLOR.filled(),
                )))
                .map_err(|e| {
                    PeacoQCError::ExportError(format!("Failed to draw legend background: {:?}", e))
                })?;
            let fill_color = RGBAColor(
                config.unstable_color.0,
                config.unstable_color.1,
                config.unstable_color.2,
                0.5,
            );
            chart
                .draw_series(std::iter::once(Rectangle::new(
                    [
                        (legend_x_start, legend_y_start - rect_h),
                        (legend_x_start + rect_w, legend_y_start),
                    ],
                    fill_color.filled(),
                )))
                .map_err(|e| {
                    PeacoQCError::ExportError(format!("Failed to draw legend rect: {:?}", e))
                })?;
            chart
                .plotting_area()
                .draw(&Text::new(
                    "Removed events".to_string(),
                    (legend_x_start + rect_w + text_gap, legend_y_start),
                    (font_family, config.legend_font_size)
                        .into_font()
                        .color(&fg),
                ))
                .map_err(|e| {
                    PeacoQCError::ExportError(format!("Failed to draw legend text: {:?}", e))
                })?;
        }
    }

    // Plot channels: Signal value vs cell event
    let total_cells = n_rows * n_cols;
    let channel_iter: Box<dyn Iterator<Item = (usize, &String)>> = match plot_index {
        Some(i) if i >= 1 && i <= channels.len() => {
            Box::new(std::iter::once((i - 1, &channels[i - 1])))
        }
        Some(_) => Box::new(std::iter::empty()),
        None => Box::new(channels.iter().enumerate()),
    };
    for (plot_idx, channel) in channel_iter {
        let subplot_idx = if plot_index.is_some() {
            0 // Single-plot mode: use first (only) cell
        } else {
            plot_idx + 1 // +1 because first plot is time plot
        };

        if subplot_idx >= total_cells {
            break;
        }

        let channel_data = get_channel_data(fcs, channel)?;
        if channel_data.is_empty() {
            continue;
        }

        let n_events = channel_data.len();
        let cell_indices: Vec<f64> = (0..n_events).map(|i| i as f64).collect();

        // Calculate ranges
        let x_range = 0.0..(n_events as f64);
        let y_min = channel_data.iter().fold(f64::INFINITY, |a, &b| a.min(b));
        let y_max = channel_data
            .iter()
            .fold(f64::NEG_INFINITY, |a, &b| a.max(b));
        let y_range = if y_min == y_max {
            (y_min - 1.0)..(y_max + 1.0)
        } else {
            y_min..y_max
        };

        let subplot_area = &subplot_areas[subplot_idx];

        // Calculate MAD percentage for title
        let mad_pct = qc_result
            .mad_percentage
            .and_then(|_| {
                // Calculate channel-specific MAD percentage
                qc_result.peaks.get(channel).map(|_| {
                    // This is approximate - we'd need to track per-channel MAD
                    qc_result.mad_percentage.unwrap_or(0.0)
                })
            })
            .unwrap_or(0.0);

        let title = if mad_pct > 0.0 {
            format!("{} MAD {:.2}%", channel, mad_pct)
        } else {
            channel.to_string()
        };

        let mut chart = ChartBuilder::on(&subplot_area)
            .margin(12)
            .caption(&title, (font_family, config.caption_font_size).into_font().color(&fg))
            .x_label_area_size(48)
            .y_label_area_size(82)
            .build_cartesian_2d(x_range.clone(), y_range.clone())
            .map_err(|e| PeacoQCError::ExportError(format!("Failed to build chart: {:?}", e)))?;

        chart
            .configure_mesh()
            .axis_desc_style(
                (font_family, config.axis_label_size)
                    .into_font()
                    .color(&fg),
            )
            .label_style(
                (font_family, config.tick_label_size)
                    .into_font()
                    .color(&fg),
            )
            .light_line_style(GRID_LINE_COLOR.stroke_width(1))
            .x_desc("Cell index")
            .y_desc("Signal (a.u.)")
            .x_label_formatter(&|v: &f64| format!("{:>8.0}", v))
            .y_label_formatter(&|v: &f64| format!("{:>8.2}", v))
            .draw()
            .map_err(|e| PeacoQCError::ExportError(format!("Failed to draw mesh: {:?}", e)))?;

        // Highlight unstable regions
        let unstable_regions = find_unstable_regions(&qc_result.good_cells);
        for (start_idx, end_idx) in unstable_regions {
            if start_idx < n_events {
                let start_cell = start_idx as f64;
                let end_cell = (end_idx.min(n_events)) as f64;

                let fill_color = RGBAColor(
                    config.unstable_color.0,
                    config.unstable_color.1,
                    config.unstable_color.2,
                    0.3,
                );
                chart
                    .draw_series(std::iter::once(Rectangle::new(
                        [(start_cell, y_range.start), (end_cell, y_range.end)],
                        fill_color.filled(),
                    )))
                    .map_err(|e| {
                        PeacoQCError::ExportError(format!("Failed to draw rectangle: {:?}", e))
                    })?;
            }
        }

        // Draw scatter plot: bad events first (red), then good (grey), same sampling for performance
        let sample_size = 10000.min(n_events);
        let step = (n_events / sample_size.max(1)).max(1);
        let mut good_points = Vec::new();
        let mut bad_points = Vec::new();

        for i in (0..n_events).step_by(step) {
            let pt = (cell_indices[i], channel_data[i]);
            if qc_result.good_cells[i] {
                good_points.push(pt);
            } else {
                bad_points.push(pt);
            }
        }

        let alpha: f64 = config.scatter_alpha.unwrap_or(1.0) as f64;
        let alpha = alpha.clamp(0.0, 1.0);
        let use_alpha = alpha < 1.0;

        if !bad_points.is_empty() {
            chart
                .draw_series(bad_points.iter().map(|(x, y)| {
                    Circle::new(
                        (*x, *y),
                        1,
                        if use_alpha {
                            RGBAColor(
                                config.bad_color.0,
                                config.bad_color.1,
                                config.bad_color.2,
                                alpha,
                            )
                            .filled()
                        } else {
                            config.bad_color.filled()
                        },
                    )
                }))
                .map_err(|e| {
                    PeacoQCError::ExportError(format!("Failed to draw bad-event circles: {:?}", e))
                })?;
        }
        if !good_points.is_empty() {
            chart
                .draw_series(good_points.iter().map(|(x, y)| {
                    Circle::new(
                        (*x, *y),
                        1,
                        if use_alpha {
                            RGBAColor(
                                config.good_color.0,
                                config.good_color.1,
                                config.good_color.2,
                                alpha,
                            )
                            .filled()
                        } else {
                            config.good_color.filled()
                        },
                    )
                }))
                .map_err(|e| {
                    PeacoQCError::ExportError(format!("Failed to draw circles: {:?}", e))
                })?;
        }

        // Draw median line per bin
        let medians = calculate_median_per_bin(&channel_data, qc_result.events_per_bin);
        if !medians.is_empty() {
            let median_points: Vec<(f64, f64)> = medians
                .iter()
                .map(|(bin_idx, median)| {
                    let cell_idx = (*bin_idx * qc_result.events_per_bin) as f64;
                    (cell_idx, *median)
                })
                .collect();

            chart
                .draw_series(LineSeries::new(
                    median_points.clone(),
                    config.median_color.stroke_width(2),
                ))
                .map_err(|e| {
                    PeacoQCError::ExportError(format!("Failed to draw median line: {:?}", e))
                })?;

            // Draw bin boundaries (if enabled)
            if config.show_bin_boundaries {
                let n_bins = (n_events + qc_result.events_per_bin - 1) / qc_result.events_per_bin;
                let boundary_color = RGBColor(200, 200, 200);
                for bin_idx in 0..=n_bins {
                    let cell_idx = (bin_idx * qc_result.events_per_bin) as f64;
                    if cell_idx <= n_events as f64 {
                        chart
                            .draw_series(std::iter::once(plotters::prelude::PathElement::new(
                                vec![(cell_idx, y_range.start), (cell_idx, y_range.end)],
                                boundary_color.stroke_width(1),
                            )))
                            .map_err(|e| {
                                PeacoQCError::ExportError(format!(
                                    "Failed to draw bin boundary: {:?}",
                                    e
                                ))
                            })?;
                    }
                }
            }

            // Draw smoothed spline and MAD threshold lines (if enabled)
            if config.show_spline_and_mad && medians.len() >= 3 {
                let bin_medians: Vec<f64> = medians.iter().map(|(_, m)| *m).collect();
                let bin_indices: Vec<f64> = medians.iter().map(|(i, _)| *i as f64).collect();

                // Apply smoothing (using default spar=0.5)
                if let Ok(smoothed) =
                    crate::stats::spline::smooth_spline(&bin_indices, &bin_medians, 0.5)
                {
                    let smoothed_points: Vec<(f64, f64)> = smoothed
                        .iter()
                        .enumerate()
                        .map(|(i, &y)| {
                            let cell_idx = (i * qc_result.events_per_bin) as f64;
                            (cell_idx, y)
                        })
                        .collect();

                    chart
                        .draw_series(LineSeries::new(
                            smoothed_points.clone(),
                            config.smoothed_spline_color.stroke_width(2),
                        ))
                        .map_err(|e| {
                            PeacoQCError::ExportError(format!(
                                "Failed to draw smoothed spline: {:?}",
                                e
                            ))
                        })?;

                    // Compute and draw MAD threshold lines
                    if let Ok((median, mad)) =
                        crate::stats::median_mad::median_mad_scaled(&smoothed)
                    {
                        let mad_threshold = 6.0; // Default MAD threshold
                        let upper_threshold = median + mad_threshold * mad;
                        let lower_threshold = median - mad_threshold * mad;

                        // Draw threshold lines (dashed: longer dashes, narrower gaps)
                        let threshold_points_upper: Vec<(f64, f64)> =
                            vec![(0.0, upper_threshold), (n_events as f64, upper_threshold)];
                        let threshold_points_lower: Vec<(f64, f64)> =
                            vec![(0.0, lower_threshold), (n_events as f64, lower_threshold)];
                        let mad_style = config.mad_threshold_color.stroke_width(2);

                        chart
                            .draw_series(std::iter::once(
                                plotters::element::DashedPathElement::new(
                                    threshold_points_upper,
                                    MAD_DASH_LEN,
                                    MAD_DASH_GAP,
                                    mad_style.clone(),
                                ),
                            ))
                            .map_err(|e| {
                                PeacoQCError::ExportError(format!(
                                    "Failed to draw upper threshold: {:?}",
                                    e
                                ))
                            })?;

                        chart
                            .draw_series(std::iter::once(
                                plotters::element::DashedPathElement::new(
                                    threshold_points_lower,
                                    MAD_DASH_LEN,
                                    MAD_DASH_GAP,
                                    mad_style,
                                ),
                            ))
                            .map_err(|e| {
                                PeacoQCError::ExportError(format!(
                                    "Failed to draw lower threshold: {:?}",
                                    e
                                ))
                            })?;
                    }
                }
            }

            // Draw legend in top-right corner: first "Removed events" (rect), then line items
            let legend_rects: Vec<(&str, RGBColor)> =
                vec![("Removed events", config.unstable_color)];
            let mut legend_items: Vec<(&str, RGBColor, u32)> =
                vec![("Median", config.median_color, 2)];

            if config.show_spline_and_mad {
                legend_items.push(("Spline", config.smoothed_spline_color, 2));
                legend_items.push(("MAD ±6", config.mad_threshold_color, 2));
            }

            let x_range_size = x_range.end - x_range.start;
            let y_range_size = y_range.end - y_range.start;
            let legend_margin_right_pct = 0.10;
            let legend_margin_top_pct = 0.02;
            let legend_x_start = x_range.end - (x_range_size * legend_margin_right_pct);
            let legend_y_step = y_range_size * 0.032;
            let line_length = x_range_size * 0.035;
            let text_gap = x_range_size * 0.008;
            let rect_w = x_range_size * 0.02;
            let rect_h = y_range_size * 0.025;
            let legend_initial_y = y_range.end - (y_range_size * legend_margin_top_pct);
            let n_legend_rows = legend_rects.len() + legend_items.len();
            let pad_x = x_range_size * 0.006;
            let pad_y = y_range_size * 0.006;
            // Semi-transparent white background behind full legend
            let legend_bg_left = legend_x_start - pad_x;
            let legend_bg_right =
                legend_x_start + line_length + text_gap + x_range_size * 0.10 + pad_x;
            let legend_bg_bottom = legend_initial_y
                - rect_h
                - (n_legend_rows.saturating_sub(1) as f64 * legend_y_step)
                - pad_y;
            let legend_bg_top = legend_initial_y + pad_y;
            chart
                .draw_series(std::iter::once(Rectangle::new(
                    [
                        (legend_bg_left, legend_bg_bottom),
                        (legend_bg_right, legend_bg_top),
                    ],
                    LEGEND_BG_COLOR.filled(),
                )))
                .map_err(|e| {
                    PeacoQCError::ExportError(format!("Failed to draw legend background: {:?}", e))
                })?;

            let mut legend_y = legend_initial_y;

            for (label, color) in &legend_rects {
                let fill_color = RGBAColor(color.0, color.1, color.2, 0.5);
                chart
                    .draw_series(std::iter::once(Rectangle::new(
                        [
                            (legend_x_start, legend_y - rect_h),
                            (legend_x_start + rect_w, legend_y),
                        ],
                        fill_color.filled(),
                    )))
                    .map_err(|e| {
                        PeacoQCError::ExportError(format!("Failed to draw legend rect: {:?}", e))
                    })?;
                chart
                    .plotting_area()
                    .draw(&Text::new(
                        (*label).to_string(),
                        (legend_x_start + rect_w + text_gap, legend_y),
                        (font_family, config.legend_font_size)
                            .into_font()
                            .color(&fg),
                    ))
                    .map_err(|e| {
                        PeacoQCError::ExportError(format!("Failed to draw legend text: {:?}", e))
                    })?;
                legend_y -= legend_y_step;
            }

            for (label, color, stroke_width) in &legend_items {
                let line_pts = vec![
                    (legend_x_start, legend_y),
                    (legend_x_start + line_length, legend_y),
                ];
                let stroke = color.stroke_width(*stroke_width);
                if *label == "MAD ±6" {
                    chart
                        .draw_series(std::iter::once(plotters::element::DashedPathElement::new(
                            line_pts,
                            MAD_DASH_LEN,
                            MAD_DASH_GAP,
                            stroke,
                        )))
                        .map_err(|e| {
                            PeacoQCError::ExportError(format!(
                                "Failed to draw legend line: {:?}",
                                e
                            ))
                        })?;
                } else {
                    chart
                        .draw_series(std::iter::once(plotters::prelude::PathElement::new(
                            line_pts, stroke,
                        )))
                        .map_err(|e| {
                            PeacoQCError::ExportError(format!(
                                "Failed to draw legend line: {:?}",
                                e
                            ))
                        })?;
                }
                chart
                    .plotting_area()
                    .draw(&Text::new(
                        label.to_string(),
                        (legend_x_start + line_length + text_gap, legend_y),
                        (font_family, config.legend_font_size)
                            .into_font()
                            .color(&fg),
                    ))
                    .map_err(|e| {
                        PeacoQCError::ExportError(format!("Failed to draw legend text: {:?}", e))
                    })?;
                legend_y -= legend_y_step;
            }
        }
    }

    root.present()
        .map_err(|e| PeacoQCError::ExportError(format!("Failed to present plot: {:?}", e)))?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_find_unstable_regions() {
        let good_cells = vec![false, false, true, true, true, false, true, true];
        let regions = find_unstable_regions(&good_cells);
        assert_eq!(regions, vec![(0, 2), (5, 6)]);
    }

    #[test]
    fn test_calculate_median_per_bin() {
        let values = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let medians = calculate_median_per_bin(&values, 2);
        assert_eq!(medians.len(), 4);
        assert_eq!(medians[0], (0, 1.5));
        assert_eq!(medians[1], (1, 3.5));
    }

    #[test]
    fn test_calculate_grid_dimensions() {
        // Test various plot counts
        assert_eq!(calculate_grid_dimensions(1), (1, 1));
        assert_eq!(calculate_grid_dimensions(4), (2, 2));
        assert!(calculate_grid_dimensions(5) == (3, 2) || calculate_grid_dimensions(5) == (2, 3)); // or (2, 3) - alternates
        assert_eq!(calculate_grid_dimensions(9), (3, 3));
        assert_eq!(calculate_grid_dimensions(25), (5, 5));
        assert!(calculate_grid_dimensions(30) == (6, 5) || calculate_grid_dimensions(30) == (5, 6)); // or (5, 6) - alternates
        assert_eq!(calculate_grid_dimensions(36), (6, 6));

        // Verify the grid can fit all plots
        let (rows, cols) = calculate_grid_dimensions(25);
        assert!(rows * cols >= 25);
        assert_eq!(rows, 5);
        assert_eq!(cols, 5);

        let (rows, cols) = calculate_grid_dimensions(30);
        assert!(rows * cols >= 30);

        let (rows, cols) = calculate_grid_dimensions(24);
        assert!(rows * cols >= 24);
    }
}