use crate::city;
use crate::db::{SqliteLayoutDB, SavedLayout, MergeLowerBound, LayoutDB};
use crate::city::{City, House, HouseDistances};
use itertools::Itertools;
use itertools::iproduct;
use std::collections::{VecDeque, HashMap};
use std::collections::vec_deque::Iter;

struct LeftState<'a> {
    layout: &'a SavedLayout,
    sorted_houses: Vec<House>,
    line: LeftLine<'a>
}

struct RightState<'a> {
    layout: &'a SavedLayout,
    line: RightLine<'a>
}

pub struct CompatibilityCache {
    map: HashMap<(usize, usize, usize, usize, bool), bool>
}

impl CompatibilityCache {
    pub fn new() -> CompatibilityCache {
        CompatibilityCache {
            map: HashMap::new()
        }
    }

    pub fn is_compatible(&self, left_layout: &SavedLayout, right_layout: &SavedLayout, left_update_index: usize, right_update_index: usize, y_axis: bool) -> Option<&bool> {
        self.map.get(&(left_layout.id(), right_layout.id(), left_update_index, right_update_index, y_axis))
    }

    pub fn set_compatible(&mut self, left_layout: &SavedLayout, right_layout: &SavedLayout, left_update_index: usize, right_update_index: usize, y_axis: bool, compatible: bool) {
        self.map.insert((left_layout.id(), right_layout.id(), left_update_index, right_update_index, y_axis), compatible);
    }
}

fn choose_layouts<TDB: LayoutDB>(db: &TDB, city: &City, top_layout_count: usize) -> Vec<SavedLayout> {
    // We do a secondary sort by ID to make the result stable in case of ties
    let top_deduplicated_layouts: Vec<SavedLayout> = db.layouts().iter()
        .sorted_by_key(|x| (city::get_price(&city, x.houses()), x.id()))
        .unique_by(|x| x.houses().iter().sorted_by_key(|x| (x.x, x.y)).collect::<Vec<_>>())
        .map(|layout| (*layout).clone())
        .take(top_layout_count)
        .collect();

    top_deduplicated_layouts
}

pub fn iterate_combines<TDB: LayoutDB>(mut db: &mut TDB, top_layout_count: usize, city: &City, mut cache: &mut CompatibilityCache, print_progress: bool) {
    let mut no_improvements_in_a_row = 0;
    if print_progress { eprintln!("Building a transposed city..."); }
    let transposed_city = transpose_city(&city);
    if print_progress { eprintln!("Finished building a transposed city"); }

    loop {
        if no_improvements_in_a_row == 2 { break; }
        if print_progress { eprintln!("Starting to combine {} top houses DB; vertical cuts", top_layout_count); }

        {
            let chosen_layouts = choose_layouts(db, &city, top_layout_count);
            if print_progress { eprintln!("Building right distances..."); }
            let right_distances = city::build_distances_right(&city);
            if print_progress { eprintln!("Finished building right distances"); }
            if create_new_best_combination(&city, &chosen_layouts, &chosen_layouts, db, &mut cache, &right_distances, false, print_progress) {
                no_improvements_in_a_row = 0;
            } else {
                no_improvements_in_a_row += 1;
            }
            if print_progress { eprintln!("Finished vertical cuts, improvement: {}", no_improvements_in_a_row == 0); }
        }

        if no_improvements_in_a_row == 2 { break; }

        {
            let chosen_layouts = choose_layouts(db, &city, top_layout_count);
            let transposed_chosen_layouts: Vec<_> = chosen_layouts.iter().map(|x| transpose_saved_layout(x)).collect();
            if print_progress { eprintln!("Building right distances (transposed)..."); }
            let right_distances_transposed = city::build_distances_right(&transposed_city);
            if print_progress { eprintln!("Finished building right distances"); }

            if print_progress { eprintln!("Starting to combine {} top houses DB; horizontal cuts", top_layout_count); }
            if create_new_best_combination(&transposed_city, &transposed_chosen_layouts, &transposed_chosen_layouts, db, &mut cache, &right_distances_transposed, true, print_progress) {
                no_improvements_in_a_row = 0;
            } else {
                no_improvements_in_a_row += 1;
            }
            if print_progress { eprintln!("Finished horizontal cuts, improvement: {}", no_improvements_in_a_row == 0); }
        }

        if no_improvements_in_a_row == 2 { break; }
    }
}

