nanocov 0.1.0

// src/plotting/mod.rs
// Plotting module for nanocov - handles coverage visualization and statistics display

mod stats;
mod themes;
mod utils;

use stats::{CoverageStats, calculate_coverage_stats, calculate_per_base_stats};
use themes::{CATPPUCCIN_FRAPPE, CATPPUCCIN_LATTE, ColorTheme, GRUVBOX_LIGHT, NORD};
use utils::{blend_colors, format_number};

// Default theme (can be overridden via CLI)
pub static mut CURRENT_THEME: &ColorTheme = &CATPPUCCIN_LATTE;

// Shorthand for accessing current theme colors
#[inline]
fn theme() -> &'static ColorTheme {
    unsafe { CURRENT_THEME }
}

use plotters::coord::Shift;
use plotters::prelude::*;
use plotters::style::FontTransform;
use std::collections::HashMap;

/// Plots per-base coverage as a bar graph for a single chromosome.
///
/// This function takes a chromosome name, a map of base positions to coverage counts,
/// and an output path for the PNG file. It sorts the positions, determines the y-axis
/// scaling based on the data, and draws a bar for each covered base. The plot is saved
/// as a PNG file using the plotters crate.
///
/// # Arguments
/// * `chrom` - Chromosome name (for labeling the plot)
/// * `coverage` - Map from base position (u32, starting at 0) to coverage count (u32)
/// * `output_path` - Path to save the output PNG file
/// * `read_stats` - Optional read statistics to display in the sidebar
/// * `show_zero_regions` - Whether to show regions with zero coverage
///
/// # Details
/// - The y-axis is automatically scaled based on the data range
/// - Each bar spans its bin width, with color gradient based on coverage
/// - The plot includes detailed statistics panels and legends
/// - The function is robust to empty or sparse data
use crate::utils::ReadStats;

/// Set the global color theme for all plots
///
/// # Arguments
/// * `theme_name` - One of: "latte", "frappe", "nord", "gruvbox"
#[inline]
pub fn set_theme(theme_name: &str) {
    unsafe {
        CURRENT_THEME = match theme_name.to_lowercase().as_str() {
            "frappe" => &CATPPUCCIN_FRAPPE,
            "nord" => &NORD,
            "gruvbox" => &GRUVBOX_LIGHT,
            _ => &CATPPUCCIN_LATTE, // Default to latte
        };
    }
}

/// Simple entry point for plotting coverage with automatic range determination
/// This is externally used by the io module
#[inline]
#[allow(dead_code)] // Used by io module
#[allow(clippy::too_many_arguments)]
pub fn plot_per_base_coverage(
    chrom: &str,
    coverage: &HashMap<u32, u32>,
    output_path: &str,
    read_stats: Option<&ReadStats>,
    show_zero_regions: bool,
    use_log_scale: bool,
    plot_bin_size: Option<u32>,
    title_prefix: Option<&str>,
) -> Result<(), Box<dyn std::error::Error>> {
    // Use the min/max of coverage as the range
    let mut positions: Vec<u32> = coverage.keys().copied().collect();
    positions.sort_unstable();

    if positions.is_empty() {
        eprintln!("Warning: No coverage data for chromosome {chrom}");
        return Ok(());
    }

    let min_x = positions.iter().min().copied().unwrap_or(0);
    let max_x = positions.iter().max().copied().unwrap_or(0);
    plot_per_base_coverage_with_range(
        chrom,
        coverage,
        output_path,
        min_x,
        max_x,
        read_stats,
        show_zero_regions,
        use_log_scale,
        plot_bin_size,
        title_prefix,
    )
}

