sb-rust-library 0.1.7

use image::{ImageError, RgbImage};
use std::ops::Range;

use crate::math::{Bounds, Point};
use super::color::{BLACK, Color, GREY, WHITE};
use super::plot_frame::PlotFrame;
use super::shapes::{Circle, Orientation};

/// Struct used to plot data and generate images.
///
/// A plot represents a single drawing area. The region of the image that data
/// is plotted to is refered to as the `canvas`, which by default, is inset
/// within the full image. The size of the canvas can be changed with
/// [`with_drawing_bounds`](#method.with_drawing_bounds).
///
/// A plot has two notions of coordinates: "logical" coordinates which are real-
/// valued coordinates that are plotted and "physical" coordinates which are the
/// integer locations of actual pixels. The range of logical coordinates
/// displayed can be changed with [`with_plotting_range`](#method.with_plotting_range).
///
/// Currently data can only be plotted with the low level primitives
/// [`Circle`](struct.Circle.html) and [`Line`](struct.Line.html). This makes
/// plotting somewhat cumbersome, but it also makes this system flexible.
///
/// # Examples
///
/// Basic sine wave without axis labels
///
/// ```
/// let mut p = Plot::new(300, 200)
///   .with_canvas_color(GREY)
///   .with_bg_color(BLUE)
///   .with_plotting_range(0.0..(2.0 * 3.14), -1.1..1.1)
///   .with_drawing_bounds(0.05..0.95, 0.1..0.95);
///
/// let x_values: Vec<f64> = (0..1000).map(|i| (i as f64) / 1000.0 * 2.0 * 3.14).collect();
/// let points: Vec<Point> = x_values.into_iter().map(|x| {
///   (x, x.sin()).into()
/// }).collect();
///
/// let circle = Circle::new(RED, 1);
/// p.draw_circles(&circle, &points);
/// p.save("sine-wave.bmp").unwrap();
/// ```
pub struct Plot {
  /// The width of the output image.
  pub width: i64, // more efficient to store as i64 due to `put_pixel_safe`
  /// The width of the output image.
  pub height: i64, // more efficient to store as i64 due to `put_pixel_safe`
  /// Internal only: TODO: Make private
  pub frame: PlotFrame,
  bg_color: Color,
  canvas_color: Color,
  canvas_bounds: Bounds<i64>,
  frame_color: Option<Color>,
  image: image::RgbImage,
}

impl Plot {

  /// Creates a new plot with the given image dimensions.
  pub fn new(width: u32, height: u32) -> Plot {
    let mut p = Plot {
      bg_color: WHITE,
      canvas_color: GREY,
      width: width as i64,
      height: height as i64,
      canvas_bounds: (0..(width as i64), 0..(height as i64)).into(),
      frame_color: Some(BLACK),
      frame: PlotFrame::new(
        (0..(width as i64), 0..(height as i64)).into(),
        (0.0..1.0, 0.0..1.0).into()),
      image: RgbImage::new(width, height),
    }
      .with_drawing_bounds(0.05..0.95, 0.05..0.95);
    p.clear();
    p
  }

  /// Clears all plotted primitives, redrawing the plot frame and canvas.
  pub fn clear(&mut self) {
    let x_bounds = self.canvas_bounds.x.clone();
    let y_bounds = self.canvas_bounds.y.clone();
    for y in 0..self.height {
      for x in 0..self.width {
        if x_bounds.contains(&x) && y_bounds.contains(&y) {
          self.image.put_pixel(x as u32, y as u32, self.canvas_color);
        } else {
          self.image.put_pixel(x as u32, y as u32, self.bg_color);
        }
      }
    }
    if let Some(color) = self.frame_color {
      for y in y_bounds.clone() {
        self.image.put_pixel(x_bounds.start as u32, y as u32, color);
        self.image.put_pixel(x_bounds.end as u32, y as u32, color);
      }
      for x in x_bounds.clone() {
        self.image.put_pixel(x as u32, y_bounds.start as u32, color);
        self.image.put_pixel(x as u32, y_bounds.end as u32, color);
      }
    }
  }

  /// Draws a line segment with the given logical coordinate endpoints.
  pub fn draw_line<P>(&mut self, point1: P, point2: P, color: Color)
  where P: Into<(f64, f64)> {
    let p1 = self.frame.pt_to_px(point1.into());
    let p2 = self.frame.pt_to_px(point2.into());
    let dx = p2.0 - p1.0;
    let dy = p2.1 - p1.1;

    if dx.abs() > dy.abs() {
      if dx > 0 {
        for x in p1.0..p2.0 {
          let y = p1.1 + dy * (x - p1.0) / dx;
          self.put_pixel_safe(x, y, color);
        }
      } else {
        for x in p2.0..p1.0 {
          let y = p1.1 + dy * (x - p1.0) / dx;
          self.put_pixel_safe(x, y, color);
        }
      }
    } else {
      if dy > 0 {
        for y in p1.1..p2.1 {
          let x = p1.0 + dx * (y - p1.1) / dy;
          self.put_pixel_safe(x, y, color);
        }
      } else {
        for y in p2.1..p1.1 {
          let x = p1.0 + dx * (y - p1.1) / dy;
          self.put_pixel_safe(x, y, color);
        }
      }
    }
    self.put_pixel_safe(p2.0, p2.1, color);
  }

