use rand::prelude::{StdRng, SliceRandom};
use rand::{SeedableRng, Rng, thread_rng};
use std::fmt;
use std::fmt::Formatter;
use std::collections::{HashMap, HashSet};

pub const SIZE: usize = 16384;
pub const HOUSE_RANGE: usize = 500;

pub struct City {
    prices: Vec<u16>,
    buyable_house_count: usize
}

impl City {
    pub fn read_from_file(filename: &str) -> Self {
        let values = std::fs::read(filename).unwrap();
        let mut prices: Vec<u16> = Vec::new();

        for y in 0..SIZE {
            for x in 0..SIZE {
                let price = (values[(y * SIZE + x) * 2] as u16) | ((values[(y * SIZE + x) * 2 + 1] as u16) << 8);
                prices.push(price);
            }
        }

        City::new(prices)
    }

    pub fn new(prices: Vec<u16>) -> Self {
        let mut buyable_house_count = 0;
        for &price in &prices {
            if price > 0 {
                buyable_house_count += 1;
            }
        }
        City { prices, buyable_house_count }
    }

    pub fn get_price(&self, house: &House) -> u16 {
        self.prices[house.y * SIZE + house.x]
    }

    pub fn get_price_xy(&self, x: usize, y: usize) -> u16 {
        self.prices[y * SIZE + x]
    }

    pub fn is_house(&self, house: &House) -> bool {
        self.get_price(&house) > 0
    }

    pub fn is_house_xy(&self, x: usize, y: usize) -> bool {
        self.get_price_xy(x, y) > 0
    }

    pub fn get_house_count(&self) -> usize {
        self.buyable_house_count
    }
}

#[derive(Eq, PartialEq, Hash, Copy, Clone)]
pub struct House {
    x: usize,
    y: usize,
}

impl House {
    pub fn new(x: usize, y: usize) -> Self {
        House { x, y }
    }

    pub fn range_rectangle(&self) -> Rectangle {
        let top = if self.y <= HOUSE_RANGE { 0 } else { self.y - HOUSE_RANGE };
        let bottom = if self.y >= SIZE - 1 - HOUSE_RANGE { SIZE - 1 } else { self.y + HOUSE_RANGE };
        let left = if self.x <= HOUSE_RANGE { 0 } else { self.x - HOUSE_RANGE };
        let right = if self.x >= SIZE - 1 - HOUSE_RANGE { SIZE - 1 } else { self.x + HOUSE_RANGE };
        Rectangle {top, bottom, left, right}
    }
}

/// Rectangle - a 2D range with inclusive bounds
pub struct Rectangle {
    /// The smaller x coordinate.
    left: usize,
    /// The bigger x coordinate.
    right: usize,
    /// The smaller y coordinate.
    top: usize,
    /// The bigger y coordinate.
    bottom: usize,
}

impl fmt::Display for Rectangle {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "L{}-{}R T{}-{}B", self.left, self.right, self.top, self.bottom)
    }
}

impl Rectangle {
    pub fn is_inside(&self, x: usize, y: usize) -> bool {
        self.left <= x && x <= self.right && self.top <= y && y <= self.bottom
    }

    pub fn width(&self) -> usize {
        self.right - self.left
    }

    pub fn height(&self) -> usize {
        self.bottom - self.top
    }
}

pub struct HouseLayout<'a> {
    city: &'a City,
    reachable: Vec<u16>,
    houses: Vec<House>,
    reachable_houses: usize
}

impl<'a> HouseLayout<'a> {
    pub fn new(city: &'a City) -> Self {
        HouseLayout { city, reachable: vec![0; SIZE * SIZE], houses: Vec::new(), reachable_houses: 0 }
    }

    pub fn cover_count(&self, house: House) -> u16 {
        self.reachable[house.y * SIZE + house.x]
    }

    pub fn cover_count_xy(&self, x: usize, y: usize) -> u16 {
        self.reachable[y * SIZE + x]
    }