#[allow(clippy::too_many_arguments)]
pub fn plot_per_base_coverage_with_range(
    chrom: &str,
    coverage: &HashMap<u32, u32>,
    output_path: &str,
    plot_start: u32,
    plot_end: u32,
    read_stats: Option<&ReadStats>,
    show_zero_regions: bool,
    use_log_scale: bool,
    plot_bin_size: Option<u32>,
    title_prefix: Option<&str>,
) -> Result<(), Box<dyn std::error::Error>> {
    // Adaptive binning logic for large ranges
    let range = plot_end.saturating_sub(plot_start);
    let bin_size = if let Some(bin_size) = plot_bin_size {
        std::cmp::max(1, bin_size)
    } else if range > 10_000_000 {
        10_000 // 10kb bins for large regions (>10Mb)
    } else if range > 1_000_000 {
        1_000 // 1kb bins for medium regions (>1Mb)
    } else if range > 100_000 {
        100 // 100bp bins for smaller regions (>100kb)
    } else if range > 10_000 {
        10 // 10bp bins for tiny regions (>10kb)
    } else {
        1 // No binning for very small regions
    };

    let mut binned: std::collections::BTreeMap<u32, u64> = std::collections::BTreeMap::new();
    for (&pos, &count) in coverage.iter() {
        if pos < plot_start || pos > plot_end {
            continue;
        }
        let bin = ((pos - plot_start) / bin_size) * bin_size + plot_start;
        let entry = binned.entry(bin).or_insert(0);
        *entry += count as u64;
    }
    let chrom_points: Vec<(i64, f64)> = binned
        .iter()
        .map(|(&bin, &sum)| {
            let bin_end = bin.saturating_add(bin_size.saturating_sub(1)).min(plot_end);
            let bin_len = (bin_end - bin + 1) as u64;
            let avg = if bin_len > 0 {
                sum as f64 / bin_len as f64
            } else {
                0.0
            };
            (bin as i64, avg)
        })
        .collect();

    // Calculate and add zero coverage regions if requested
    let chrom_points = if show_zero_regions && !chrom_points.is_empty() {
        // Create a more complete dataset with zero regions
        let mut full_points = Vec::new();
        let mut last_pos = plot_start;

        for &(pos, coverage) in &chrom_points {
            // Add zeros for any gaps
            if pos as u32 > last_pos + bin_size {
                // Add zero points for missing regions
                let mut gap_pos = last_pos + bin_size;
                while gap_pos < pos as u32 {
                    full_points.push((gap_pos as i64, 0.0));
                    gap_pos += bin_size;
                }
            }
            full_points.push((pos, coverage));
            last_pos = pos as u32;
        }
        full_points
    } else {
        chrom_points
    };

    // Determine the maximum y (coverage) value with smarter scaling
    let y_max = chrom_points.iter().map(|&(_, y)| y).fold(0.0, f64::max);
    let y_max = if y_max < 3.0 {
        3.0 // Minimum scale for very low coverage
    } else if y_max < 10.0 {
        (y_max * 1.2).ceil() // Slightly larger for clearer view of small values
    } else if y_max < 100.0 {
        (y_max * 1.1).ceil() // 10% headroom for medium values
    } else {
        (y_max * 1.05).ceil() // 5% headroom for large values
    };

    // Debug output
    eprintln!("[DEBUG] Max coverage for {chrom}: {y_max}");

    // Set up the drawing area with BitMapBackend (PNG output)
    let root = BitMapBackend::new(output_path, (2200, 1000)).into_drawing_area();

    root.fill(&theme().base)?;

    // Split horizontally: left panel (stats), middle panel (main plot), right panel (scale bars)
    let (left_panel, remaining) = root.split_horizontally(400);
    let (plot_area, right_panel) = remaining.split_horizontally(1400);

    // Format bin size for display (kb or Mb for large bins)
    let bin_size_label = if bin_size >= 1_000_000 {
        format!("{} Mb", bin_size / 1_000_000)
    } else if bin_size >= 1_000 {
        format!("{} kb", bin_size / 1_000)
    } else {
        format!("{bin_size} bp")
    };

    // Handle logarithmic and linear scaling separately due to incompatible chart types
    let title_prefix = title_prefix
        .filter(|p| !p.is_empty())
        .map(|p| format!("{p} - "))
        .unwrap_or_default();

    if use_log_scale {
        // Find minimum non-zero value for log scale
        let min_coverage = chrom_points
            .iter()
            .map(|&(_, y)| y)
            .filter(|&v| v > 0.0)
            .fold(y_max, f64::min)
            .max(0.1); // Minimum of 0.1 for log scale

        let mut chart = ChartBuilder::on(&plot_area)
            .x_label_area_size(10)
            .y_label_area_size(10)
            .set_label_area_size(LabelAreaPosition::Left, 75)
            .set_label_area_size(LabelAreaPosition::Bottom, 50)
            .margin(50)
            .caption(
                format!(
                    "{title_prefix}Chromosome {chrom} Coverage (bin: {bin_size_label}, log scale)"
                ),
                ("sans-serif", 40).into_font().color(&theme().text),
            )
            .build_cartesian_2d(
                plot_start as i64..plot_end as i64,
                (min_coverage..y_max * 1.1).log_scale(),
            )?;

        // Configure the mesh (axes, grid, and labels) for log scale
        chart
            .configure_mesh()
            .x_desc("Chromosome Position (Mb)")
            .y_desc("Coverage (log scale)")
            .axis_desc_style(("sans-serif", 25).into_font().color(&theme().text))
            .x_label_formatter(&|x| format!("{:.2}", (*x as f64) / 1_000_000.0)) // Format x axis in Mb
            .x_labels(20) // Increase number of x-axis ticks
            .x_label_style(("sans-serif", 18).into_font().color(&theme().text))
            .y_label_style(("sans-serif", 18).into_font().color(&theme().text))
            .y_label_formatter(&|y| format!("{y:.1}")) // Clean y-axis formatting
            .light_line_style(RGBAColor(100, 100, 100, 0.3)) // Subtle grid lines
            .bold_line_style(RGBAColor(100, 100, 100, 0.5)) // Bolder major grid lines
            .axis_style(theme().text)
            .draw()?;

        // Draw each bar as a filled rectangle with gradient color based on coverage (log scale)
        let plot_end_i64 = plot_end as i64;
        chart.draw_series(chrom_points.windows(2).filter_map(|w| {
            let (x, y) = w[0];

            // Skip zero coverage in log scale
            if y <= 0.0 {
                return None;
            }

            let bar_end = (x + bin_size as i64).min(plot_end_i64);

            // Color gradient based on coverage level
            let fill_color = if y <= y_max * 0.3 {
                // Low coverage: blend from low coverage color to main color
                let blend_factor = y / (y_max * 0.3);
                let blend_factor = blend_factor.clamp(0.0, 1.0);
                RGBColor(
                    ((theme().low.0 as f64) * (1.0 - blend_factor)
                        + (theme().primary.0 as f64) * blend_factor) as u8,
                    ((theme().low.1 as f64) * (1.0 - blend_factor)
                        + (theme().primary.1 as f64) * blend_factor) as u8,
                    ((theme().low.2 as f64) * (1.0 - blend_factor)
                        + (theme().primary.2 as f64) * blend_factor) as u8,
                )
            } else if y >= y_max * 0.7 {
                // High coverage: blend from main color to high coverage color
                let blend_factor = (y - y_max * 0.7) / (y_max * 0.3);
                let blend_factor = blend_factor.clamp(0.0, 1.0);
                RGBColor(
                    ((theme().primary.0 as f64) * (1.0 - blend_factor)
                        + (theme().high.0 as f64) * blend_factor) as u8,
                    ((theme().primary.1 as f64) * (1.0 - blend_factor)
                        + (theme().high.1 as f64) * blend_factor) as u8,
                    ((theme().primary.2 as f64) * (1.0 - blend_factor)
                        + (theme().high.2 as f64) * blend_factor) as u8,
                )
            } else {
                // Medium coverage: use main color
                theme().primary
            };

            // For log scale, use minimum coverage as baseline instead of 0
            Some(Rectangle::new(
                [(x, min_coverage), (bar_end, y)],
                fill_color.filled(),
            ))
        }))?;

        // Draw the last bar if only one point or for the last position (log scale)
        if let Some(&(x, y)) = chrom_points.last() {
            // Skip zero coverage in log scale
            if y > 0.0 {
                // Apply same color logic for last bar
                let fill_color = if y <= y_max * 0.3 {
                    let blend_factor = y / (y_max * 0.3);
                    let blend_factor = blend_factor.clamp(0.0, 1.0);
                    RGBColor(
                        ((theme().low.0 as f64) * (1.0 - blend_factor)
                            + (theme().primary.0 as f64) * blend_factor)
                            as u8,
                        ((theme().low.1 as f64) * (1.0 - blend_factor)
                            + (theme().primary.1 as f64) * blend_factor)
                            as u8,
                        ((theme().low.2 as f64) * (1.0 - blend_factor)
                            + (theme().primary.2 as f64) * blend_factor)
                            as u8,
                    )
                } else if y >= y_max * 0.7 {
                    let blend_factor = (y - y_max * 0.7) / (y_max * 0.3);
                    let blend_factor = blend_factor.clamp(0.0, 1.0);
                    RGBColor(
                        ((theme().primary.0 as f64) * (1.0 - blend_factor)
                            + (theme().high.0 as f64) * blend_factor) as u8,
                        ((theme().primary.1 as f64) * (1.0 - blend_factor)
                            + (theme().high.1 as f64) * blend_factor) as u8,
                        ((theme().primary.2 as f64) * (1.0 - blend_factor)
                            + (theme().high.2 as f64) * blend_factor) as u8,
                    )
                } else {
                    theme().primary
                };

                let bar_end = (x + bin_size as i64).min(plot_end as i64);
                chart.draw_series(std::iter::once(Rectangle::new(
                    [(x, min_coverage), (bar_end, y)],
                    fill_color.filled(),
                )))?;
            }
        }

        // Add threshold line if average coverage is available (log scale)
        let y_vals: Vec<f64> = chrom_points
            .iter()
            .map(|&(_, y)| y)
            .filter(|&y| y > 0.0)
            .collect();
        if !y_vals.is_empty() {
            let mean = y_vals.iter().sum::<f64>() / y_vals.len() as f64;
            if mean >= min_coverage {
                chart.draw_series(std::iter::once(PathElement::new(
                    vec![(plot_start as i64, mean), (plot_end as i64, mean)],
                    theme().accent.stroke_width(2),
                )))?;
            }
        }
    } else {
        // Linear scale (original implementation)
        let mut chart = ChartBuilder::on(&plot_area)
            .x_label_area_size(10)
            .y_label_area_size(10)
            .set_label_area_size(LabelAreaPosition::Left, 75)
            .set_label_area_size(LabelAreaPosition::Bottom, 50)
            .margin(50)
            .caption(
                format!("{title_prefix}Chromosome {chrom} Coverage (bin: {bin_size_label})"),
                ("sans-serif", 40).into_font().color(&theme().text),
            )
            .build_cartesian_2d(plot_start as i64..plot_end as i64, 0f64..y_max)?;

        // Configure the mesh (axes, grid, and labels)
        chart
            .configure_mesh()
            .x_desc("Chromosome Position (Mb)")
            .y_desc("Coverage")
            .axis_desc_style(("sans-serif", 25).into_font().color(&theme().text))
            .x_label_formatter(&|x| format!("{:.2}", (*x as f64) / 1_000_000.0)) // Format x axis in Mb
            .x_labels(20) // Increase number of x-axis ticks
            .x_label_style(("sans-serif", 18).into_font().color(&theme().text))
            .y_label_style(("sans-serif", 18).into_font().color(&theme().text))
            .y_label_formatter(&|y| format!("{y:.1}")) // Clean y-axis formatting
            .light_line_style(RGBAColor(100, 100, 100, 0.3)) // Subtle grid lines
            .bold_line_style(RGBAColor(100, 100, 100, 0.5)) // Bolder major grid lines
            .axis_style(theme().text)
            .draw()?;

        // Draw each bar as a filled rectangle with gradient color based on coverage
        let plot_end_i64 = plot_end as i64;
        chart.draw_series(chrom_points.windows(2).map(|w| {
            let (x, y) = w[0];
            let bar_end = (x + bin_size as i64).min(plot_end_i64);

            // Color gradient based on coverage level
            let fill_color = if y <= y_max * 0.3 {
                // Low coverage: blend from low coverage color to main color
                let blend_factor = y / (y_max * 0.3);
                let blend_factor = blend_factor.clamp(0.0, 1.0);
                RGBColor(
                    ((theme().low.0 as f64) * (1.0 - blend_factor)
                        + (theme().primary.0 as f64) * blend_factor) as u8,
                    ((theme().low.1 as f64) * (1.0 - blend_factor)
                        + (theme().primary.1 as f64) * blend_factor) as u8,
                    ((theme().low.2 as f64) * (1.0 - blend_factor)
                        + (theme().primary.2 as f64) * blend_factor) as u8,
                )
            } else if y >= y_max * 0.7 {
                // High coverage: blend from main color to high coverage color
                let blend_factor = (y - y_max * 0.7) / (y_max * 0.3);
                let blend_factor = blend_factor.clamp(0.0, 1.0);
                RGBColor(
                    ((theme().primary.0 as f64) * (1.0 - blend_factor)
                        + (theme().high.0 as f64) * blend_factor) as u8,
                    ((theme().primary.1 as f64) * (1.0 - blend_factor)
                        + (theme().high.1 as f64) * blend_factor) as u8,
                    ((theme().primary.2 as f64) * (1.0 - blend_factor)
                        + (theme().high.2 as f64) * blend_factor) as u8,
                )
            } else {
                // Medium coverage: use main color
                theme().primary
            };

            Rectangle::new([(x, 0.0), (bar_end, y)], fill_color.filled())
        }))?;

        // Draw the last bar if only one point or for the last position
        if let Some(&(x, y)) = chrom_points.last() {
            // Apply same color logic for last bar
            let fill_color = if y <= y_max * 0.3 {
                let blend_factor = y / (y_max * 0.3);
                let blend_factor = blend_factor.clamp(0.0, 1.0);
                RGBColor(
                    ((theme().low.0 as f64) * (1.0 - blend_factor)
                        + (theme().primary.0 as f64) * blend_factor) as u8,
                    ((theme().low.1 as f64) * (1.0 - blend_factor)
                        + (theme().primary.1 as f64) * blend_factor) as u8,
                    ((theme().low.2 as f64) * (1.0 - blend_factor)
                        + (theme().primary.2 as f64) * blend_factor) as u8,
                )
            } else if y >= y_max * 0.7 {
                let blend_factor = (y - y_max * 0.7) / (y_max * 0.3);
                let blend_factor = blend_factor.clamp(0.0, 1.0);
                RGBColor(
                    ((theme().primary.0 as f64) * (1.0 - blend_factor)
                        + (theme().high.0 as f64) * blend_factor) as u8,
                    ((theme().primary.1 as f64) * (1.0 - blend_factor)
                        + (theme().high.1 as f64) * blend_factor) as u8,
                    ((theme().primary.2 as f64) * (1.0 - blend_factor)
                        + (theme().high.2 as f64) * blend_factor) as u8,
                )
            } else {
                theme().primary
            };

            let bar_end = (x + bin_size as i64).min(plot_end as i64);
            chart.draw_series(std::iter::once(Rectangle::new(
                [(x, 0.0), (bar_end, y)],
                fill_color.filled(),
            )))?;
        }

        // Add threshold line if average coverage is available
        let y_vals: Vec<f64> = chrom_points.iter().map(|&(_, y)| y).collect();
        if !y_vals.is_empty() {
            let mean = y_vals.iter().sum::<f64>() / y_vals.len() as f64;
            chart.draw_series(std::iter::once(PathElement::new(
                vec![(plot_start as i64, mean), (plot_end as i64, mean)],
                theme().accent.stroke_width(2),
            )))?;
        }
    }

    // --- Apply annotations and legends in the right panel ---

    // Fill the right panel with the base color
    right_panel.fill(&theme().base)?;

    // Coverage scale gradient removed (intentionally blank)

    // --- Draw read stats and coverage stats panels on the left ---
    let font = ("sans-serif", 24)
        .into_font()
        .style(FontStyle::Bold)
        .color(&theme().text);
    let padding = 30;
    let line_height = 32;
    let box_width = (400.0 * 1.1) as i32;
    let box_height = 6 * line_height + padding * 2;
    let num_boxes = 3;
    let spacing = 40;
    let total_boxes_height = num_boxes * box_height + (num_boxes - 1) * spacing;
    let available_height = 1000;
    let box_x = 30;
    let box_y_top = (available_height - total_boxes_height) / 2;
    let box_y_bottom = box_y_top + box_height + spacing;
    let per_base_box_y = box_y_bottom + box_height + spacing;
    // --- Top box: Read stats ---
    if let Some(stats) = read_stats {
        let stats_labels = [
            "N50:",
            "Mean Qual:",
            "Median Qual:",
            "Mean Length:",
            "Median Len:",
        ];
        let stats_values = [
            format_number(stats.n50 as u64),
            format!("{:.2}", stats.mean_qual),
            format!("{:.2}", stats.median_qual),
            format_number(stats.mean_len as u64),
            format_number(stats.median_len as u64),
        ];

        // Draw box background and border
        left_panel.draw(&Rectangle::new(
            [
                (box_x, box_y_top),
                (box_x + box_width, box_y_top + box_height),
            ],
            ShapeStyle {
                color: theme().overlay.to_rgba(),
                filled: true,
                stroke_width: 0,
            },
        ))?;
        left_panel.draw(&Rectangle::new(
            [
                (box_x, box_y_top),
                (box_x + box_width, box_y_top + box_height),
            ],
            ShapeStyle {
                color: theme().accent.to_rgba(),
                filled: false,
                stroke_width: 3,
            },
        ))?;

        // Draw title and stats
        left_panel.draw_text("Read Stats", &font, (box_x + padding, box_y_top + padding))?;
        for (i, (label, value)) in stats_labels.iter().zip(stats_values.iter()).enumerate() {
            left_panel.draw_text(
                label,
                &font,
                (
                    box_x + padding,
                    box_y_top + padding + ((i as i32 + 1) * line_height),
                ),
            )?;
            left_panel.draw_text(
                value,
                &font,
                (
                    box_x + box_width - padding - 160,
                    box_y_top + padding + ((i as i32 + 1) * line_height),
                ),
            )?;
        }
    }

    // Calculate statistics for display in the stats boxes
    let coverage_stats = calculate_coverage_stats(&chrom_points);
    let mean = coverage_stats.mean;
    let median = coverage_stats.median;
    let min = coverage_stats.min;
    let max = coverage_stats.max;

    // --- Bottom box: Coverage stats ---
    let stats_labels = ["Mean:", "Median:", "Min:", "Max:", "Bin:"];
    let stats_values = [
        format!("{mean:.2}"),
        format!("{median:.2}"),
        format!("{min:.2}"),
        format!("{max:.2}"),
        bin_size_label.clone(),
    ];
    left_panel.draw(&Rectangle::new(
        [
            (box_x, box_y_bottom),
            (box_x + box_width, box_y_bottom + box_height),
        ],
        ShapeStyle {
            color: theme().overlay.to_rgba(),
            filled: true,
            stroke_width: 0,
        },
    ))?;
    left_panel.draw(&Rectangle::new(
        [
            (box_x, box_y_bottom),
            (box_x + box_width, box_y_bottom + box_height),
        ],
        ShapeStyle {
            color: theme().accent.to_rgba(),
            filled: false,
            stroke_width: 3,
        },
    ))?;
    left_panel.draw_text(
        "Coverage Stats",
        &font,
        (box_x + padding, box_y_bottom + padding),
    )?;
    for (i, (label, value)) in stats_labels.iter().zip(stats_values.iter()).enumerate() {
        left_panel.draw_text(
            label,
            &font,
            (
                box_x + padding,
                box_y_bottom + padding + ((i as i32 + 1) * line_height),
            ),
        )?;
        left_panel.draw_text(
            value,
            &font,
            (
                box_x + box_width - padding - 160,
                box_y_bottom + padding + ((i as i32 + 1) * line_height),
            ),
        )?;
    }

    // --- Per-base coverage stats box ---
    let per_base_labels = [
        "Per-base Mean:",
        "Per-base Median:",
        "Per-base Min:",
        "Per-base Max:",
        "Per-base Stddev:",
    ];
    // Calculate per-base statistics using the helper function
    let per_base_stats = calculate_per_base_stats(coverage);
    let per_base_values = [
        format!("{:.2}", per_base_stats.mean),
        format!("{:.2}", per_base_stats.median),
        format!("{:.2}", per_base_stats.min),
        format!("{:.2}", per_base_stats.max),
        format!("{:.2}", per_base_stats.stddev),
    ];

    left_panel.draw(&Rectangle::new(
        [
            (box_x, per_base_box_y),
            (box_x + box_width, per_base_box_y + box_height),
        ],
        ShapeStyle {
            color: theme().overlay.to_rgba(),
            filled: true,
            stroke_width: 0,
        },
    ))?;
    left_panel.draw(&Rectangle::new(
        [
            (box_x, per_base_box_y),
            (box_x + box_width, per_base_box_y + box_height),
        ],
        ShapeStyle {
            color: theme().accent.to_rgba(),
            filled: false,
            stroke_width: 3,
        },
    ))?;
    left_panel.draw_text(
        "Per-base Coverage",
        &font,
        (box_x + padding, per_base_box_y + padding),
    )?;
    for (i, (label, value)) in per_base_labels
        .iter()
        .zip(per_base_values.iter())
        .enumerate()
    {
        left_panel.draw_text(
            label,
            &font,
            (
                box_x + padding,
                per_base_box_y + padding + ((i as i32 + 1) * line_height),
            ),
        )?;
        left_panel.draw_text(
            value,
            &font,
            (
                box_x + box_width - padding - 160,
                per_base_box_y + padding + ((i as i32 + 1) * line_height),
            ),
        )?;
    }

    // This repeated computation was unnecessary and can be removed

    Ok(())
}

