use std::collections::{HashMap, BinaryHeap};
use std::ops::{Index, IndexMut};
use std::cmp::{Reverse};

#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
enum Tile {
    Wall,
    Ground,
}

impl Tile {
    pub fn cost(&self) -> u32 {
        match self {
            Tile::Wall => 0,
            Tile::Ground => 1,
        }
    }
}

#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
struct Cell {
    tile: Tile,
    x: usize,
    y: usize,
}

struct Grid {
    data: Vec<Cell>,
    size: usize,
}

impl Grid {
    pub fn new(size: usize) -> Self {
        let data = (0..size).map(|x| {
            (0..size).map(move |y| {
                Cell { tile: Tile::Ground, x, y }
            })
        }).flatten().collect::<Vec<_>>();
        Self {
            data,
            size: size,
        }
    }

    pub fn is_passable(&self, index: (usize, usize)) -> bool {
        self.in_bounds(index) && self[index].tile != Tile::Wall
    }

    pub fn in_bounds(&self, index: (usize, usize)) -> bool {
        index.0 < self.size && index.1 < self.size
    }

    pub fn neighbors_for(&self, index: (usize, usize)) -> Vec<Cell> {
        let direct = [
            (index.0 as i32 + 1, index.1 as i32 + 0), // right
            (index.0 as i32 + 0, index.1 as i32 + 1), // top
            (index.0 as i32 - 1, index.1 as i32 + 0), // left
            (index.0 as i32 + 0, index.1 as i32 - 1), // bottom
        ];
        let direct = direct.iter().filter(|i| {
            let i = *i;
            i.0 >= 0 && i.1 >= 0
        }).map(|i| {
            (i.0 as usize, i.1 as usize)
        }).filter(|i| {
            self.is_passable(*i)
        });

        direct.map(|i| {
            self[i]
        }).collect()
    }
}

impl Index<(usize, usize)> for Grid {
    type Output = Cell;
    fn index<'a>(&'a self, i: (usize, usize)) -> &'a Self::Output {
        &self.data[i.0 + self.size * i.1]
    }
}

impl IndexMut<(usize, usize)> for Grid {
    fn index_mut<'a>(&'a mut self, i: (usize, usize)) -> &'a mut Cell {
        &mut self.data[i.0 + self.size * i.1]
    }
}

fn heuristic(a: Cell, b: Cell) -> u32 {
    (a.x as i32 - b.x as i32).abs() as u32
        + (a.y as i32 - b.y as i32).abs() as u32
}

fn a_star_search(grid: &Grid, start: (usize, usize), goal: (usize, usize))
    -> Option<(HashMap<Cell, Cell>, HashMap<Cell, u32>, Vec<Cell>)>
{
    if grid.is_passable(start) && grid.is_passable(goal) {
        let start = grid[start];
        let goal = grid[goal];

        let mut frontier = BinaryHeap::new();
        // frontier.push(start, 0);
        frontier.push((Reverse(0), start));

        let mut came_from = HashMap::new();
        came_from.insert(start, start);

        let mut cost_so_far = HashMap::new();
        cost_so_far.insert(start, 0);

        while let Some((Reverse(priority), current)) = frontier.pop() {
            println!("current = {:?}, priority = {:?}", current, priority);
            if current == goal {
                break;
            }

            for next in grid.neighbors_for((current.x, current.y)) {
                let new_cost = cost_so_far[&current] + next.tile.cost();

                let next_cost = cost_so_far.get(&next);
                if next_cost.is_none() || new_cost < *next_cost.unwrap() {
                    println!("new_cost = {:?}", new_cost);
                    println!("heuristic = {:?}", heuristic(next, goal));
                    let priority = new_cost + heuristic(next, goal);
                    // frontier.push(next, priority);
                    frontier.push((Reverse(priority), next));
                    cost_so_far.insert(next, new_cost);
                    came_from.insert(next, current);
                }
            }
        }


        let mut path = Vec::new();
        let mut current = goal;
        while current != start {
            println!("current = {:?}", current);
            path.push(current);
            current = came_from[&current];
        }
        path.reverse();

        Some((came_from, cost_so_far, path))
    } else {
        None
    }
}

fn main() {
    let grid = Grid::new(20);

    // grid[(1,1)].tile = Tile::Wall;
    // grid[(1,2)].tile = Tile::Wall;
    // grid[(1,3)].tile = Tile::Wall;
    // grid[(1,4)].tile = Tile::Wall;

    if let Some((_came_from, cost_so_far, path)) = a_star_search(&grid, (0, 0), (19, 19)) {
        let mut graph = vec![0; grid.size * grid.size];

        for (cell, cost) in cost_so_far.iter() {
            graph[cell.x + cell.y * grid.size] = *cost;
        }

        println!("Cost grid");
        for x in 0..grid.size {
            for y in 0..grid.size {
                print!("{:3}", graph[x + y * grid.size]);
            }
            println!();
        }

        println!("Path grid");
        println!("{:?}", grid[(1,0)]);
        let mut graph = grid.data.iter().map(|c| match c.tile {
            Tile::Wall => '#',
            Tile::Ground => '-',
        }).collect::<Vec<_>>();
        for cell in path {
            graph[cell.x + cell.y * grid.size] = '*';
        }
        graph[0] = 's';
        graph[19 + 19 * grid.size] = 'e';

        for x in 0..grid.size {
            for y in 0..grid.size {
                print!("{}", graph[x + y * grid.size]);
            }
            println!();
        }

        // for x in _came_from {
        //     println!("{:?}", x);
        // }

        // for (cf, csf) in came_from.iter().zip(cost_so_far.iter()) {
        //     println!("{:?}: {:?}", cf, csf);
        // }
    }
}