    pub fn is_covered(&self, house: House) -> bool {
        self.cover_count(house) > 0
    }

    pub fn add_house(&mut self, house: House) -> usize {
        let range_rect = house.range_rectangle();
        for y in range_rect.top..=range_rect.bottom {
            for x in range_rect.left..=range_rect.right {
                let index = y as usize * SIZE + x as usize;

                if self.reachable[index] == 0 && self.city.is_house_xy(x as usize, y as usize) {
                    self.reachable_houses += 1;
                }

                self.reachable[index] += 1;
            }
        }
        self.houses.push(house);
        self.houses.len() - 1
    }

    pub fn remove_house(&mut self, index: usize) {
        let house = self.houses.swap_remove(index);

        let range_rect = house.range_rectangle();
        for y in range_rect.top..=range_rect.bottom {
            for x in range_rect.left..=range_rect.right {
                let index = y as usize * SIZE + x as usize;

                self.reachable[index] -= 1;

                if self.reachable[index] == 0 && self.city.is_house_xy(x as usize, y as usize) {
                    self.reachable_houses -= 1;
                }
            }
        }
    }

    pub fn is_valid(&self) -> bool {
        self.reachable_houses == self.city.buyable_house_count
    }

    pub fn price(&self) -> u32 {
        get_price(self.city, &self.houses)
    }

    pub fn houses(&self) -> &Vec<House> {
        &self.houses
    }
}

fn dump_layout(layout: &HouseLayout, best_price: &mut Option<u32>, seed: u64) {
    let price = layout.price();
    if best_price.is_none() || price < best_price.unwrap() {
        *best_price = Some(price);
        eprintln!("Printing {} - new best", price);
        println!("New best!");
        println!("Price {}, seed {}", price, seed);
        print_houses(&layout.houses());
        println!();
    } else {
        eprintln!("Printing {}", price);
        println!("Price {}, seed {}", price, seed);
        print_houses(&layout.houses());
        println!();
    }
}


fn main() {
    let city = City::read_from_file("01.in");
    let mut best_price: Option<u32> = None;

    loop {
        let seed: u64 = thread_rng().gen();
        eprintln!("Starting seed {}", seed);
        let mut rng = StdRng::seed_from_u64(seed);
        let mut layout = HouseLayout::new(&city);
        eprintln!("Starting random population...");
        populate_random(&mut layout, &mut rng);
        eprintln!("Finished random init, price: {}", layout.price());
        loop {
            let mut improved = false;
            eprintln!("Starting moving individual houses...");
            if improve_move_individual_houses(&mut layout, &mut rng) {
                dump_layout(&layout, &mut best_price, seed);
                improved = true;
            }
            eprintln!("Finished moving individual houses...");
            eprintln!("Starting pairwise house merge...");
            if improve_merge_pairwise(&mut layout) {
                dump_layout(&layout, &mut best_price, seed);
                improved = true;
            }
            eprintln!("Finished pairwise house merge");
            if !improved {
                break;
            }
        }
    }
}

fn populate_random(layout: &mut HouseLayout, rng: &mut StdRng) {
    loop {
        loop {
            let x = rng.gen_range(0..SIZE);
            let y = rng.gen_range(0..SIZE);
            let house = House::new(x, y);
            if layout.city.is_house_xy(x, y) && !layout.is_covered(house) {
                layout.add_house(house);
                break;
            }
        }

        if layout.is_valid() {
            break;
        }
    }
}