/// Plot genome overview with canonical chromosomes plus separate mito/EBV panels.
pub fn plot_overview_coverage(
    coverage: &HashMap<String, HashMap<u32, u32>>,
    output_path: &str,
    title_prefix: Option<&str>,
    read_stats: Option<&ReadStats>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let (canonical, mito, ebv) = partition_chromosomes(coverage);
    let has_any = !canonical.is_empty() || !mito.is_empty() || !ebv.is_empty();
    if !has_any {
        eprintln!("Warning: No coverage data available for overview plot");
        return Ok(());
    }

    let root = BitMapBackend::new(output_path, (2400, 1400)).into_drawing_area();
    root.fill(&theme().base)?;

    let (top_area, remainder) = root.split_vertically(880);
    let (stats_area, plot_area) = top_area.split_horizontally(420);
    let (gap_area, bottom_area) = remainder.split_vertically(40);
    gap_area.fill(&theme().base)?;
    let (mito_area, ebv_area) = bottom_area.split_horizontally(1200);

    let title_prefix = title_prefix
        .filter(|p| !p.is_empty())
        .map(|p| format!("{p} - "))
        .unwrap_or_default();

    draw_overview_panel(
        &plot_area,
        Some(&stats_area),
        &format!("{title_prefix}Canonical chromosomes"),
        &canonical,
        coverage,
        10_000,
        read_stats,
    )?;
    draw_overview_panel(
        &mito_area,
        None,
        &format!("{title_prefix}Mitochondrial (MT)"),
        &mito,
        coverage,
        1_000,
        read_stats,
    )?;
    draw_overview_panel(
        &ebv_area,
        None,
        &format!("{title_prefix}EBV"),
        &ebv,
        coverage,
        1_000,
        read_stats,
    )?;

    Ok(())
}

