xiron 0.5.0

use crate::behaviour::traits::{Collidable, Drawable};
use crate::prelude::RESOLUTION;
use crate::utils::{interpolate_pose, XironError};

use image::{DynamicImage, GrayImage};
use macroquad::prelude::*;
use parry2d::math::{Isometry, Vector};
use serde::{Deserialize, Serialize};

use rayon::iter::IntoParallelIterator;
use rayon::iter::IntoParallelRefIterator;
use rayon::iter::ParallelIterator;

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
struct YamlConfig {
    pub image: String,
    pub mode: String,
    pub resolution: f32,
    pub origin: (f32, f32, f32),
    pub negate: i32,
    pub occupied_thresh: f32,
    pub free_thresh: f32,
}

impl Default for YamlConfig {
    fn default() -> Self {
        YamlConfig {
            image: "".into(),
            mode: "".into(),
            resolution: 0.0,
            origin: (0.0, 0.0, 0.0),
            negate: 0,
            occupied_thresh: 0.0,
            free_thresh: 0.0,
        }
    }
}

#[derive(Debug, Clone)]
pub struct OccupancyGrid {
    pub yaml_file_path: String,

    // Private Members
    config: YamlConfig,
    map_texture: Texture2D,
    map_image: GrayImage,
}

impl OccupancyGrid {
    pub async fn new(file_path: &str) -> Result<OccupancyGrid, XironError> {
        let file = std::fs::File::open(file_path);
        match file {
            Ok(file) => {
                let config: Result<YamlConfig, serde_yaml::Error> = serde_yaml::from_reader(file);
                match config {
                    Ok(config) => {
                        let map_name = &config.image;
                        let split_file_path: Vec<&str> = file_path.split('/').collect();

                        let mut map_file_path = String::new();
                        for i in 0..split_file_path.len() - 1 {
                            map_file_path += split_file_path[i];
                            map_file_path += "/";
                        }
                        map_file_path += map_name;

                        // Load raw image using `image` crate
                        let img_result = image::open(&map_file_path);
                        match img_result {
                            Ok(img) => {
                                let gray_image = img.to_luma8(); // convert to grayscale

                                // Convert to RGBA8 so it can be used by macroquad
                                let rgba_image =
                                    DynamicImage::ImageLuma8(gray_image.clone()).to_rgba8();
                                let texture = Texture2D::from_rgba8(
                                    rgba_image.width() as u16,
                                    rgba_image.height() as u16,
                                    &rgba_image,
                                );
                                texture.set_filter(FilterMode::Nearest);

                                Ok(OccupancyGrid {
                                    yaml_file_path: file_path.to_string(),
                                    config,
                                    map_texture: texture,
                                    map_image: gray_image,
                                })
                            }
                            Err(e) => Err(XironError::new(&format!("Failed to load image: {}", e))),
                        }
                    }
                    Err(e) => Err(XironError::new(&format!("YAML parse error: {}", e))),
                }
            }
            Err(e) => Err(XironError::new(&format!("Failed to open YAML: {}", e))),
        }
    }

    /// Get Pixel value given world coorinates
    pub fn world_to_pixels(&self, pos: &(f32, f32)) -> (i32, i32) {
        let resolution = self.config.resolution;
        let origin = &self.config.origin;
        let vmax = self.map_texture.height();

        let x_orig = origin.0;
        let y_orig = origin.1;

        let x = pos.0;
        let y = pos.1;

        let u = (-x_orig + x) / resolution;
        let v = vmax - (-y_orig + y) / resolution;

        return (u.floor() as i32, v.floor() as i32);
    }

    /// Get World coordinates given pixel values
    pub fn pixels_to_world(&self, pixels: &(i32, i32)) -> (f32, f32) {
        let resolution = self.config.resolution;
        let origin = &self.config.origin;
        let vmax = self.map_texture.height();

        let x_orig = origin.0;
        let y_orig = origin.1;

        let u = pixels.0;
        let v = pixels.1;

        let x = u as f32 * resolution + x_orig;
        let y = (vmax - v as f32) * resolution + y_orig;

        return (x, y);
    }

    pub fn check_within_limits(&self, pos: &(f32, f32)) -> bool {
        let (px, py) = self.world_to_pixels(pos);

        // println!("Pixel Value: {}, {}", px, py);

        // Convert to unsigned to compare against texture dimensions
        if px < 0 || py < 0 {
            return false;
        }

        let px = px as f32;
        let py = py as f32;

        let width = self.map_texture.width();
        let height = self.map_texture.height();

        px < width && py < height
    }

    pub fn get_occupancy_value(&self, pos: &(f32, f32)) -> f32 {
        let pixels = self.world_to_pixels(pos);
        let (px, py) = pixels;

        if px < 0 || py < 0 {
            return 1.0; // treat out-of-bounds as occupied
        }

        let (px, py) = (px as u32, py as u32);
        let (width, height) = self.map_image.dimensions();

        if px >= width || py >= height {
            return 1.0;
        }

        let pixel = self.map_image.get_pixel(px, py).0[0];

        // Normalize to [0.0, 1.0]
        pixel as f32 / 255.0
    }

