compb-dla-model/src/system/spaces/square_grid.rs

use std::ops::Add;
use num_integer::Integer;
use crate::system::{GriddedPosition, Position};
use serde::{Serialize, Deserialize};

#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)]
pub struct Grid2D {
    pub x: i32,
    pub y: i32,
}

impl Grid2D {
    pub fn in_direction(direction: u32, value: i32) -> Self {
        match direction {
            0 => Grid2D { x: value, y: 0 },
            1 => Grid2D { x: 0, y: value },
            _ => panic!("Invalid direction"),
        }
    }
}

impl Add for Grid2D {
    type Output = Grid2D;

    fn add(self, rhs: Self) -> Self::Output {
        Grid2D { x: self.x + rhs.x, y: self.y + rhs.y }
    }
}

impl Position for Grid2D {
    const DIM: usize = 2;

    fn zero() -> Self {
        Grid2D { x: 0, y: 0 }
    }

    fn abs(&self) -> f32 {
        (((self.x * self.x) + (self.y * self.y)) as f32).powf(0.5)
    }

    fn from_cartesian(cartesian: &[f32]) -> Self {
        Grid2D {
            x: cartesian[0] as i32,
            y: cartesian[1] as i32,
        }
    }

    fn to_cartesian(&self) -> Vec<f32> {
        vec![self.x as f32, self.y as f32]
    }
}

impl GriddedPosition for Grid2D {
    const NEIGHBOURS: u32 = 4;

    fn neighbour(&self, neighbour_index: u32) -> Self {
        let (dim, sign) = neighbour_index.div_rem(&2);
        let sign = if sign == 0 { 1 } else { -1 };
        let offset = Self::in_direction(dim, sign);

        self.clone() + offset
    }

    fn linear_index(&self, grid_size: u32) -> usize {
        let grid_size = grid_size as i32;

        assert!(self.x < grid_size && -(grid_size) < self.x);
        assert!(self.y < grid_size && -(grid_size) < self.y);

        let x = (self.x + (grid_size) / 2) as usize;
        let y = (self.y + (grid_size) / 2) as usize;

        let linear_index = grid_size as usize * y + x;

        if linear_index >= (grid_size * grid_size) as usize {
            eprintln!("AHHH SOMETHING WENT WRONG {:?} gives {}", self, linear_index);
        }

        return linear_index;
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct Grid3D {
    pub x: i32,
    pub y: i32,
    pub z: i32,
}

impl Grid3D {
    pub fn in_direction(direction: u32, value: i32) -> Self {
        match direction {
            0 => Grid3D { x: value, y: 0, z: 0 },
            1 => Grid3D { x: 0, y: value, z: 0 },
            2 => Grid3D { x: 0, y: 0, z: value },
            _ => panic!("Invalid direction"),
        }
    }
}

impl Add for Grid3D {
    type Output = Grid3D;

    fn add(self, rhs: Self) -> Self::Output {
        Grid3D {
            x: self.x + rhs.x,
            y: self.y + rhs.y,
            z: self.z + rhs.z,
        }
    }
}

#[test]
fn grid3_add_test() {
    assert_eq!(
        Grid3D::from_cartesian(&[0f32, 0f32, 0f32]) + Grid3D::from_cartesian(&[0f32, 1f32, 0f32]),
        Grid3D::from_cartesian(&[0f32, 1f32, 0f32])
    );
    assert_eq!(
        Grid3D::from_cartesian(&[5.0, 3.0, 1.0]) + Grid3D::from_cartesian(&[-2.0, 5.0, 100.0]),
        Grid3D::from_cartesian(&[3.0, 8.0, 101.0])
    );
}

impl Position for Grid3D {
    const DIM: usize = 3;

    fn zero() -> Self {
        Grid3D { x: 0, y: 0, z: 0 }
    }

    fn abs(&self) -> f32 {
        (((self.x * self.x) + (self.y * self.y) + (self.z * self.z)) as f32).powf(0.5)
    }

    fn from_cartesian(cartesian: &[f32]) -> Self {
        Self {
            x: cartesian[0] as i32,
            y: cartesian[1] as i32,
            z: cartesian[2] as i32,
        }
    }

    fn to_cartesian(&self) -> Vec<f32> {
        vec![self.x as f32, self.y as f32, self.z as f32]
    }
}

#[test]
fn grid3_neighbours_test() {
    let neighbours = (0..Grid3D::NEIGHBOURS)
        .map(|n| Grid3D::neighbour(&Grid3D::zero(), n))
        .collect::<Vec<_>>();

    assert!(neighbours.contains(&Grid3D { x: 1, y: 0, z: 0 }));
    assert!(neighbours.contains(&Grid3D { x: -1, y: 0, z: 0 }));
    assert!(neighbours.contains(&Grid3D { x: 0, y: 1, z: 0 }));
    assert!(neighbours.contains(&Grid3D { x: 0, y: -1, z: 0 }));
    assert!(neighbours.contains(&Grid3D { x: 0, y: 0, z: 1 }));
    assert!(neighbours.contains(&Grid3D { x: 0, y: 0, z: -1 }));
}

impl GriddedPosition for Grid3D {
    const NEIGHBOURS: u32 = 6;

    fn neighbour(&self, neighbour_index: u32) -> Self {
        let (dim, sign) = neighbour_index.div_rem(&2);
        let sign = if sign == 0 { 1 } else { -1 };
        let offset = Self::in_direction(dim, sign);

        self.clone() + offset
    }

    fn linear_index(&self, grid_size: u32) -> usize {
        let grid_size = grid_size as i32;

        assert!(self.x < grid_size && -(grid_size) < self.x);
        assert!(self.y < grid_size && -(grid_size) < self.y);
        assert!(self.z < grid_size && -(grid_size) < self.z);

        let half_grid = (grid_size) / 2;
        let x = (self.x + half_grid) as usize;
        let y = (self.y + half_grid) as usize;
        let z = (self.z + half_grid) as usize;

        let grid_size_usize = grid_size as usize;

        let linear_index =
            (grid_size_usize * grid_size_usize) * z
                + (grid_size_usize) * y
                + x;

        if linear_index >= (grid_size_usize * grid_size_usize * grid_size_usize) as usize {
            eprintln!("AHHH SOMETHING WENT WRONG {:?} gives {}", self, linear_index);
        }

        return linear_index;
    }
}