fn draw_overview_panel(
    area: &DrawingArea<BitMapBackend, Shift>,
    stats_area: Option<&DrawingArea<BitMapBackend, Shift>>,
    title: &str,
    chroms: &[String],
    coverage: &HashMap<String, HashMap<u32, u32>>,
    max_bin_size: u32,
    read_stats: Option<&ReadStats>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    if chroms.is_empty() {
        if let Some(stats_area) = stats_area {
            draw_empty_panel(stats_area, "Overview Stats")?;
        }
        return draw_empty_panel(area, title);
    }

    let panel_padding = 30;
    let panel = area.margin(panel_padding, panel_padding, panel_padding, panel_padding);

    let (points, total_len, bin_size, spans) =
        build_overview_points(chroms, coverage, max_bin_size);
    if total_len == 0 || points.is_empty() {
        if let Some(stats_area) = stats_area {
            draw_empty_panel(stats_area, "Overview Stats")?;
        }
        return draw_empty_panel(area, title);
    }

    let coverage_stats = calculate_coverage_stats(&points);
    let bin_size_label = if bin_size >= 1_000_000 {
        format!("{} Mb", bin_size / 1_000_000)
    } else if bin_size >= 1_000 {
        format!("{} kb", bin_size / 1_000)
    } else {
        format!("{bin_size} bp")
    };
    if let Some(stats_area) = stats_area {
        draw_overview_stats_panel(stats_area, read_stats, &coverage_stats, &bin_size_label)?;
    }

    let y_max = points.iter().map(|&(_, y)| y).fold(0.0, f64::max).max(1.0);
    let y_min = 0.1;
    let y_top = y_max * 1.1;

    let mut chart = ChartBuilder::on(&panel)
        .set_label_area_size(LabelAreaPosition::Left, 60)
        .set_label_area_size(LabelAreaPosition::Bottom, 40)
        .margin(10)
        .caption(title, ("sans-serif", 28).into_font().color(&theme().text))
        .build_cartesian_2d(0i64..total_len as i64, (y_min..y_top).log_scale())?;

    chart
        .configure_mesh()
        .x_desc("")
        .y_desc("Depth (log scale)")
        .axis_desc_style(("sans-serif", 18).into_font().color(&theme().text))
        .x_labels(0)
        .x_label_style(("sans-serif", 14).into_font().color(&theme().text))
        .y_label_style(("sans-serif", 14).into_font().color(&theme().text))
        .light_line_style(RGBAColor(80, 80, 80, 0.55))
        .bold_line_style(RGBAColor(80, 80, 80, 0.8))
        .axis_style(theme().text)
        .draw()?;

    let overlay = theme().overlay;
    let base = theme().base;
    let alt_rgb = blend_colors(&overlay, &base, 0.5);
    let bg_colors = [
        RGBAColor(overlay.0, overlay.1, overlay.2, 0.9),
        RGBAColor(alt_rgb.0, alt_rgb.1, alt_rgb.2, 0.9),
    ];
    chart.draw_series(spans.iter().enumerate().map(|(idx, span)| {
        let color = bg_colors[idx % bg_colors.len()];
        Rectangle::new(
            [(span.start as i64, y_min), (span.end as i64, y_top)],
            color.filled(),
        )
    }))?;

    let bin_width = bin_size as i64;
    chart.draw_series(points.iter().map(|&(x, y)| {
        let end = (x + bin_width).min(total_len as i64);
        Rectangle::new([(x, y_min), (end, y)], theme().primary.filled())
    }))?;

    let rotate = spans.len() > 18;
    let label_font = if rotate {
        ("sans-serif", 12)
            .into_font()
            .transform(FontTransform::Rotate90)
    } else {
        ("sans-serif", 14).into_font()
    }
    .color(&theme().text);
    let label_y = y_min * 1.0;
    let label_offset_px = 12;
    let plotting_area = chart.plotting_area();
    let text_area = plotting_area.strip_coord_spec();
    let (base_x, base_y) = text_area.get_base_pixel();
    let label_value = label_y.max(y_min);
    for span in spans {
        let center = span.start.saturating_add((span.end - span.start) / 2);
        let (px, py) = chart.backend_coord(&(center as i64, label_value));
        let x = px - base_x;
        let y = py - base_y + label_offset_px;
        text_area.draw_text(span.name.as_str(), &label_font, (x, y))?;
    }

    let axis_title = "Chromosome";
    let axis_font = ("sans-serif", 18).into_font().color(&theme().text);
    let axis_area = area.strip_coord_spec();
    let (w, h) = axis_area.dim_in_pixel();
    let (tw, th) = axis_area.estimate_text_size(axis_title, &axis_font)?;
    let axis_offset_px = 10i32;
    let x = ((w as i32 - tw as i32) / 2).max(0);
    let y = (h as i32 - th as i32 + axis_offset_px).max(0);
    axis_area.draw_text(axis_title, &axis_font, (x, y))?;

    Ok(())
}

