use std::cmp::min;
use std::collections::HashMap;
use std::f32::consts::TAU;
use std::ptr;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;

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;

use rand_gen::RngExt;
use rand_gen::SeedableRng;
use rand_gen::distr::Distribution;
use rand_gen::distr::Uniform;
use rand_gen::rngs::SmallRng;
use rayon::iter::IntoParallelIterator;
use rayon::iter::ParallelIterator;

use clap::Parser;

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

trait Mortal {
    fn collider(&self) -> (&Pose, &dyn Shape);
    fn kill(&self);
    fn living(&self) -> bool;

    fn draw(&self);
}

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

impl Mortal for Robot {
    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 collider(&self) -> (&Pose, &dyn Shape) {
        (&self.pose, &self.collider)
    }

    fn kill (&self) {
        self.alive.store(false, Ordering::Relaxed);
    }
    fn living(&self) -> bool {
        self.alive.load(Ordering::Relaxed)
    }
}

impl Robot {
    fn new(pose: Pose, team: Color) -> Self {
        Self { pose, alive: AtomicBool::new(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)
    }

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


    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))
    }

    fn crater(&self) -> (f32, f32, f32, Color) {
        let tl = self.pose.translation;
        (tl.x, tl.y, KILL_RADIUS, ORANGE)
    }

    fn draw_weapon(&self) {
        let tl = self.pose.translation;
        draw_circle(tl.x, tl.y, KILL_RADIUS, ORANGE);
    }

    fn weapon<'a>(&self, peers: impl Iterator<Item = &'a dyn Mortal>) {
        let weapon = Ball::new(KILL_RADIUS);

        for peer in peers {
            let (p_p, p_c) = peer.collider();

            if let Some(_) = contact(&self.pose, &weapon, p_p , p_c, 0.0).unwrap() {
                peer.kill();
            }

        }
    }

    pub fn policy(&mut self, sensor: &Sensor) -> bool{
        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, -0.02),
            Sensor::None => self.move_step(0.3, 0.01),
        }

        return *sensor == Sensor::NonKin;
    }
}

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

#[derive(PartialEq, Eq)]
enum Sensor {
    Food,
    Kin,
    NonKin,
    None
}

impl Mortal for Food { 
    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);
    }

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

    fn kill (&self) {
        self.alive.store(false, Ordering::Relaxed);
    }

    fn living(&self) -> bool {
        self.alive.load(Ordering::Relaxed)
    }
}

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


    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)
    }
}

struct Simulation {
    robots: Vec<Robot>,
    food: Vec<Food>,
    craters: Vec<(f32, f32, f32, Color)>,
    new_craters: Vec<(f32, f32, f32, Color)>,
    watchdog: i32,
    score: HashMap<[u32;4], i32>,
}

impl Simulation {
    fn new(seed: u64) -> Self {
        let mut rng = SmallRng::seed_from_u64(seed);
        let watchdog = 0; // iterations since something interesting
        let mut robots = Vec::new();

        let craters: Vec<(f32, f32, f32, Color)> = Vec::new();
        let new_craters: Vec<(f32, f32, f32, Color)> = Vec::new();

        let score = HashMap::new();

        for i in 5..20 {
            let angle = Uniform::new(0.0, TAU).unwrap();
            let rand_angle = angle.sample(&mut rng);
            let pose:Pose = Pose::from_parts(Vec2 { x: 100.0, y: i as f32*50.0 }, Rot2::from_angle(rand_angle));
            robots.push(Robot::new(pose, BLUE));

            let rand_angle = angle.sample(&mut rng);
            let pose:Pose = Pose::from_parts(Vec2 { x: 1200.0, y: i as f32*50.0 }, Rot2::from_angle(rand_angle));
            robots.push(Robot::new(pose, RED));
        }

        let mut food = Vec::new();

        for _ in 0..120 {
            let x = rng.random_range(200.0 .. 1100.0);
            let y = rng.random_range(200.0 .. 900.0);
            food.push(Food::new(Vec2 { x, y }));
        }
        Self { robots, food, craters, new_craters, watchdog, score }
    }

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

