visualisation.rs

//! Spectrogram & ESD plotting helpers using plotters + colorous (v‑5)
//!
//! *Shared style constants* → spectrogram and ESD now default to the **same
//! base canvas size** (1800×900 at 100dpi) so exported PNGs line up perfectly.
//! Adjust one pair of numbers and both functions follow.
//!
//! # Crate features
//! No special features required (we avoid `full_palette`).
//!
//! ---------------------------------------------------------------------------

use plotters::prelude::*;
use plotters::style::text_anchor::{HPos, VPos, Pos};
use plotters::style::FontTransform;
use colorous::MAGMA;
use std::fs;
use std::path::Path;

use crate::analysis::spectrogram::Spectrogram;

// ==================== Style constants =========================================
const BASE_W:      u32 = 1800; // logical canvas size at 100 dpi
const BASE_H:      u32 =  900;
const LABEL_FONT_PT: u32 = 28;
const TICK_FONT_PT:  u32 = 22;
const LEGEND_GUTTER: u32 = 30;
const GRADIENT_W:    u32 = 36;
const LABEL_PAD:     i32 = 60;
const ORANGE_RED: RGBColor = RGBColor(255, 69, 0);

/// Helper – convert a `colorous::Color` into a `plotters::RGBColor`
#[inline]
fn to_rgb(c: colorous::Color) -> RGBColor { RGBColor(c.r, c.g, c.b) }

/// Helper - Converts a plot title string to a clean filename.
fn title_to_filename(title: &str) -> String {
    title
        .to_lowercase()
        .replace(&[' ', ':', '/'][..], "_") // replace spaces, colons, slashes with _
        + ".png"
}

// ╔══════════════════════════════════════════════════════════════════════╗
// ║ 1) SPECTROGRAM WITH MAGMA COLOUR-BAR                                 ║
// ╚══════════════════════════════════════════════════════════════════════╝

pub fn plot_spectrogram(
    spec: &Spectrogram,
    output_path: &str,
    title: &str,
    dpi: u32,
) -> Result<(), Box<dyn std::error::Error>> {
    // ---------- canvas size ------------------------------------------
    let scale = dpi as f32 / 100.0;
    let (width, height) = ((BASE_W as f32 * scale) as u32, (BASE_H as f32 * scale) as u32);

    // make sure output dir exists
    if let Some(parent) = Path::new(output_path).parent() { fs::create_dir_all(parent)?; }
    let root = BitMapBackend::new(output_path,(width,height)).into_drawing_area();
    root.fill(&WHITE)?;

    // ---------- split main + legend, then add gutter -----------------
    let legend_panel_w = (width as f32 * 0.10) as u32;
    let (main_area, 
        raw_legend) = root.split_horizontally(width - legend_panel_w - LEGEND_GUTTER);
    let (_gutter, 
        legend_panel) = raw_legend.split_horizontally(LEGEND_GUTTER); // drop the gutter area

    // further split legend into gradient + ticks/label
    let (gradient_area, 
        tick_area) = legend_panel.split_horizontally(GRADIENT_W);

    // ---------- axis ranges -----------------------------------------
    let (freq_min,freq_max) = (*spec.freqs.first().unwrap_or(&0.0),*spec.freqs.last().unwrap_or(&3000.0));
    let (time_min,time_max) = (*spec.times.first().unwrap_or(&0.0),*spec.times.last().unwrap_or(&4.0));
    
    let db_min = -140.0; 
    let db_max = *spec.data.iter().max_by(|a,b| a.total_cmp(b)).unwrap_or(&0.0);

    // ---------- build chart -----------------------------------------
    let mut chart = ChartBuilder::on(&main_area)
        .caption(title,("sans-serif", 42))
        .margin(20)
        .x_label_area_size(60)
        .y_label_area_size(110)
        .build_cartesian_2d(time_min..time_max,freq_min..freq_max)?;
    
    chart.configure_mesh()
        .x_desc("Time (s)") .y_desc("Frequency (Hz)")
        .axis_desc_style(("sans-serif", LABEL_FONT_PT).into_font().style(FontStyle::Bold))
        .label_style(("sans-serif", TICK_FONT_PT))
        .light_line_style(&WHITE.mix(0.2))
        .draw()?;

    // heat‑map tiles
    chart.draw_series(spec.data.indexed_iter().map(|((fi,ti),&db)| {
        let x0 = spec.times[ti]; let x1 = *spec.times.get(ti+1).unwrap_or(&time_max);
        let y0 = spec.freqs[fi]; let y1 = *spec.freqs.get(fi+1).unwrap_or(&freq_max);
        let t = ((db.clamp(db_min,db_max)-db_min)/(db_max-db_min)) as f64;
        Rectangle::new([(x0,y0),(x1,y1)],to_rgb(MAGMA.eval_continuous(t)).filled())
    }))?;

    // ---------- colour‑bar ------------------------------------------
    let (_,grad_y) = gradient_area.get_pixel_range();
    let (_, plot_y) = chart.plotting_area().get_pixel_range();
    let top_offset = (plot_y.start - grad_y.start) as i32;
    let grad_h = (plot_y.end - plot_y.start) as i32;

    let n_steps = 256;
    for i in 0..n_steps {
        let interp = i as f64 / (n_steps-1) as f64;
        let color = to_rgb(MAGMA.eval_continuous(1.0-interp));
        let y0 = top_offset + (i * grad_h / n_steps as i32);
        let y1 = top_offset + ((i+1)*grad_h / n_steps as i32);
        gradient_area.draw(&Rectangle::new([(0,y0),(GRADIENT_W as i32,y1)],color.filled()))?;
    }

    // ticks & numbers --------------------------------------------------
    let tick_style = TextStyle::from(("sans-serif",TICK_FONT_PT).into_font())
        .pos(Pos::new(HPos::Left,VPos::Center));
    for db in (db_min as i32..=db_max as i32).step_by(20) {
        let interp = (db_max-db as f32) as f64 / (db_max-db_min) as f64;
        let y = top_offset + (interp*grad_h as f64) as i32;
        gradient_area.draw(&PathElement::new(vec![(GRADIENT_W as i32 + 2,y),(GRADIENT_W as i32 + 6,y)], &BLACK))?;
        tick_area.draw(&Text::new(format!("{db:>3}"), (8,y), tick_style.clone()))?;
    }

    // vertical label (rotated 270° so text reads top→bottom) ----------
    let (_,ty) = tick_area.get_pixel_range();
    let mid_y = (ty.start+ty.end) as i32 / 2;
    let axis_style = TextStyle::from(("sans-serif",LABEL_FONT_PT).into_font().style(FontStyle::Bold))
        .transform(FontTransform::Rotate270)
        .pos(Pos::new(HPos::Center,VPos::Center));
    tick_area.draw(&Text::new("Power (dB)",(LABEL_PAD, mid_y),axis_style))?;

    Ok(())
}