fn stats_box_height(rows: usize, line_height: i32, padding: i32) -> i32 {
    (rows as i32 + 1) * line_height + padding * 2
}

fn draw_overview_stats_panel(
    area: &DrawingArea<BitMapBackend, Shift>,
    read_stats: Option<&ReadStats>,
    coverage_stats: &CoverageStats,
    bin_size_label: &str,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    area.fill(&theme().base)?;

    let (panel_w, panel_h) = area.dim_in_pixel();
    let padding = 24;
    let line_height = 28;
    let spacing = 32;
    let box_width = (panel_w as i32 - padding * 2).max(240);
    let box_x = padding;

    let font = ("sans-serif", 22)
        .into_font()
        .style(FontStyle::Bold)
        .color(&theme().text);

    let mut boxes: Vec<(&str, Vec<&str>, Vec<String>)> = Vec::new();
    if let Some(stats) = read_stats {
        let labels = vec![
            "N50:",
            "Mean Qual:",
            "Median Qual:",
            "Mean Length:",
            "Median Len:",
        ];
        let values = vec![
            format_number(stats.n50 as u64),
            format!("{:.2}", stats.mean_qual),
            format!("{:.2}", stats.median_qual),
            format_number(stats.mean_len as u64),
            format_number(stats.median_len as u64),
        ];
        boxes.push(("Read Stats", labels, values));
    }

    let coverage_labels = vec!["Mean:", "Median:", "Min:", "Max:", "Bin:"];
    let coverage_values = vec![
        format!("{:.2}", coverage_stats.mean),
        format!("{:.2}", coverage_stats.median),
        format!("{:.2}", coverage_stats.min),
        format!("{:.2}", coverage_stats.max),
        bin_size_label.to_string(),
    ];
    boxes.push(("Coverage Stats", coverage_labels, coverage_values));

    let heights: Vec<i32> = boxes
        .iter()
        .map(|(_, labels, _)| stats_box_height(labels.len(), line_height, padding))
        .collect();
    let total_height =
        heights.iter().sum::<i32>() + spacing * (boxes.len().saturating_sub(1) as i32);
    let start_y = ((panel_h as i32 - total_height) / 2).max(padding);

    let mut y = start_y;
    for (idx, (title, labels, values)) in boxes.iter().enumerate() {
        let height = heights[idx];
        area.draw(&Rectangle::new(
            [(box_x, y), (box_x + box_width, y + height)],
            ShapeStyle {
                color: theme().overlay.to_rgba(),
                filled: true,
                stroke_width: 0,
            },
        ))?;
        area.draw(&Rectangle::new(
            [(box_x, y), (box_x + box_width, y + height)],
            ShapeStyle {
                color: theme().accent.to_rgba(),
                filled: false,
                stroke_width: 3,
            },
        ))?;
        area.draw_text(title, &font, (box_x + padding, y + padding))?;

        let mut value_x = box_x + ((box_width as f32) * 0.58) as i32;
        let value_max = box_x + box_width - padding - 10;
        if value_x > value_max {
            value_x = value_max;
        }

        for (i, (label, value)) in labels.iter().zip(values.iter()).enumerate() {
            area.draw_text(
                label,
                &font,
                (
                    box_x + padding,
                    y + padding + ((i as i32 + 1) * line_height),
                ),
            )?;
            area.draw_text(
                value,
                &font,
                (value_x, y + padding + ((i as i32 + 1) * line_height)),
            )?;
        }
        y += height + spacing;
    }

    Ok(())
}

