Initial commitHEADmaster

3 files changed, 327 insertions, 0 deletions

diff --git a/src/bin/tc-calmap.rs b/src/bin/tc-calmap.rs
new file mode 100644
index 0000000..566815a
--- /dev/null
+++ b/src/bin/tc-calmap.rs
@@ -0,0 +1,110 @@
+use std::io::{Write, stderr, stdin, stdout};
+use std::process::exit;
+
+use chrono::{Datelike, NaiveDate};
+use clap::{Parser, ValueEnum};
+
+use timechart::heatmap::{CalendarProjection, Heatmap, IsoProjection, Projection};
+use timechart::{Data, Result};
+
+/// Render a calendar-like heatmap from a list of dates and values.
+///
+/// Values are read from stdin. Output is written to stdout.
+#[derive(Parser, Debug)]
+struct Args {
+    /// The size of each square in the heat map (in px)
+    #[arg(long, default_value = "10")]
+    size: u32,
+
+    /// The padding between the squares (in px)
+    #[arg(long, default_value = "2")]
+    padding: u32,
+
+    /// The color scale
+    #[arg(
+        long,
+        default_values = ["#e7ebef","#deebf7","#c6dbef","#9ecae1","#6baed6","#2171b5"]
+    )]
+    colors: Vec<String>,
+
+    /// The unit description to use in the captions
+    #[arg(long)]
+    unit: Option<String>,
+
+    /// Use a log scale instead of a linear scale
+    #[arg(long, default_value = "true")]
+    log_scale: bool,
+
+    #[arg(
+        long("mode"),
+        value_enum,
+        default_value_t = Mode::Calendar
+    )]
+    mode: Mode,
+}
+
+#[derive(Debug, Clone, ValueEnum)]
+enum Mode {
+    Calendar,
+    Iso,
+}
+
+fn main() {
+    match run() {
+        Ok(_) => {}
+        Err(e) => {
+            let _ = writeln!(stderr(), "Error: {e}");
+            exit(1);
+        }
+    }
+}
+
+fn run() -> Result<()> {
+    let args = Args::parse();
+
+    // Read data from stdin.
+    let mut data = Data::new();
+    for line in stdin().lines() {
+        if let Some((ts, value)) = line?.split_once(' ') {
+            let date: NaiveDate = ts.parse()?;
+            let value: f64 = value.parse()?;
+            data.entry(date)
+                .and_modify(|v| *v += value)
+                .or_insert(value);
+        }
+    }
+
+    // Find the range of dates that we will render.
+    let (start, end) = data
+        .keys()
+        .copied()
+        .fold((NaiveDate::MAX, NaiveDate::MIN), |acc, v| {
+            let min = if acc.0 < v {
+                acc.0
+            } else {
+                NaiveDate::from_ymd_opt(v.year(), 1, 1).unwrap()
+            };
+            let max = if acc.1 > v {
+                acc.1
+            } else {
+                NaiveDate::from_ymd_opt(v.year(), 12, 31).unwrap()
+            };
+            (min, max)
+        });
+
+    // Choose a drawing style.
+    let projection: Box<dyn Projection> = match args.mode {
+        Mode::Calendar => Box::new(CalendarProjection::new(start, end, args.size, args.padding)),
+        Mode::Iso => Box::new(IsoProjection::new(start, end, args.size, args.padding)),
+    };
+
+    // Render the heatmap.
+    let heatmap = Heatmap {
+        projection,
+        colors: args.colors,
+        unit: args.unit,
+        log_scale: args.log_scale,
+    };
+
+    heatmap.render(&data, stdout())
+}
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)
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..d1ae1a9
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,14 @@
+use std::collections::BTreeMap;
+
+use chrono::NaiveDate;
+
+pub mod heatmap;
+
+/// The type used for the collection of date/value pairs.
+pub type Data = BTreeMap<NaiveDate, f64>;
+
+/// The Error type used within the library.
+pub type Error = Box<dyn std::error::Error>;
+
+/// The Result type used within the library.
+pub type Result<T> = std::result::Result<T, Error>;