diff --git a/grid_image.png b/grid_image.png
index 32e6587..90b5563 100644
Binary files a/grid_image.png and b/grid_image.png differ
diff --git a/src/function.rs b/src/function.rs
index 7964f4b..e201e8f 100644
--- a/src/function.rs
+++ b/src/function.rs
@@ -1,4 +1,4 @@
-pub fn default_points(points: &mut [[f64;2];3]) {
+fn default_points(points: &mut [[f64;2];3]) {
     if points.iter().all(|row| row.iter().all(|&x| x == 0.0)) {
         *points = [[-1.0,1.0], [2.0,4.0], [3.0,9.0]];
     }
@@ -8,7 +8,7 @@ pub fn group_by_x(arr: &mut [[f64;2];3]) {
     arr.sort_by(|a, b| a[0].partial_cmp(&b[0]).unwrap_or(std::cmp::Ordering::Equal));
 }
 
-pub fn get_points(args: Vec<String>) -> [[f64;2];3] {
+pub fn get_args(args: Vec<String>) -> ([[f64;2];3], Option<u32>) {
     let mut points: [[f64; 2]; 3] = [[0.0; 2]; 3];
 
     for (i, part) in args.iter().enumerate().take(3) {
@@ -26,10 +26,19 @@ pub fn get_points(args: Vec<String>) -> [[f64;2];3] {
         }
     }
 
+    let mut img_size: Option<u32> = None;
+    if args.len() > 3 {
+        img_size = Some(args.get(3).unwrap().parse().unwrap());
+    }
+
     #[cfg(debug_assertions)]
     default_points(&mut points);
 
-    points
+    (points, img_size)
+}
+
+pub fn get_args_new(args: Vec<String>) -> (Vec<[f64;3]>, Option<u32>) {
+    todo!()
 }
 
 fn nearest_to_zero(arr: &[[f64;2];3]) -> usize {
diff --git a/src/image.rs b/src/img.rs
similarity index 99%
rename from src/image.rs
rename to src/img.rs
index 0c8e285..0035883 100644
--- a/src/image.rs
+++ b/src/img.rs
@@ -17,5 +17,3 @@ pub fn get_image(width: u32, height: u32) -> ImageBuffer<Rgb<u8>, Vec<u8>> {
     }
     img
 }
-
-
diff --git a/src/main.rs b/src/main.rs
index 527ef8f..7470087 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,13 +1,15 @@
 use std::env;
+use image::Rgb;
 
 mod function;
-mod image;
+mod img;
 use crate::function::*;
+use crate::img::*;
 
 fn main() {
-    let args: Vec<String> = env::args().skip(1).take(3).collect();
+    let args: Vec<String> = env::args().skip(1).take(4).collect();
 
-    let mut points = get_points(args);
+    let (mut points, img_size) = get_args(args);
     group_by_x(&mut points);
 
     let m1 = (points[0][1] - points[1][1]) / (points[0][0] - points[1][0]);
@@ -25,14 +27,14 @@ fn main() {
     println!("{}", function);
 
     let rust_function = function.parse::<meval::Expr>().unwrap().bind("x").unwrap();
-    let result = rust_function(5.0);
-    println!("I can also calculate with this function: f(5) = {}", result);
+    let result = rust_function(3.0);
+    println!("I can also calculate with this function: f(3) = {}", result);
 
-    /*
-    let width = 7000;
-    let height = 7000;
-    let color = [255, 0, 0];
-    let mut img = get_image(4000, 4000);
+    // /*
+    let width = img_size.unwrap_or(5000);
+    let height = img_size.unwrap_or(5000);
+    let color = [0, 255, 0];
+    let mut img = get_image(width, height);
     let step: f64 = format!("0.{}1", "0".repeat(6)).parse().unwrap();
     let mut points = vec![];
 
@@ -48,13 +50,12 @@ fn main() {
 
     let mut x;
     let mut y;
-    // let mut x_result_tmp;
-    // let mut y_result_tmp;
     for point in points {
         // NOTE: convert later to u32 because deviding may change actual value
         x = point * 1000.0_f64;
         y = rust_function(point) * 1000.0;
-        if width as f64 / 2.0 + y < (width - 5) as f64 {
+        if width as f64 / 2.0 + x < (width - 10) as f64
+        && width as f64 / 2.0 - y < (width - 10) as f64 {
             img.put_pixel(
                 (width as f64 / 2.0 + x).round() as u32,
                 (height as f64 / 2.0 - y) as u32,
@@ -69,7 +70,6 @@ fn main() {
     }
 
     img.save("grid_image.png").expect("Failed to save image");
-    */
 }