fn draw_empty_panel(
    area: &DrawingArea<BitMapBackend, Shift>,
    title: &str,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    area.fill(&theme().base)?;
    let font = ("sans-serif", 24).into_font().color(&theme().text);
    let note_font = ("sans-serif", 16).into_font().color(&theme().text);
    let (_w, h) = area.dim_in_pixel();
    area.draw_text(title, &font, (30, 30))?;
    area.draw_text("No data", &note_font, (30, (h / 2) as i32))?;
    Ok(())
}

struct ChromSpan {
    name: String,
    start: u32,
    end: u32,
}

fn build_overview_points(
    chroms: &[String],
    coverage: &HashMap<String, HashMap<u32, u32>>,
    max_bin_size: u32,
) -> (Vec<(i64, f64)>, u32, u32, Vec<ChromSpan>) {
    let mut spans: Vec<ChromSpan> = Vec::new();
    let mut total_len = 0u32;

    for chrom in chroms {
        if let Some(chrom_coverage) = coverage.get(chrom)
            && let Some(&max_pos) = chrom_coverage.keys().max()
        {
            let chrom_len = max_pos.max(1);
            spans.push(ChromSpan {
                name: chrom.clone(),
                start: total_len,
                end: total_len.saturating_add(chrom_len),
            });
            total_len = total_len.saturating_add(chrom_len);
        }
    }

    if total_len == 0 {
        return (Vec::new(), 0, 1, Vec::new());
    }

    let target_bins = 5000u32;
    let mut bin_size = (total_len / target_bins).max(1);
    if bin_size > max_bin_size {
        bin_size = max_bin_size;
    }

    let mut bin_sums: HashMap<u32, u64> = HashMap::new();
    let mut bin_lengths: HashMap<u32, u32> = HashMap::new();

    for span in &spans {
        let chrom_start = span.start;
        let chrom_end = span.end;
        if chrom_start >= chrom_end {
            continue;
        }

        let mut bin_start = (chrom_start / bin_size) * bin_size;
        while bin_start < chrom_end {
            let bin_end = bin_start.saturating_add(bin_size).min(chrom_end);
            let overlap_start = chrom_start.max(bin_start);
            let overlap_end = chrom_end.min(bin_end);
            let len = overlap_end.saturating_sub(overlap_start);
            if len > 0 {
                let bin_idx = bin_start / bin_size;
                *bin_lengths.entry(bin_idx).or_insert(0) += len;
            }
            bin_start = bin_start.saturating_add(bin_size);
        }

        if let Some(chrom_coverage) = coverage.get(&span.name) {
            for (&pos, &count) in chrom_coverage.iter() {
                if pos == 0 {
                    continue;
                }
                let global_pos = span.start.saturating_add(pos - 1);
                let bin_idx = global_pos / bin_size;
                let entry = bin_sums.entry(bin_idx).or_insert(0);
                *entry += count as u64;
            }
        }
    }

    let mut points = Vec::new();
    for (bin_idx, sum) in bin_sums {
        if let Some(len) = bin_lengths.get(&bin_idx) {
            if *len == 0 {
                continue;
            }
            let avg = sum as f64 / *len as f64;
            if avg > 0.0 {
                points.push((bin_idx as i64 * bin_size as i64, avg));
            }
        }
    }
    points.sort_by_key(|&(x, _)| x);

    (points, total_len, bin_size, spans)
}