fn transpose_city(city: &City) -> City {
    let mut transposed_prices = vec![0u16; city.width() * city.height()];
    for y in 0..city.height() {
        for x in 0..city.width() {
            // Sorry, cache! Not worth optimizing with blocks,
            // this is not going to be ran often.
            transposed_prices[x * city.height() + y] = city.get_price_xy(x, y);
        }
    }

    City::new(transposed_prices, city.height(), city.width())
}

fn transpose_saved_layout(layout: &SavedLayout) -> SavedLayout {
    let mut transposed = Vec::new();
    for house in layout.houses() {
        transposed.push(House::new(house.y, house.x));
    }

    SavedLayout::new(layout.id(), transposed)
}

pub fn create_new_best_combination<TDB: LayoutDB>(city: &City, left_layouts: &Vec<SavedLayout>, right_layouts: &Vec<SavedLayout>, db: &mut TDB, cache: &mut CompatibilityCache, right_distances: &HouseDistances, transposed: bool, print_progress: bool) -> bool {
    let mut best_price = left_layouts.iter().chain(right_layouts.iter())
        .map(|layout| city::get_price(&city, layout.houses()))
        .min();

    let mut improved = false;

    let mut lefts = Vec::new();
    let mut rights = Vec::new();

    for left_layout in left_layouts {
        // We make sure that there is at least one other layout we want to compare this one with
        // to avoid unnecessary house updates
        let mut needed = false;
        for right_layout in right_layouts {
            if let Some(bound) = db.get_merge_lower_bound(left_layout, right_layout, transposed) {
                if bound < best_price.expect("No best price set while lower bounds exist") {
                    needed = true;
                    break;
                }
            } else {
                needed = true;
                break;
            }
        }

        if !needed {
            continue;
        }

        // Sorted in reverse so we can remove from the end
        let sorted_houses: Vec<_> = left_layout.houses().iter().sorted_by(|h1, h2| h2.x.cmp(&h1.x)).map(|x| *x).collect();

        lefts.push(LeftState {
            layout: &left_layout,
            sorted_houses,
            line: LeftLine::new(&city)
        });
    }

    for right_layout in right_layouts {
        // We make sure that there is at least one other layout we want to compare this one with
        // to avoid unnecessary house updates
        let mut needed = false;
        for left_layout in left_layouts {
            if let Some(bound) = db.get_merge_lower_bound(left_layout, right_layout, transposed) {
                if bound < best_price.expect("No best price set while lower bounds exist") {
                    needed = true;
                }
            } else {
                needed = true;
            }
        }

        if !needed {
            continue;
        }

        let mut sorted_houses: Vec<_> = right_layout.houses().iter().sorted_by(|h1, h2| h2.x.cmp(&h1.x)).map(|x| *x).collect();

        let mut line = RightLine::new(&city);
        // Make sure that we include all houses initially
        while let Some(house) = sorted_houses.pop() {
            line.add_house(house, &city);
        }
        rights.push(RightState {
            layout: right_layout,
            line
        });
    }

    let axis = if transposed { "y" } else { "x" };

    // We construct a quick lookup table for skipped pairs of layouts - these are kept the same
    // the entire run. We also assume that the best price does not get worse throughout the run.
    let checked_pair_table: Vec<Vec<bool>> = lefts.iter()
        .map(|left| rights.iter().map(|right| {
            if left.layout.id() == right.layout.id() {
                // We do not want to combine a layout with itself, it would never improve the score.
                return false;
            }
            if let Some(bound) = db.get_merge_lower_bound(left.layout, right.layout, transposed) {
                if bound >= best_price.expect("No best price set while lower bounds exist") {
                    // This combination can never attain a price better than the current best.
                    return false;
                }
            }
            return true;
        }).collect()).collect();

    if print_progress {
        let total_pairs = checked_pair_table.iter().flatten().count();
        let checked_pairs = checked_pair_table.iter().flatten().filter(|x| **x).count();
        eprintln!("Checking {}/{} layout pairs", checked_pairs, total_pairs);
    }

    // x is the last left coordinate, x+1 is right
    for x in 0..city.width() {
        if print_progress { eprintln!("Starting {} {}", axis, x); }

        // Update the lines
        for left in lefts.iter_mut() {
            while let Some(house) = left.sorted_houses.last() {
                if house.x == x {
                    left.line.add_house(*house, &city);
                    left.sorted_houses.pop();
                } else {
                    break;
                }
            }
        }

        for right_line in rights.iter_mut() {
            right_line.line.remove_houses(x, &city);
        }

        // Check compatibility of lines
        if x == 0 {
            // Cannot check this due to limitations in the implementation of LeftLine,
            // it wouldn't be very interesting anyway.
            continue;
        }

        let pairs = iproduct!(lefts.iter().enumerate(), rights.iter().enumerate())
            .filter(|((left_i, _), (right_i, _))| checked_pair_table[*left_i as usize][*right_i as usize])
            .map(|((left_i, left), (right_i, right))| (left, right, left.line.price + right.line.price, left_i, right_i))
            .sorted_by(|(_, _, price1, _, _), (_, _, price2, _, _)| price1.cmp(&price2));

        let mut compatibles = 0;
        let mut incompatibles = 0;
        let mut cache_hits = 0;
        let mut best_incompatible_price = best_price;
        for (left, right, price, left_i, right_i) in pairs {
            if let Some(min_price) = best_price {
                if price >= min_price {
                    break;
                }
            }
            let compatible = match cache.is_compatible(left.layout, right.layout, left.line.last_update_x, right.line.last_update_x, transposed) {
                None => {
                    let compatible = is_compatible(city, right_distances, &left.line, &right.line);
                    cache.set_compatible(left.layout, right.layout, left.line.last_update_x, right.line.last_update_x, transposed, compatible);
                    compatible
                }
                Some(compatible) => {
                    cache_hits += 1;
                    *compatible
                }
            };
            if compatible {
                if best_price.is_none() || price < best_price.unwrap() {
                    best_price = Some(price);
                    if print_progress {
                        eprintln!("{} - new best score, cut on {} {}, left {} - right {}, printing", price, axis, x, left_i, right_i);
                        println!("{} - new best score, cut on {} {}, left {} - right {}", price, axis, x, left_i, right_i);
                    }
                    let mut new_houses: Vec<_> = left.line.houses().copied().chain(right.line.houses().copied()).collect();
                    if transposed {
                        new_houses = transpose_layout(&new_houses);
                    }
                    if print_progress {
                        println!("{}", new_houses.len());
                        for house in &new_houses {
                            println!("{} {}", house.y, house.x);
                        }
                    }

                    // We only add best results to avoid overfilling the database with similar layouts
                    db.add_layout(&new_houses, false);
                    improved = true;
                }
                compatibles += 1;

                // All other pairs would be more expensive
                break;
            } else {
                best_incompatible_price = Some(best_incompatible_price.unwrap_or(u32::MAX).min(price));
                incompatibles += 1;
            }
        }

        if print_progress { eprintln!("{} incompatibles checked before {} compatible ({} cache hits) [{}-{}]", incompatibles, compatibles, cache_hits, best_incompatible_price.unwrap_or(u32::MAX), best_price.unwrap_or(u32::MAX)); }
    }

    if let Some(lowest_price) = best_price {
        let mut new_bounds = Vec::new();

        for left in &lefts {
            for right in &rights {
                new_bounds.push(MergeLowerBound::new(left.layout.id(), right.layout.id(), transposed, lowest_price));
            }
        }
        db.add_merge_lower_bounds(new_bounds);
    }

    improved
}