fn improve_move_individual_houses(layout: &mut HouseLayout, mut rng: &mut StdRng) -> bool {
    let mut improved = false;
    let mut untried_houses = layout.houses().clone();
    untried_houses.shuffle(&mut rng);

    while untried_houses.len() > 0 {
        let house = untried_houses.pop().unwrap();
        let house_index = layout.houses().iter().position(|x| *x == house).unwrap();

        let move_rectangle = match get_valid_move_rectangle(&layout, house) {
            Ok(move_rectangle) => move_rectangle,
            Err(RectangleSearchError::Useless) => {
                //let old_price = layout.price();
                layout.remove_house(house_index);
                //let new_price = layout.price();
                //let price_diff = new_price as i64 - old_price as i64;
                //eprintln!(" candidate is valid, price diff: {}.", price_diff);
                //eprintln!("Removed a house (useless), diff {}", price_diff);
                //eprintln!("Improved price: {}", new_price);
                improved = true;
                untried_houses = layout.houses().clone();
                untried_houses.shuffle(&mut rng);
                continue;
            }
            _ => unreachable!()
        };

        // TODO: Not needed, can just store best
        let mut new_candidates = Vec::new();
        for new_y in move_rectangle.top..=move_rectangle.bottom {
            for new_x in move_rectangle.left..=move_rectangle.right {
                if layout.city.is_house_xy(new_x, new_y) && layout.city.get_price_xy(new_x, new_y) < layout.city.get_price(&house) {
                    new_candidates.push(House::new(new_x, new_y));
                }
            }
        }

        new_candidates.sort_by(|a, b| layout.city.get_price(&a).cmp(&layout.city.get_price(&b)));
        if new_candidates.len() == 0 {
            //eprintln!("Did not find candidate");
        } else {
            for (i, &candidate) in new_candidates.iter().enumerate() {
                //eprint!("Found candidate {}...", i);

                //let old_price = layout.price();
                layout.remove_house(house_index);
                layout.add_house(candidate);

                assert!(layout.is_valid());
                //let new_price = layout.price();
                //let price_diff = new_price as i64 - old_price as i64;
                //eprintln!(" candidate is valid, price diff: {}.", price_diff);
                //eprintln!("Improved price: {}", new_price);
                improved = true;
                untried_houses = layout.houses().clone();
                untried_houses.shuffle(&mut rng);
                break;
            }
        }
    }

    improved
}


pub fn improve_merge_pairwise(layout: &mut HouseLayout) -> bool {
    let mut improved = false;

    loop {
        // This here is a hack for being unable to modify the houses while looping through them.
        // We instead go through the houses repeatedly and remember which pairs we have already
        // tried by hashing their values because they can and do move throughout the layout Vec
        // as it's being modified.
        let mut checked = HashSet::new();

        let mut loop_improved = false;
        loop {
            let mut merge = None;

            'outer_houses: for i in 0..layout.houses().len() {
                for j in i + 1..layout.houses().len() {
                    let house1 = layout.houses()[i];
                    let house2 = layout.houses()[j];

                    let x_dist = (house1.x as i32 - house2.x as i32).abs() as usize;
                    let y_dist = (house1.y as i32 - house2.y as i32).abs() as usize;
                    if x_dist > 4 * HOUSE_RANGE || y_dist > 4 * HOUSE_RANGE {
                        // Never close enough to merge
                        continue;
                    }

                    if checked.contains(&(house1, house2)) || checked.contains(&(house2, house1)) {
                        continue;
                    } else {
                        checked.insert((house1, house2));
                    }

                    match get_valid_move_rectangle_multiple(&layout, &vec! {house1, house2}) {
                        Ok(rect) => {
                            let mut cheapest = None;
                            for y in rect.top..=rect.bottom {
                                for x in rect.left..=rect.right {
                                    if !layout.city.is_house_xy(x, y) { continue; }
                                    let price = layout.city.get_price_xy(x, y);
                                    match cheapest {
                                        None => cheapest = Some((x, y, price)),
                                        Some((_, _, cheapest_price)) if price < cheapest_price => cheapest = Some((x, y, price)),
                                        _ => {}
                                    };
                                }
                            }
                            if let Some((x, y, price)) = cheapest {
                                if price >= layout.city.get_price(&house1) + layout.city.get_price(&house2) {
                                    // Merging not worth
                                    //eprintln!("Merging not worth!");
                                } else {
                                    merge = Some((i, j, House::new(x, y)));
                                    break 'outer_houses;
                                }
                            }
                        }
                        Err(RectangleSearchError::Useless) => eprintln!("Found useless pair of houses, not solving!"),
                        Err(RectangleSearchError::Unsatisfiable) => {}
                    }
                }
            }

            if let Some((i, j, house)) = merge {
                let old_price = layout.price();
                assert!(i < j);
                layout.remove_house(j);
                layout.remove_house(i);
                layout.add_house(house);

                assert!(layout.is_valid());
                let new_price = layout.price();
                let price_diff = new_price as i32 - old_price as i32;
                eprintln!("Merged two houses, new price {}, diff {}", new_price, price_diff);
                improved = true;
                loop_improved = true;
            } else {
                break;
            }
        }

        if !loop_improved {
            break;
        }
    }

    improved
}