fn partition_chromosomes(
    coverage: &HashMap<String, HashMap<u32, u32>>,
) -> (Vec<String>, Vec<String>, Vec<String>) {
    let mut canonical: Vec<String> = Vec::new();
    let mut mito: Vec<String> = Vec::new();
    let mut ebv: Vec<String> = Vec::new();

    for chrom in coverage.keys() {
        let name = chrom.to_string();
        if is_mito(name.as_str()) {
            mito.push(name);
        } else if is_ebv(name.as_str()) {
            ebv.push(name);
        } else if is_canonical(name.as_str()) {
            canonical.push(name);
        }
    }

    canonical.sort_by_key(|c| canonical_order_key(c).unwrap_or(1000));
    mito.sort();
    ebv.sort();

    (canonical, mito, ebv)
}

fn is_canonical(name: &str) -> bool {
    canonical_order_key(name).is_some()
}

fn canonical_order_key(name: &str) -> Option<u32> {
    let mut s = name.to_ascii_lowercase();
    if let Some(stripped) = s.strip_prefix("chr") {
        s = stripped.to_string();
    }
    if s == "x" {
        return Some(23);
    }
    if s == "y" {
        return Some(24);
    }
    if let Ok(val) = s.parse::<u32>()
        && (1..=22).contains(&val)
    {
        return Some(val);
    }
    None
}

