Initial commitHEADmaster

1 files changed, 203 insertions, 0 deletions

diff --git a/src/heatmap.rs b/src/heatmap.rs
new file mode 100644
index 0000000..dd3ac3b
--- /dev/null
+++ b/src/heatmap.rs
@@ -0,0 +1,203 @@
+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)
+    }
+}