From 9bd654a8e6ac1c34308bbe0e466a1ff06ab8ad1a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ji=C5=99=C3=AD=20Sejkora?= <sejsel@sejsel.cz>
Date: Mon, 11 Jan 2021 05:25:40 +0100
Subject: [PATCH] Draw from the layout DB

---
 src/city.rs       |  2 +-
 src/db.rs         |  7 ++++++
 src/main.rs       | 14 +++++++++--
 src/population.rs | 61 ++++++++++++++++++++++++++++++++++++++++++++---
 4 files changed, 78 insertions(+), 6 deletions(-)

diff --git a/src/city.rs b/src/city.rs
index ad723bc..964a80f 100644
--- a/src/city.rs
+++ b/src/city.rs
@@ -179,7 +179,7 @@ impl<'a> HouseLayout<'a> {
     }
 }
 
-fn get_price(city: &City, houses: &Vec<House>) -> u32 {
+pub fn get_price(city: &City, houses: &Vec<House>) -> u32 {
     let mut price = 0u32;
     for &house in houses {
         price += city.get_price(house) as u32;
diff --git a/src/db.rs b/src/db.rs
index b496d09..b59bf37 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -12,6 +12,12 @@ pub struct SavedLayout {
     houses: Vec<House>
 }
 
+impl SavedLayout {
+    pub fn houses(&self) -> &Vec<House> {
+        &self.houses
+    }
+}
+
 impl LayoutDB {
     pub fn from_file(filename: &str) -> Result<Self> {
         let connection = Connection::open(filename)?;
@@ -55,6 +61,7 @@ impl LayoutDB {
                                 params![layout_id, house.x as u32, house.y as u32])?;
         }
         transaction.commit()?;
+        self.layouts.push(SavedLayout {id: layout_id as usize, houses: houses.clone()});
         Ok(())
     }
 }
diff --git a/src/main.rs b/src/main.rs
index 5dc7cf2..5cf7861 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,6 +5,7 @@ use std::fmt::Formatter;
 use std::collections::{HashMap, HashSet};
 use city::{HouseLayout, City, House};
 use crate::db::LayoutDB;
+use crate::population::build_house_probabilities;
 
 mod optimization;
 mod population;
@@ -16,6 +17,12 @@ fn main() {
     eprintln!("Loaded the DB, {} stored layouts", db.layouts().len());
 
     let city = City::read_from_file("01.in");
+    eprintln!("Loaded the city file, {} houses", city.get_house_count());
+
+    const MIN_WEIGHT_SCORE: f64 = 600000.;
+    const MAX_WEIGHT_SCORE: f64 = 700000.;
+    const DB_CHOICE_PROBABILITY: f64 = 0.9;
+
     let mut best_price: Option<u32> = None;
 
     loop {
@@ -23,8 +30,11 @@ fn main() {
         eprintln!("Starting seed {}", seed);
         let mut rng = StdRng::seed_from_u64(seed);
         let mut layout = HouseLayout::new(&city);
-        eprintln!("Starting random population...");
-        population::populate_random(&mut layout, &mut rng);
+        //eprintln!("Starting random population...");
+        //population::populate_random(&mut layout, &mut rng);
+        eprintln!("Starting random population, using DB ({}-{} {})...",
+                  MIN_WEIGHT_SCORE, MAX_WEIGHT_SCORE, DB_CHOICE_PROBABILITY);
+        population::populate_using_db(&mut layout, &mut rng, &db, MIN_WEIGHT_SCORE, MAX_WEIGHT_SCORE, DB_CHOICE_PROBABILITY);
         eprintln!("Finished random init, price: {}", layout.price());
         loop {
             let mut improved = false;
diff --git a/src/population.rs b/src/population.rs
index 06e61b2..5079660 100644
--- a/src/population.rs
+++ b/src/population.rs
@@ -1,8 +1,14 @@
 use rand::Rng;
-use crate::city::{SIZE, House, HouseLayout};
-use rand::prelude::StdRng;
+use crate::city::{SIZE, House, HouseLayout, City};
+use rand::prelude::*;
+use crate::db::LayoutDB;
+use crate::city;
+use std::collections::HashMap;
+use std::ops::Add;
+use itertools::Itertools;
+use rand::distributions::WeightedIndex;
 
-pub(crate) fn populate_random(layout: &mut HouseLayout, rng: &mut StdRng) {
+pub fn populate_random(layout: &mut HouseLayout, rng: &mut StdRng) {
     loop {
         loop {
             let x = rng.gen_range(0..SIZE);
@@ -19,3 +25,52 @@ pub(crate) fn populate_random(layout: &mut HouseLayout, rng: &mut StdRng) {
         }
     }
 }
+
+pub fn build_house_probabilities(city: &City, db: &LayoutDB, min_score: f64, max_score: f64) -> (WeightedIndex<f64>, Vec<House>){
+    let mut counts: HashMap<House, f64> = HashMap::new();
+    for layout in db.layouts() {
+        let price = city::get_price(&city, layout.houses()) as f64;
+        let value_range = max_score - min_score;
+        let value = 1. - ((price - min_score) / value_range).max(0.).min(1.);
+
+        for house in layout.houses() {
+            *counts.entry(*house).or_default() += value;
+        }
+    }
+
+    let houses: Vec<_> = counts.into_iter().filter(|(house, price)| *price != 0.).collect();
+    let index = WeightedIndex::new(houses.iter().map(|(house, price)| price)).unwrap();
+
+    (index, houses.iter().map(|(house, _)| *house).collect())
+}
+
+pub fn populate_using_db(layout: &mut HouseLayout, mut rng: &mut StdRng, db: &LayoutDB, min_score: f64, max_score: f64, db_probability: f64) {
+    let (mut index, houses) = build_house_probabilities(&layout.city, &db, min_score, max_score);
+
+    loop {
+        if rng.gen::<f64>() < db_probability {
+            // TODO: Test without allowing buying covered tiles
+            let house = houses[index.sample(&mut rng)];
+
+            layout.add_house(house);
+
+            if layout.is_valid() {
+                break;
+            }
+        } else {
+            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;
+            }
+        }
+    }
+}