  /// Draws a line segment with the given physical pixel endpoints.
  pub fn draw_pixel_line(&mut self, p1: (f64, f64), p2: (f64, f64), color: Color) {
    let dx = p2.0 - p1.0;
    let dy = p2.1 - p1.1;

    if dx.abs() > dy.abs() {
      if dx > 0.0 {
        let px_x_1 = p1.0.round() as i64;
        let px_x_2 = p2.0.round() as i64;
        for x in px_x_1..px_x_2 {
          let y = p1.1 + dy * (x as f64 - p1.0) / dx;
          self.put_pixel_safe(x, y.round() as i64, color);
        }
      } else {
        let px_x_1 = p1.0.round() as i64;
        let px_x_2 = p2.0.round() as i64;
        for x in px_x_2..px_x_1 {
          let y = p1.1 + dy * (x as f64 - p1.0) / dx;
          self.put_pixel_safe(x, y.round() as i64, color);
        }
      }
    } else {
      if dy > 0.0 {
        let px_y_1 = p1.1.round() as i64;
        let px_y_2 = p2.1.round() as i64;
        for y in px_y_1..px_y_2 {
          let x = p1.0 + dx * (y as f64 - p1.1) / dy;
          self.put_pixel_safe(x.round() as i64, y, color);
        }
      } else {
        let px_y_1 = p1.1.round() as i64;
        let px_y_2 = p2.1.round() as i64;
        for y in px_y_2..px_y_1 {
          let x = p1.0 + dx * (y as f64 - p1.1) / dy;
          self.put_pixel_safe(x.round() as i64, y, color);
        }
      }
    }
    self.put_pixel_safe(p2.0.round() as i64, p2.1.round() as i64, BLACK);
  }

  /// Plots data as line segments with the given logical positions and angles.
  pub fn draw_orientations(&mut self, orientation: &Orientation, positions: &[Point], angles: &[f64]) {
    for i in 0..positions.len() {
      orientation.draw(self, positions[i], angles[i]);
    }
  }

  /// Plots data as circles with the given logical positions.
  pub fn draw_circles(&mut self, circle: &Circle, positions: &[Point]) {
    for i in 0..positions.len() {
      circle.draw(self, positions[i]);
    }
  }

  /// Saves contents of plot as an image.
  pub fn save(&self, filename: &str) -> Result<(), ImageError> {
    self.image.save(filename)
  }

  /// Sets a pixel value, respecting bounds of plotting frame.
  pub fn put_pixel_safe(&mut self, x: i64, y: i64, c: Color) {
    if self.canvas_bounds.x.contains(&x) && self.canvas_bounds.y.contains(&y) {
      self.image.put_pixel(x as u32, y as u32, c)
    }
  }

  /// Changes the fraction of the plot image that the plotting canvas takes up.
  ///
  /// This accepts ranges of fractions between 0.0 and 1.0 indicating the
  /// portion of the plot that the plottable area takes up. All plotted data is
  /// displayed in this area. A bordering frame is drawn around the plotting
  /// canvas.
  pub fn with_drawing_bounds<R: Into<Range<f64>>>(mut self, x_bounds: R, y_bounds: R) -> Self {
    let w = self.width as f64;
    let h = self.height as f64;
    let x = x_bounds.into();
    let y = y_bounds.into();
    self.canvas_bounds.x = ((x.start * w + 0.5) as i64)..((x.end * w + 0.5) as i64);
    self.canvas_bounds.y = ((y.start * h + 0.5) as i64)..((y.end * h + 0.5) as i64);
    self.frame.set_drawing_bounds(self.canvas_bounds.x.clone(), self.canvas_bounds.y.clone());
    self
  }

  /// Changes the logical coordinate range displayed in the plot.
  ///
  /// For example, `(0.0, 2.0 * PI), (-1.0, 1.0)` would correspond to displaying
  /// a single period of a standard sine function.
  pub fn with_plotting_range<R: Into<Range<f64>>>(mut self, x_range: R, y_range: R) -> Self {
    self.frame.set_plotting_range(x_range, y_range);
    self
  }

  /// Changes the background color of the plot image (behind the canvas).
  pub fn with_bg_color(mut self, color: Color) -> Self {
    self.bg_color = color;
    self
  }

  /// Changes the background color of the plotting canvas.
  pub fn with_canvas_color(mut self, color: Color) -> Self {
    self.canvas_color = color;
    self
  }

  /// Changes the color of the border around the plotting canvas.
  pub fn with_frame_color(mut self, color: Option<Color>) -> Self {
    self.frame_color = color;
    self
  }

}