// ╔══════════════════════════════════════════════════════════════════════╗
// ║ 2) ENERGY SPECTRAL DENSITY LINE PLOT                                 ║
// ╚══════════════════════════════════════════════════════════════════════╝

/// Plot the energy spectral density with styling consistent with the spectrogram.
pub fn plot_energy_spectral_density(
    freqs: &[f32],
    power_spectrum: &[f32],
    title: &str,
    x_lim: Option<(f32, f32)>,
    y_lim: Option<(f32, f32)>,
    dpi: u32,
) -> Result<(), Box<dyn std::error::Error>> {
    //---------------- output path --------------------------------------------
    let out_dir = Path::new("output/plots");
    fs::create_dir_all(out_dir)?;
    let filename = out_dir.join(title_to_filename(title));

    //---------------- canvas --------------------------------------------------
    let scale = if dpi == 0 { 1.0 } else { dpi as f32 / 100.0 };
    let (w, h) = ((BASE_W as f32 * scale) as u32, (BASE_H as f32 * scale) as u32);
    let root = BitMapBackend::new(filename.to_str().unwrap(), (w, h)).into_drawing_area();
    root.fill(&WHITE)?;

    //---------------- axis limits --------------------------------------------
    let (x_min, x_max) = x_lim.unwrap_or_else(|| {
        let min = freqs.iter().cloned().fold(f32::INFINITY, f32::min);
        let max = freqs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        (min, max)
    });
    let (y_min, y_max) = y_lim.unwrap_or_else(|| {
        (0.0, power_spectrum.iter().cloned().fold(f32::NEG_INFINITY, f32::max))
    });

    //---------------- chart ---------------------------------------------------
    let mut chart = ChartBuilder::on(&root)
        .caption(title, ("sans-serif", 30))
        .margin(20)
        .x_label_area_size(60)
        .y_label_area_size(80)
        .build_cartesian_2d(x_min..x_max, y_min..y_max)?;

    chart.configure_mesh()
        .x_desc("Frequency (Hz)")
        .y_desc("Energy Spectral Density")
        .axis_desc_style(("sans-serif", LABEL_FONT_PT - 2).into_font().style(FontStyle::Bold))
        .label_style(("sans-serif", TICK_FONT_PT))
        .y_label_formatter(&|v| format!("{:.1e}", v))
        .light_line_style(&WHITE.mix(0.2))
        .draw()?;

    // Line series (propagate error with `?`)
    chart.draw_series(LineSeries::new(
        freqs.iter().copied().zip(power_spectrum.iter().copied()),
        ORANGE_RED.stroke_width(2),
    ))?;

    Ok(())
}