pub enum RectangleSearchError {
    Useless,
    Unsatisfiable
}

pub fn get_valid_move_rectangle_multiple(layout: &HouseLayout, houses: &Vec<House>) -> Result<Rectangle, RectangleSearchError> {
    // This is a generalization of get_valid_move_rectangle, it's basically the same thing,
    // just with a dynamic rectangles_containing_count

    // We first establish a bounding box for an that has to be covered if all houses are removed.
    let mut covered_rect: Option<Rectangle> = None;
    for house in houses {
        let range_rect = house.range_rectangle();
        for y in range_rect.top..=range_rect.bottom {
            for x in range_rect.left..=range_rect.right {
                // We count how many rectangles of houses contain this xy position.
                let mut rectangles_containing_count = 0;
                for house in houses {
                    let rect = house.range_rectangle();
                    if rect.is_inside(x, y) {
                        rectangles_containing_count += 1;
                    }
                }

                // If this house is covered by the exact amount of rectangles,
                // then removing all input houses would uncover this position.
                // It cannot be less than the rectangle count, and more means there
                // is another house covering it as well.
                if layout.cover_count_xy(x, y) == rectangles_containing_count && layout.city.is_house_xy(x, y) {
                    if let Some(cover) = &mut covered_rect {
                        cover.left = cover.left.min(x);
                        cover.right = cover.right.max(x);
                        cover.top = cover.top.min(y);
                        cover.bottom = cover.bottom.max(y);
                    } else {
                        covered_rect = Some(Rectangle { left: x, right: x, top: y, bottom: y });
                    }
                }
            }
        }
    };

    if covered_rect.is_none() {
        // Unnecessary set of houses.
        return Err(RectangleSearchError::Useless);
    }

    let covered_rect = covered_rect.unwrap();

    let height_margin = HOUSE_RANGE as i32 - covered_rect.height() as i32;
    let width_margin = HOUSE_RANGE as i32 - covered_rect.width() as i32;

    let top = (covered_rect.top as i32 - height_margin).max(0) as usize;
    let left = (covered_rect.left as i32 - width_margin).max(0) as usize;
    let bottom = (covered_rect.bottom + height_margin as usize).min(SIZE - 1);
    let right = (covered_rect.right + width_margin as usize).min(SIZE - 1);

    if top > bottom || left > right {
        // Unsatisfiable rectangle by one house
        return Err(RectangleSearchError::Unsatisfiable);
    }

    Ok(Rectangle { left, right, top, bottom })
}

