use super::Config;
use num_traits::{cast::AsPrimitive, Float};
use plotters::prelude::*;
type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
type Data<'a, T> = (&'a [T], (usize, usize));
pub struct Heatmap {}
impl<'a, T: Float + AsPrimitive<f64>> From<(Data<'a, T>, Option<Config>)> for Heatmap {
fn from((data, config): (Data<T>, Option<Config>)) -> Self {
fn inner<T>((data, config): (Data<T>, Option<Config>)) -> Result<()>
where
T: Float + AsPrimitive<f64>,
{
let (map, (rows, cols)) = data;
assert_eq!(rows, cols, "rectangular heatmap unimplemented");
let config = config.unwrap_or_default();
let osf = config.osf;
let res = rows;
let size = res * osf;
let filename = config
.filename
.unwrap_or_else(|| "complot-heatmap.png".to_string());
let cmap = config.cmap;
let width = size as u32 + 50;
let height = size as u32 + 90;
let root = BitMapBackend::new(&filename, (width, height)).into_drawing_area();
let (plot, colorbar) = root.split_vertically(size as u32 + 30);
plot.fill(&BLACK)?;
let mut chart = ChartBuilder::on(&plot);
chart
.margin_left(20)
.margin_right(20)
.margin_top(0)
.margin_bottom(0);
if let Some(value) = config.title {
chart.caption(value, ("sans-serif", 16, &WHITE));
}
let mut chart_ctx = chart
.build_cartesian_2d(0i32..(size - 1) as i32, 0i32..(size - 1) as i32)
.expect("Failed building chart");
let (cells_min, cells_max) = match config.cmap_minmax {
Some(value) => value,
None => (
map.iter()
.cloned()
.fold(Float::infinity(), Float::min)
.as_(),
map.iter()
.cloned()
.fold(Float::neg_infinity(), Float::max)
.as_(),
),
};
chart_ctx
.configure_mesh()
.disable_x_mesh()
.disable_y_mesh()
.draw()?;
chart_ctx.draw_series(map.iter().enumerate().map(|(k, &v)| {
let j = (k / res) as i32;
let i = (k % res) as i32;
let u = (v.as_() - cells_min) / (cells_max - cells_min);
let c = cmap.eval_continuous(u as f64).as_tuple();
Rectangle::new(
[
(osf as i32 * i, osf as i32 * j),
(osf as i32 * (i + 1), osf as i32 * (j + 1)),
],
RGBColor(c.0, c.1, c.2).filled(),
)
}))?;
colorbar.fill(&BLACK)?;
let mut colorbar_chart = ChartBuilder::on(&colorbar)
.margin_left(20)
.margin_right(20)
.set_label_area_size(LabelAreaPosition::Bottom, 40)
.build_cartesian_2d(cells_min..cells_max, 0f64..1f64)?;
let mut mesh = colorbar_chart.configure_mesh();
mesh.axis_style(WHITE)
.set_tick_mark_size(LabelAreaPosition::Bottom, 5)
.x_label_style(("sans-serif", 14, &WHITE));
if let Some(value) = config.xaxis.label {
mesh.x_desc(value);
}
mesh.draw()?;
let dx = (cells_max - cells_min) / (size - 1) as f64;
colorbar_chart.draw_series((0..size).map(|k| {
let x = cells_min + k as f64 * dx;
let c = cmap.eval_rational(k, size).as_tuple();
Rectangle::new([(x, 0.), (x + dx, 1.)], RGBColor(c.0, c.1, c.2).filled())
}))?;
Ok(())
}
if let Err(e) = inner((data, config)) {
eprintln!("Complot failed in Heatmap: {}", e);
}
Heatmap {}
}
}