fn is_compatible(city: &City, right_distances: &HouseDistances, left: &LeftLine, right: &RightLine) -> bool {
    for y in 0..city.height() {
        let max_left_covered_x = left.get_max_covered_x(y);
        let min_right_covered_x = right.get_min_covered_x(y);

        if let Some(right_house) = right_distances.get_closest_house_xy(max_left_covered_x, y) {
            if right_house.x < min_right_covered_x {
                return false;
            }
        }
    }

    true
}

struct LeftLine<'a> {
    covers: Vec<usize>,
    houses: Vec<House>,
    price: u32,
    last_update_x: usize,
    city: &'a City
}

struct RightLine<'a> {
    covers: Vec<usize>,
    houses: VecDeque<House>,
    price: u32,
    last_update_x: usize,
    city: &'a City
}

impl<'a> LeftLine<'a> {
    pub fn new(city: &'a City) -> Self {
        // XXX: Careful, default of 0 includes covering first vertical line
        let covers = vec![0; city.height()];
        let houses = Vec::new();
        LeftLine { covers, houses, price: 0, last_update_x: 0, city }
    }

    pub fn add_house(&mut self, house: House, city: &City) {
        let range_rect = house.range_rectangle(city);
        for y in range_rect.top..=range_rect.bottom {
            // Should always be the max variant
            self.covers[y] = self.covers[y].max(range_rect.right);
        }
        self.price += city.get_price(house) as u32;
        self.houses.push(house);
        self.last_update_x = house.x;
    }

