use bland::{Figure, Marker, PaperSize, Stroke};
fn main() {
let cluster_a_x: Vec<f64> = (0..50).map(|i| 0.4 * i as f64 + 1.0).collect();
let cluster_a_y: Vec<f64> = cluster_a_x
.iter()
.enumerate()
.map(|(i, x)| 0.7 * x + 1.0 + ((i as f64 * 0.71) % 1.0 - 0.5) * 2.0)
.collect();
let cluster_b_x: Vec<f64> = (0..50).map(|i| 0.4 * i as f64 + 1.0).collect();
let cluster_b_y: Vec<f64> = cluster_b_x
.iter()
.enumerate()
.map(|(i, x)| 0.5 * x + 4.0 + ((i as f64 * 1.31 + 0.4) % 1.0 - 0.5) * 1.5)
.collect();
let fit_x = vec![1.0, 21.0];
let fit_y = vec![1.7, 15.7];
let fig = Figure::new()
.size(PaperSize::A5Landscape)
.title("Linear fit comparison")
.xlabel("input")
.ylabel("response")
.scatter(&cluster_a_x, &cluster_a_y, |s| {
s.label("control").marker(Marker::CircleOpen)
})
.scatter(&cluster_b_x, &cluster_b_y, |s| {
s.label("treated").marker(Marker::TriangleFilled)
})
.line(&fit_x, &fit_y, |l| l.label("y = 0.7x + 1").stroke(Stroke::DashDot))
.legend_bottom_right();
std::fs::create_dir_all("out").expect("create out/");
std::fs::write("out/scatter.svg", fig.to_svg()).expect("write svg");
println!("wrote out/scatter.svg");
}