use crate::color_scheme::{default_color_scheme_data, get_color_count};
pub struct HeatmapStamp {
w: u32,
h: u32,
buffer: Vec<f32>,
}
impl HeatmapStamp {
pub fn new(width: u32, height: u32, data: Vec<f32>) -> Self {
assert_eq!(data.len(), (width * height) as usize);
Self {
w: width,
h: height,
buffer: data,
}
}
pub fn new_round(radius: u32) -> Self {
let d = 2 * radius + 1;
let w = d;
let h = d;
let mut buffer = vec![0.0; (w * h) as usize];
for y in 0..h {
for x in 0..w {
let dx = x as f32 - radius as f32;
let dy = y as f32 - radius as f32;
let dist = (dx * dx + dy * dy).sqrt() / (radius + 1) as f32;
let clamped_ds = dist.clamp(0.0, 1.0);
buffer[(y * w + x) as usize] = 1.0 - clamped_ds;
}
}
Self { w, h, buffer }
}
pub fn new_with_shape<F>(radius: u32, distshape: F) -> Self
where
F: Fn(f32) -> f32,
{
let d = 2 * radius + 1;
let w = d;
let h = d;
let mut buffer = vec![0.0; (w * h) as usize];
for y in 0..h {
for x in 0..w {
let dx = x as f32 - radius as f32;
let dy = y as f32 - radius as f32;
let dist = (dx * dx + dy * dy).sqrt() / (radius + 1) as f32;
let ds = distshape(dist);
let clamped_ds = ds.clamp(0.0, 1.0);
buffer[(y * w + x) as usize] = 1.0 - clamped_ds;
}
}
Self { w, h, buffer }
}
pub fn get_width(&self) -> u32 {
self.w
}
pub fn get_height(&self) -> u32 {
self.h
}
pub fn get_buffer(&self) -> &[f32] {
&self.buffer
}
}
fn default_stamp_data() -> Vec<f32> {
vec![
0.0, 0.0, 0.1055728, 0.1753789, 0.2, 0.1753789, 0.1055728, 0.0, 0.0, 0.0, 0.1514719,
0.2788897, 0.3675445, 0.4, 0.3675445, 0.2788897, 0.1514719, 0.0, 0.1055728, 0.2788897,
0.4343146, 0.5527864, 0.6, 0.5527864, 0.4343146, 0.2788897, 0.1055728, 0.1753789,
0.3675445, 0.5527864, 0.7171573, 0.8, 0.7171573, 0.5527864, 0.3675445, 0.1753789, 0.2, 0.4,
0.6, 0.8, 1.0, 0.8, 0.6, 0.4, 0.2, 0.1753789, 0.3675445, 0.5527864, 0.7171573, 0.8,
0.7171573, 0.5527864, 0.3675445, 0.1753789, 0.1055728, 0.2788897, 0.4343146, 0.5527864,
0.6, 0.5527864, 0.4343146, 0.2788897, 0.1055728, 0.0, 0.1514719, 0.2788897, 0.3675445, 0.4,
0.3675445, 0.2788897, 0.1514719, 0.0, 0.0, 0.0, 0.1055728, 0.1753789, 0.2, 0.1753789,
0.1055728, 0.0, 0.0,
]
}
fn default_heatmap_stamp() -> HeatmapStamp {
HeatmapStamp::new(9, 9, default_stamp_data())
}
pub struct Heatmap {
pub width: u32,
pub height: u32,
pub max_heat: f32,
pub buffer: Vec<f32>,
}
impl Heatmap {
pub fn new(width: u32, height: u32) -> Self {
Self {
width,
height,
max_heat: 0.0,
buffer: vec![0.0; (width * height) as usize],
}
}
pub fn add_point(&mut self, x: u32, y: u32) {
let stamp = default_heatmap_stamp();
self.add_point_with_stamp(x, y, &stamp);
}
pub fn add_point_with_stamp(&mut self, x: u32, y: u32, stamp: &HeatmapStamp) {
if x >= self.width || y >= self.height {
return;
}
let stamp_w = stamp.get_width();
let stamp_h = stamp.get_height();
let stamp_buf = stamp.get_buffer();
let x0 = if x < stamp_w / 2 { stamp_w / 2 - x } else { 0 };
let y0 = if y < stamp_h / 2 { stamp_h / 2 - y } else { 0 };
let x1 = if (x + stamp_w / 2) < self.width {
stamp_w
} else {
stamp_w / 2 + (self.width - x)
};
let y1 = if (y + stamp_h / 2) < self.height {
stamp_h
} else {
stamp_h / 2 + (self.height - y)
};
for iy in y0..y1 {
let buf_y = (y + iy) - stamp_h / 2;
let buf_line_start = (buf_y * self.width + (x + x0) - stamp_w / 2) as usize;
let stamp_line_start = (iy * stamp_w + x0) as usize;
for ix in 0..(x1 - x0) {
let buf_idx = buf_line_start + ix as usize;
let stamp_idx = stamp_line_start + ix as usize;
let stamp_value = stamp_buf[stamp_idx];
self.buffer[buf_idx] += stamp_value;
if self.buffer[buf_idx] > self.max_heat {
self.max_heat = self.buffer[buf_idx];
}
}
}
}
pub fn add_weighted_point(&mut self, x: u32, y: u32, weight: f32) {
let stamp = default_heatmap_stamp();
self.add_weighted_point_with_stamp(x, y, weight, &stamp);
}
pub fn add_weighted_point_with_stamp(
&mut self,
x: u32,
y: u32,
weight: f32,
stamp: &HeatmapStamp,
) {
if x >= self.width || y >= self.height || weight < 0.0 {
return;
}
let stamp_w = stamp.get_width();
let stamp_h = stamp.get_height();
let stamp_buf = stamp.get_buffer();
let x0 = if x < stamp_w / 2 { stamp_w / 2 - x } else { 0 };
let y0 = if y < stamp_h / 2 { stamp_h / 2 - y } else { 0 };
let x1 = if (x + stamp_w / 2) < self.width {
stamp_w
} else {
stamp_w / 2 + (self.width - x)
};
let y1 = if (y + stamp_h / 2) < self.height {
stamp_h
} else {
stamp_h / 2 + (self.height - y)
};
for iy in y0..y1 {
let buf_y = (y + iy) - stamp_h / 2;
let buf_line_start = (buf_y * self.width + (x + x0) - stamp_w / 2) as usize;
let stamp_line_start = (iy * stamp_w + x0) as usize;
for ix in 0..(x1 - x0) {
let buf_idx = buf_line_start + ix as usize;
let stamp_idx = stamp_line_start + ix as usize;
let stamp_value = stamp_buf[stamp_idx] * weight;
self.buffer[buf_idx] += stamp_value;
if self.buffer[buf_idx] > self.max_heat {
self.max_heat = self.buffer[buf_idx];
}
}
}
}
pub fn render(&self) -> Vec<u8> {
let colors = default_color_scheme_data();
self.render_with_colors(&colors)
}
pub fn render_with_colors(&self, colors: &[u8]) -> Vec<u8> {
let saturation = if self.max_heat > 0.0 {
self.max_heat
} else {
1.0
};
self.render_saturated(colors, saturation)
}
pub fn render_saturated(&self, colors: &[u8], saturation: f32) -> Vec<u8> {
assert!(saturation > 0.0);
let total_pixels = (self.width * self.height) as usize;
let mut colorbuf = vec![0u8; total_pixels * 4];
let ncolors = get_color_count(colors);
if ncolors == 0 {
return colorbuf;
}
for idx in 0..total_pixels {
let val = self.buffer[idx];
let normalized = (val / saturation).clamp(0.0, 1.0);
let color_idx = ((ncolors - 1) as f32 * normalized + 0.5) as usize;
let color_idx = color_idx.min(ncolors - 1);
let src_start = color_idx * 4;
let dst_start = idx * 4;
colorbuf[dst_start..dst_start + 4].copy_from_slice(&colors[src_start..src_start + 4]);
}
colorbuf
}
pub fn reset(&mut self) {
self.buffer.fill(0.0);
self.max_heat = 0.0;
}
}