    pub fn get_max_covered_x(&self, y: usize) -> usize {
        self.covers[y]
    }

    pub fn get_side_price(&self) -> u32 {
        self.price
    }

    pub fn houses(&self) -> std::slice::Iter<'_, House> {
        self.houses.iter()
    }
}

impl<'a> RightLine<'a> {
    pub fn new(city: &'a City) -> Self {
        let covers = vec![usize::MAX; city.height()];
        let houses = VecDeque::new();
        RightLine { covers, houses, price: 0, last_update_x: 0, city }
    }

    pub fn add_house(&mut self, house: House, city: &City) {
        // Added houses have to always be ordered by x
        let range_rect = house.range_rectangle(city);
        for y in range_rect.top..=range_rect.bottom {
            self.covers[y] = self.covers[y].min(range_rect.left);
        }

        self.houses.push_back(house);
        self.price += city.get_price(house) as u32;
        self.last_update_x = house.x;
    }

    pub fn remove_houses(&mut self, x: usize, city: &City) {
        // Has to be called with x, x+1, x+2...
        while let Some(house) = self.houses.front() {
            if house.x == x {
                let removed_house = self.houses.pop_front().unwrap();
                let removed_rect = removed_house.range_rectangle(city);

                // Remove the now-outdated distances around the removed house
                for y in removed_rect.top..=removed_rect.bottom {
                    self.covers[y] = usize::MAX;
                }

                // Update distances around the removed house if the area of any houses
                // intersects the removed area
                for house in &self.houses {
                    let house_rect = house.range_rectangle(city);
                    let y_intersection = if removed_house.y < house.y {
                        house_rect.top..=removed_rect.bottom
                    } else {
                        removed_rect.top..=house_rect.bottom
                    };

                    for y in y_intersection {
                        self.covers[y] = self.covers[y].min(house_rect.left);
                    }
                }

                self.price -= city.get_price(removed_house) as u32;
                self.last_update_x = removed_house.x;
            } else {
                break;
            }
        }
    }

    pub fn get_min_covered_x(&self, y: usize) -> usize {
        self.covers[y].min(self.city.width())
    }

    pub fn get_side_price(&self) -> u32 {
        self.price
    }

    pub fn houses(&self) -> Iter<'_, House> {
        self.houses.iter()
    }
}

pub fn transpose_layout(houses: &Vec<House>) -> Vec<House> {
    let mut transposed = Vec::new();
    for house in houses {
        transposed.push(House::new(house.y, house.x));
    }

    transposed
}