    fn sim_step(&mut self) -> Option<(i32, i32)>{
        let Simulation { robots, food, craters, new_craters, watchdog, score } = self;

        for food in &*food{
            if let Some(team) = food.detect(robots.iter()) {
                if food.living() {
                    score.entry(Self::color_hash(team)).and_modify(|s| *s +=1 ).or_insert(1);
                    food.kill();
                    *watchdog = 0;
                }
            }
        }

        food.retain(Mortal::living);
        robots.retain(Mortal::living);

        let mut sensors = Vec::new();
        for robot in &*robots {
            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; 
                }

            }

            sensors.push(sensor);
        }
        for idx in 0..robots.len() {
            let (before, after) = robots.split_at_mut(idx);
            let (robot, after) = after.split_first_mut().unwrap();
            let robots = before.iter().chain(after.iter());

            let robots = robots.map(|r| r as &dyn Mortal);
            let peers = robots.chain(food.iter().map(|f| f as &dyn Mortal));

            let mut death = robot.policy(&sensors[idx]);
            if robot.team == BLUE {
                death = false; // pacifists
            }
            if death {
                new_craters.push(robot.crater());
                robot.weapon(peers);
            }
        }

        *watchdog += 1;
        if food.is_empty() || *watchdog > 1000 {
            let red = score.get(&Self::color_hash(RED)).unwrap_or(&0);
            let blu = score.get(&Self::color_hash(BLUE)).unwrap_or(&0);
            return Some((*red, *blu));
        }
        None
    }

    async fn render(&mut self) {
        let Simulation { robots, food, craters, new_craters, watchdog: _, score } = self;

        clear_background(LIGHTGRAY);
        for crater in &*new_craters {
            draw_circle(crater.0, crater.1, crater.2, crater.3);
        }
        craters.append(&mut *new_craters);
        for crater in &*craters {
            draw_circle(crater.0, crater.1, crater.2, Color::new(0.75, 0.75, 0.75, 1.0));
        }
        for food in &*food {
            food.draw();

        }
        for robot in &*robots {
            robot.draw();
        }
        draw_text(format!("RED: {}", score.get(&Self::color_hash(RED)).unwrap_or(&0)), 20.0, 20.0, 30.0, DARKGRAY);
        draw_text(format!("BLUE: {}", score.get(&Self::color_hash(BLUE)).unwrap_or(&0)), 20.0, 40.0, 30.0, DARKGRAY);
        next_frame().await // TODO: framerate cap
    }
}

async fn run_simulation(seed: u64) -> (i32, i32) {
    let mut sim = Simulation::new(seed);

    let mut loo = 1;
    loop {
        let rendered = loo % 20 == 0;
        let result = sim.sim_step();
        
        if rendered {
            sim.render().await
        }

        if let Some(result) = result {
            return result;
        }

        loo += 1;
    }
}

fn run_headless_simulation(seed: u64) -> (i32, i32) {
    let mut sim = Simulation::new(seed);

    loop {
        if let Some(result) = sim.sim_step() {
            return result;
        }
    }
}


#[derive(Parser)]
#[command(version, about, long_about = None)]
struct Cli {
    /// run headless, allows multithreading
    #[arg(short, long)]
    no_gui: bool,

    /// number of samples
    #[arg(short, long)]
    run_count: Option<u64>,

    ///// Sets a custom config file
    //#[arg(short, long, value_name = "FILE")]
    //config: Option<PathBuf>,

    /// Turn debugging information on
    #[arg(short, long, action = clap::ArgAction::Count)]
    debug: u8,
}

#[macroquad::main("Adversarial Foraging Simulator")]
async fn main() {
    let cli = Cli::parse();
    let iters = cli.run_count.unwrap_or(1);
    
    let (red, blu) = if cli.no_gui { 
        (0..iters).into_par_iter().map(run_headless_simulation).reduce(|| (0,0), |a,b| (a.0 + b.0, a.1 + b.1))
    } else {
        let (mut rtot, mut btot) = (0, 0);
        for i in 0..iters  {
            let (red, blu) = run_simulation(i).await;
            rtot += red;
            btot += blu;

            println!("{} wins! RED: {red}, BLU: {blu}", if red > blu {"red"} else {"blue"});
        }
        (rtot,btot)
    };

    let red = red as f32 / iters as f32;
    let blu = blu as f32 / iters as f32;
    println!("{} wins on average! RED: {red}, BLU: {blu}", if red > blu {"red"} else {"blue"});

}