timechart/

heatmap.rs

use std::io::Write;

use chrono::{Datelike, NaiveDate};

use crate::{Data, Result};

/// Renders an SVG heatmap.
pub struct Heatmap {
    pub projection: Box<dyn Projection>,
    pub colors: Vec<String>,
    pub unit: Option<String>,
    pub log_scale: bool,
}

impl Heatmap {
    pub fn render<T: Write>(&self, data: &Data, mut io: T) -> Result<()> {
        // No data, no image.
        if data.is_empty() {
            writeln!(
                io,
                r#"<svg xmlns="http://www.w3.org/2000/svg" width="0" height="0"></svg>"#
            )?;
            return Ok(());
        }

        let (y0, y1) = data.values().copied().fold((f64::MAX, f64::MIN), |acc, v| {
            let min = if acc.0 < v { acc.0 } else { v };
            let max = if acc.1 > v { acc.1 } else { v };
            (min, max)
        });

        // The color scale range.
        let y_range = if self.log_scale {
            y1.log10() - y0.log10()
        } else {
            y1 - y0
        };

        let (width, height) = self.projection.dimensions();
        write!(
            io,
            r#"<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}">"#
        )?;

        let (start, end) = self.projection.range();
        for date in start.iter_days() {
            if date > end {
                break;
            }

            let (x, y) = self.projection.map(date);
            let value = data.get(&date).copied().unwrap_or(0.0);

            // Scale
            let scaled_value = if self.log_scale {
                (value.signum() * (1.0 + value.abs()).log10())
                    / (y1.signum() * (1.0 + y1.abs()).log10())
            } else {
                (value + y0) / y_range
            };

            let color = &self.colors[(scaled_value * (self.colors.len() - 1) as f64) as usize];

            writeln!(
                io,
                r#"<rect x="{x}" y="{y}" width="{0}" height="{0}" fill="{color}" rx="2" ry="2">"#,
                self.projection.size()
            )?;
            if let Some(ref unit) = self.unit {
                writeln!(io, "<title>{value} {unit} on {date}</title>")?;
            } else {
                writeln!(io, "<title>{value} {date}</title>")?;
            }
            writeln!(io, "</rect>")?;
        }
        writeln!(io, "</svg>")?;

        Ok(())
    }
}

/// Projection instances determine where each date goes in the image.
pub trait Projection {
    /// Transform a date to x,y coordinates.
    fn map(&self, date: NaiveDate) -> (u32, u32);

    /// Return the width and height of the image.
    fn dimensions(&self) -> (u32, u32);

    /// Return the size of each cell in the image.
    fn size(&self) -> u32;

    /// Return the range of dates in the image.
    fn range(&self) -> (NaiveDate, NaiveDate);
}

/// CalendarProjection creates a heatmap that looks like a calendar, with all
/// of the months in a row.
pub struct CalendarProjection {
    start: NaiveDate,
    end: NaiveDate,
    size: u32,
    padding: u32,
}

impl CalendarProjection {
    pub fn new(start: NaiveDate, end: NaiveDate, size: u32, padding: u32) -> Self {
        Self {
            start,
            end,
            size,
            padding,
        }
    }
}

impl Projection for CalendarProjection {
    fn map(&self, date: NaiveDate) -> (u32, u32) {
        let fom = date.with_day(1).unwrap();
        let year_number = (date.year() - self.start.year()) as u32;
        let week_number = (fom.weekday().num_days_from_monday() + date.day() - 1) / 7 + 1;

        let x = self.padding
            + ((date.month() - self.start.month()) % 12)
                * (self.padding + 7 * (self.size + self.padding))
            + date.weekday().num_days_from_monday() * (self.size + self.padding);

        let y = self.padding
            + week_number * (self.size + self.padding)
            + year_number * 6 * (self.size + self.padding);

        (x, y)
    }

    fn dimensions(&self) -> (u32, u32) {
        let years = (self.end.year() - self.start.year() + 1) as u32;
        let width = self.padding + ((self.size + self.padding) * 8) * 12;
        let height = self.padding + (self.size + self.padding) * 6 * years;
        (width, height)
    }

    fn size(&self) -> u32 {
        self.size
    }

    fn range(&self) -> (NaiveDate, NaiveDate) {
        (self.start, self.end)
    }
}

/// IsoProjection creates a denser heatmap where each column is a week of the
/// ISO calendar year and each row is a day of the week (starting with Monday).
pub struct IsoProjection {
    start: NaiveDate,
    end: NaiveDate,
    size: u32,
    padding: u32,
}

impl IsoProjection {
    pub fn new(start: NaiveDate, end: NaiveDate, size: u32, padding: u32) -> Self {
        Self {
            start,
            end,
            size,
            padding,
        }
    }
}

impl Projection for IsoProjection {
    fn map(&self, date: NaiveDate) -> (u32, u32) {
        let start = self.start.iso_week();
        let iso_date = date.iso_week();

        let year_number = (iso_date.year() - start.year()) as u32;

        let x = self.padding + iso_date.week0() * (self.size + self.padding);

        let y = self.padding
            + date.weekday().num_days_from_monday() * (self.size + self.padding)
            + year_number * 8 * (self.size + self.padding);

        (x, y)
    }

    fn dimensions(&self) -> (u32, u32) {
        let start = self.start.iso_week();
        let end = self.end.iso_week();
        let width = self.padding + (self.size + self.padding) * 53;
        let height =
            self.padding + (self.size + self.padding) * 8 * (end.year() - start.year() + 1) as u32;
        (width, height)
    }

    fn size(&self) -> u32 {
        self.size
    }

    fn range(&self) -> (NaiveDate, NaiveDate) {
        (self.start, self.end)
    }
}