    pub fn check_occupancy_status(&self, half_extents: &(f32, f32), pos: &(f32, f32, f32)) -> bool {
        let center = (pos.0, pos.1);
        let x_extent = half_extents.0;
        let y_extent = half_extents.1;

        let x_start = center.0 - x_extent;
        let x_end = center.0 + x_extent;
        let y_start = center.1 - y_extent;
        let y_end = center.1 + y_extent;

        let step = 0.1;

        // Check if all the end points are within limits
        let end_points = vec![
            (x_start, y_start),
            (x_start, y_end),
            (x_end, y_end),
            (x_end, y_start),
        ];

        for pt in end_points.iter() {
            let limits_check = self.check_within_limits(pt);
            let occupancy_at_edges = self.get_occupancy_value(pt);

            if (!limits_check) || (occupancy_at_edges <= 0.85) {
                return true;
            }
        }

        // Collect all (x, y) sample points within the AABB
        let mut sample_points = vec![];

        let mut x = x_start;
        while x <= x_end {
            let mut y = y_start;
            while y <= y_end {
                sample_points.push((x, y));
                y += step;
            }
            x += step;
        }

        // Parallel occupancy check
        let is_occupied = sample_points
            .par_iter()
            .any(|pt| !self.check_within_limits(pt) || self.get_occupancy_value(pt) <= 0.85);

        return is_occupied;
    }

    pub fn raycast_parallel(
        &self,
        origin: (f32, f32),
        direction: (f32, f32),
        max_throw: f32,
        throw_interpolation: f32,
    ) -> Option<f32> {
        let num_steps = (max_throw / throw_interpolation).ceil() as usize;

        (0..num_steps)
            .into_par_iter()
            .map(|i| i as f32 * throw_interpolation)
            .find_first(|&current_throw| {
                let world_x = origin.0 + current_throw * direction.0;
                let world_y = origin.1 + current_throw * direction.1;
                let pt = (world_x, world_y);

                let limits_check = self.check_within_limits(&pt);
                let occupancy_at_edges = self.get_occupancy_value(&pt);

                !limits_check || occupancy_at_edges <= 0.85
            })
    }
}

impl Drawable for OccupancyGrid {
    fn draw(&self, tf: fn((f32, f32)) -> (f32, f32)) {
        let map_texture = self.map_texture;

        let scaling_factor = self.config.resolution / RESOLUTION;

        let draw_size = vec2(
            map_texture.width() * scaling_factor,
            map_texture.height() * scaling_factor,
        );

        let image_top_left_in_world = self.pixels_to_world(&(0, 0));
        let world_in_screen = tf(image_top_left_in_world);

        draw_texture_ex(
            map_texture,
            world_in_screen.0,
            world_in_screen.1,
            WHITE,
            DrawTextureParams {
                dest_size: Some(draw_size),
                flip_y: false,
                ..Default::default()
            },
        );
    }

    fn draw_bounds(&self, _tf: fn((f32, f32)) -> (f32, f32)) {}
}

impl Collidable for OccupancyGrid {
    fn get_pose(&self) -> (f32, f32, f32) {
        (0.0, 0.0, 0.0)
    }

    fn get_shape(&self) -> Box<dyn parry2d::shape::Shape + Send + Sync> {
        return Box::new(parry2d::shape::Ball { radius: 0.0 });
    }

    fn get_max_extent(&self) -> f32 {
        0.0
    }

    fn collision_check_at_toi(
        &self,
        other: &dyn Collidable,
        start_pose: &(f32, f32, f32),
        end_pose: &(f32, f32, f32),
        _other_start_pose: Option<(f32, f32, f32)>,
        _other_end_pose: Option<(f32, f32, f32)>,
    ) -> Option<f32> {
        // Note: The Other {start, end} pose is used if other is also moving.
        // Since the Occupancy grid is not moving, there is no concern about this

        /*
         * 1.Interpolate the pos of object from start to end
         * 2 At each pose, get the AABB of the shape and check if the shape is in collision
         * 3. If in collision, return the interpolation idx
         * 4. If not return 1.0 ( to say that there is no collision)
         */

        // Here, use the other shape as the object shape and use the start_pose
        // and end_pose as object start pose and end pose

        let interpolation_dx = 0.1;
        let mut interpolation_current = 0.0;
        let interpolation_end = 1.0;

        while interpolation_current <= interpolation_end {
            let interpolated_pose = interpolate_pose(&start_pose, &end_pose, interpolation_current);

            let position = Isometry::new(
                Vector::new(interpolated_pose.0, interpolated_pose.1),
                interpolated_pose.2,
            );

            let half_extents = other.get_shape().compute_aabb(&position).half_extents();

            let collision_status_at_pose =
                self.check_occupancy_status(&(half_extents.x, half_extents.y), &interpolated_pose);
            if collision_status_at_pose {
                return Some(interpolation_current);
            }

            interpolation_current += interpolation_dx;
        }

        return None;
    }

    fn raycast(&self, ray: &parry2d::query::Ray) -> f32 {
        let origin = ray.origin;
        let direction = ray.dir;

        let max_throw = 15.0;
        let throw_interpolation = 0.05;

        let raycast_result = self.raycast_parallel(
            (origin.x, origin.y),
            (direction.x, direction.y),
            max_throw,
            throw_interpolation,
        );

        match raycast_result {
            Some(ray) => {
                return ray;
            }
            None => {
                return -10.0;
            }
        }
    }
}