fn is_mito(name: &str) -> bool {
    let n = name.to_ascii_lowercase();
    n == "chrm" || n == "chrmt" || n == "mt" || n == "m"
}

fn is_ebv(name: &str) -> bool {
    let n = name.to_ascii_lowercase();
    n == "chrebv" || n == "ebv"
}

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

    #[test]
    fn test_plot_output() {
        let mut coverage = HashMap::new();
        for i in 0..50u32 {
            coverage.insert(i, (i % 5) + 1);
        }

        let out_path = "test-out/coverage.test.png";
        let _ = fs::remove_file(out_path);
        plot_per_base_coverage(
            "chrTest", &coverage, out_path, None, true, false, None, None,
        )
        .expect("PNG plotting should succeed");

        assert!(fs::metadata(out_path).is_ok(), "Output PNG should exist");
        let meta = fs::metadata(out_path).unwrap();
        assert!(meta.len() > 0, "Output PNG should not be empty");
    }

    #[test]
    fn test_plot_per_base_coverage_with_different_themes() {
        let mut coverage = HashMap::new();
        // Simulate a small region with variable coverage
        for i in 0..100u32 {
            coverage.insert(i, (i % 10) + 1);
        }

        // Test with Nord theme
        set_theme("nord");
        let out_path = "test-out/coverage.test.nord.png";
        let _ = fs::remove_file(out_path);
        plot_per_base_coverage(
            "chrTest", &coverage, out_path, None, false, false, None, None,
        )
        .expect("plotting with Nord theme should succeed");

        // Test with Frappe theme
        set_theme("frappe");
        let out_path = "test-out/coverage.test.frappe.png";
        let _ = fs::remove_file(out_path);
        plot_per_base_coverage(
            "chrTest", &coverage, out_path, None, false, false, None, None,
        )
        .expect("plotting with Frappe theme should succeed");

        // Test with Gruvbox theme
        set_theme("gruvbox");
        let out_path = "test-out/coverage.test.gruvbox.png";
        let _ = fs::remove_file(out_path);
        plot_per_base_coverage(
            "chrTest", &coverage, out_path, None, false, false, None, None,
        )
        .expect("plotting with Gruvbox theme should succeed");

        // Test with Latte theme
        set_theme("latte");
        let out_path = "test-out/coverage.test.latte.png";
        let _ = fs::remove_file(out_path);
        plot_per_base_coverage(
            "chrTest", &coverage, out_path, None, false, false, None, None,
        )
        .expect("plotting with Latte theme should succeed");

        // Reset to default theme
        set_theme("latte");
    }
}