use std::cmp::min;
use std::collections::HashMap;
use std::ptr;

use macroquad::prelude::*;
use parry2d::math::Pose;
use parry2d::math::Pose2;
use parry2d::math::Rot2;
use parry2d::math::Vec2;
use parry2d::query::contact;
use parry2d::shape::Ball;
use parry2d::shape::Segment;
use parry2d::shape::Shape;

const VIEW_DIST: f32 = 180.0;
const BOT_RADIUS: f32 = 16.0;

struct Robot {
    pose: Pose,
    alive: bool,
    collider: Ball,
    sensor: Sensor,
    team: Color,
}

impl Robot {
    fn new(position: Vec2, team: Color) -> Self {
        let pose:Pose = Pose::from_translation(position.to_array().into());
        Self { pose, alive: true, collider: Ball::new(BOT_RADIUS), sensor: Sensor::None, team}
    }

    fn move_step(&mut self, forward: f32, right: f32) {
        self.pose.translation += self.ahead() * forward;
        self.pose.rotation *= Rot2::from_angle(right);
    }

    pub fn ahead(&self) -> Vec2 {
        Vec2::new(self.pose.rotation.re, self.pose.rotation.im)
    }

    pub fn draw(&self) {
        let tl = self.pose.translation;
        let tl2 = self.view_end();
        draw_circle(tl.x, tl.y, BOT_RADIUS, BLACK);
        draw_line(tl.x, tl.y, tl2.x, tl2.y, 2.0, BLUE);
        draw_circle(tl.x, tl.y, BOT_RADIUS-1.0, self.team);
    }

    fn view_end(&self) -> Vec2 {
        self.pose.translation + self.ahead() * VIEW_DIST
    }

    pub fn collider(&self) -> (&Pose, &dyn Shape) {
        (&self.pose, &self.collider)
    }

    fn detect<'a>(&self, things: impl Iterator<Item = (&'a Pose, &'a dyn Shape)>) -> Option<f32> {
        let collider = Segment::new(self.pose.translation, self.view_end());
        things
            .map(|thing| contact(&Pose2::IDENTITY, &collider, thing.0, thing.1, 0.0).unwrap()).flatten()
            .map(|contact| contact.point1.distance(self.pose.translation)).reduce(|x: f32,y: f32| x.min(y))
    }

    pub fn policy(&mut self, sensor: &Sensor){
        match sensor {
            Sensor::Food => self.move_step(1.0, 0.0),
            Sensor::Kin => self.move_step(-0.3, 0.02),
            Sensor::NonKin => self.move_step(-0.3, 1.02),
            Sensor::None => self.move_step(0.3, 0.01),
        }
    }
}

struct Food {
    pose: Pose,
    collider: Ball,
    alive: bool,
}

enum Sensor {
    Food,
    Kin,
    NonKin,
    None
}

impl Food {
    fn new(position: Vec2) -> Self {
        let pose:Pose = Pose::from_translation(position.to_array().into());
        Self { pose, alive: true, collider: Ball::new(BOT_RADIUS) }
    }

    pub fn draw(&self) {
        let tl = self.pose.translation;
        draw_circle(tl.x, tl.y, BOT_RADIUS, BLACK);
        draw_circle(tl.x, tl.y, BOT_RADIUS-1.0, GREEN);
    }

    pub fn collider(&self) -> (&Pose, &dyn Shape) {
        (&self.pose, &self.collider)
    }

    fn detect<'a>(&self, things: impl Iterator<Item = &'a Robot>) -> Option<Color> {
        things
            .map(|robot| contact(&self.pose, &self.collider, &robot.pose, &robot.collider, 0.0).unwrap().map(|a| (a,robot.team))).flatten()
            .map(|(contact,team)| (contact.point1.distance(self.pose.translation),team)).reduce(|x ,y| if x.0 < y.0 {x} else{y} ).map(|(_, team)| team)
    }
}


#[macroquad::main("Adversarial Foraging Simulator")]
async fn main() {

    let mut robots = Vec::new();

    let mut score = HashMap::new();

    for i in 5..20 {
        robots.push(Robot::new(Vec2 { x: 150.0, y: i as f32*50.0 }, RED));
        robots.push(Robot::new(Vec2 { x: 1150.0, y: i as f32*50.0 }, BLUE));
    }

    let mut food = Vec::new();
    food.push(Food::new(Vec2 { x: 160.0, y: 100.0 }));
    food.push(Food::new(Vec2 { x: 360.0, y: 100.0 }));
    food.push(Food::new(Vec2 { x: 360.0, y: 300.0 }));
    food.push(Food::new(Vec2 { x: 360.0, y: 400.0 }));
    food.push(Food::new(Vec2 { x: 460.0, y: 400.0 }));
    food.push(Food::new(Vec2 { x: 260.0, y: 400.0 }));

    for j in 5..50 {
        for i in 5..20 {
            food.push(Food::new(Vec2 { x: i as f32*60.0, y: j as f32*50.0 }));
        }
    }

    fn color_hash(color: Color) -> [u32; 4] {
        color.to_vec().to_array().map(|f| f.to_bits())
    }


    let mut loo = 0;
    loop {
        clear_background(LIGHTGRAY);

        for food in &mut food{
            if let Some(team) = food.detect(robots.iter()) {
                if food.alive {
                    score.entry(color_hash(team)).and_modify(|s| *s +=1 ).or_insert(1);
                    println!("robot detected in contact");
                    food.alive = false;
                }
            }
            food.draw();
            // TODO: die upon collision with robot
        }

        food = food.into_iter().filter(|f| f.alive).collect(); // remove dead food

        let mut sensors = Vec::new();
        for robot in &*robots {
            println!("loop");
            let detect_food = robot.detect(food.iter().map(Food::collider));
            let detect_kin = robot.detect(robots.iter().filter(|o| o.team == robot.team).filter(|r| !ptr::eq(&**r , &**&robot)).map(Robot::collider));
            let detect_non_kin = robot.detect(robots.iter().filter(|o| o.team != robot.team).filter(|r| !ptr::eq(&**r , &**&robot)).map(Robot::collider));

            let mut sensor = Sensor::None;
            let mut closest: Option<f32> = None;

            for (dist, detection) in [(detect_food, Sensor::Food), (detect_kin, Sensor::Kin), (detect_non_kin, Sensor::NonKin)] {
                if let (Some(close), Some(dist)) = (closest, dist) {
                    if close < dist {continue;}
                }                 

                if let Some(_) = dist {
                    sensor = detection;
                    closest = dist; 
                }

            }

            println!("Food detection: {:?}",detect_food);
            println!("Kin detection: {:?}",detect_kin);
            println!("NonKin detection: {:?}",detect_non_kin);
            sensors.push(sensor);
        }
        for (idx, robot) in robots.iter_mut().enumerate() {
            robot.policy(&sensors[idx]);
            robot.draw();
        }
        draw_text(format!("RED: {}", score.get(&color_hash(RED)).unwrap_or(&0)), 20.0, 20.0, 30.0, DARKGRAY);
        draw_text(format!("BLUE: {}", score.get(&color_hash(BLUE)).unwrap_or(&0)), 20.0, 90.0, 30.0, DARKGRAY);
        loo += 1;
        if loo %2 == 0 {
            next_frame().await // TODO: framerate cap
        }
    }
}