pub fn get_valid_move_rectangle(layout: &HouseLayout, house: House) -> Result<Rectangle, RectangleSearchError> {
    // We first establish a bounding box for an that has to be covered if the house is removed.
    let mut covered_rect: Option<Rectangle> = None;

    let range_rect = house.range_rectangle();
    for y in range_rect.top..=range_rect.bottom {
        for x in range_rect.left..=range_rect.right {
            if layout.cover_count_xy(x, y) == 1 && layout.city.is_house_xy(x, y) {
                // This house is only covered by the house, it has to be covered from the new position as well.
                if let Some(cover) = &mut covered_rect {
                    cover.left = cover.left.min(x);
                    cover.right = cover.right.max(x);
                    cover.top = cover.top.min(y);
                    cover.bottom = cover.bottom.max(y);
                } else {
                    covered_rect = Some(Rectangle { left: x, right: x, top: y, bottom: y });
                }
            }
        }
    }

    if covered_rect.is_none() {
        return Err(RectangleSearchError::Useless)
    }

    let covered_rect = covered_rect.unwrap();

    // The distance of the rectangle from the original box tells us how much the house can move.
    let dist_left = covered_rect.left - range_rect.left;
    let dist_right = range_rect.right - covered_rect.right;
    let dist_top = covered_rect.top - range_rect.top;
    let dist_bottom = range_rect.bottom - covered_rect.bottom;

    let left = if house.x <= dist_right { 0 } else { house.x - dist_right };
    let right = if house.x >= SIZE - 1 - dist_left { SIZE - 1 } else { house.x + dist_left };
    let top = if house.y <= dist_bottom { 0 } else { house.y - dist_bottom };
    let bottom = if house.y >= SIZE - 1 - dist_top { SIZE - 1 } else { house.y + dist_top };

    let valid_move_rectangle = Rectangle {
        left, right, top, bottom
    };

    Ok(valid_move_rectangle)
}

pub fn get_neighbors(city: &City, house: &House) -> Vec<House> {
    let mut neighbors = Vec::new();
    let range_rect = house.range_rectangle();
    for y in range_rect.top..=range_rect.bottom {
        for x in range_rect.left..=range_rect.right {
            let house = House::new(x as usize, y as usize);
            if city.get_price(&house) > 0 {
                neighbors.push(house);
            }
        }
    }

    neighbors
}

fn print_houses(houses: &Vec<House>) {
    println!("{}", houses.len());
    for house in houses {
        println!("{} {}", house.y, house.x);
    }
}

fn get_price(city: &City, houses: &Vec<House>) -> u32 {
    let mut price = 0u32;
    for house in houses {
        price += city.get_price(&house) as u32;
    }

    price
}

fn is_valid(city: &City, houses: &Vec<House>) -> Option<u32> {
    let mut reachable = vec![false; SIZE * SIZE];
    let mut price = 0u32;

    for house in houses {
        assert!(city.prices[house.y * SIZE + house.x] > 0);

        let range_rect = house.range_rectangle();
        for y in range_rect.top..=range_rect.bottom {
            for x in range_rect.left..=range_rect.right {
                reachable[y as usize * SIZE + x as usize] = true;
            }
        }
        price += city.get_price(&house) as u32;
    }

    for y in 0..SIZE {
        for x in 0..SIZE {
            if !reachable[y * SIZE + x] && city.prices[y * SIZE + x] > 0 {
                return None;
            }
        }
    }

    Some(price)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn house_rectangle_at_min() {
        let house = House::new(0, 0);
        let rect = house.range_rectangle();
        assert_eq!(rect.top, 0);
        assert_eq!(rect.left, 0);
        assert_eq!(rect.right, HOUSE_RANGE);
        assert_eq!(rect.bottom, HOUSE_RANGE);
    }

    #[test]
    fn house_rectangle_at_max() {
        let house = House::new(SIZE - 1, SIZE - 1);
        let rect = house.range_rectangle();
        assert_eq!(rect.top, SIZE - 1 - HOUSE_RANGE);
        assert_eq!(rect.left, SIZE - 1 - HOUSE_RANGE);
        assert_eq!(rect.right, SIZE - 1);
        assert_eq!(rect.bottom, SIZE - 1);
    }

    #[test]
    fn house_rect_in_middle() {
        let house = House::new(SIZE / 2, SIZE / 2);
        let rect = house.range_rectangle();
        assert_eq!(rect.top, house.y - HOUSE_RANGE);
        assert_eq!(rect.left, house.x - HOUSE_RANGE);
        assert_eq!(rect.right, house.x + HOUSE_RANGE);
        assert_eq!(rect.bottom, house.y + HOUSE_RANGE);
    }
}