Penalty for index under-usage

1 files changed, 19 insertions, 8 deletions

diff --git a/src/corpus.rs b/src/corpus.rs
index 2677cd2..9b522e0 100644
--- a/src/corpus.rs
+++ b/src/corpus.rs
@@ -8,8 +8,9 @@ const NUM_BIGRAMS: usize = NUM_KEYS * NUM_KEYS;
 
 pub struct Corpus {
     bigram_count: [u32; NUM_BIGRAMS],
-    character_count: [u32; NUM_KEYS],
+    character_count: [u32; 256],
     total_count: u32,
+    index_threshold: u32,
 }
 
 fn pair_to_index(x: char, y: char) -> usize {
@@ -23,7 +24,7 @@ impl Corpus {
         let contents = fs::read_to_string(path)?;
 
         let mut bigram_count = [0; NUM_BIGRAMS];
-        let mut character_count = [0; NUM_KEYS];
+        let mut character_count = [0; 256];
         let mut total_count = 0;
 
         let mut last_char = None;
@@ -37,22 +38,25 @@ impl Corpus {
                     }
                 }
                 last_char = Some(c);
-                character_count[CHAR_TO_INDEX[c as usize]] += 1;
+                character_count[c as usize] += 1;
                 total_count += 1;
             } else {
                 last_char = None;
             }
         }
 
+        let index_threshold = (total_count as f32 * 0.15) as u32;
+
         return Ok(Corpus {
             bigram_count,
             character_count,
             total_count,
+            index_threshold,
         });
     }
 
     pub fn prelayout_fitness(&self, layout: &Prelayout) -> u32 {
-        let mut sfb: u32 = 0;
+        let mut score: u32 = 0;
 
         for i in 0..6 {
             let keys = layout.get_standard_column(i);
@@ -60,7 +64,7 @@ impl Corpus {
             let k2 = keys[1];
             let k3 = keys[2];
 
-            sfb += self.bigram_count[pair_to_index(k1, k2)]
+            score += self.bigram_count[pair_to_index(k1, k2)]
                 + self.bigram_count[pair_to_index(k1, k3)]
                 + self.bigram_count[pair_to_index(k2, k3)];
         }
@@ -73,7 +77,12 @@ impl Corpus {
             let k4 = keys[3];
             let k5 = keys[4];
             let k6 = keys[5];
-            sfb += self.bigram_count[pair_to_index(k1, k2)]
+
+            // We want index usage!
+            let index_count: u32 = keys.iter().map(|&k| self.character_count[k as usize]).sum();
+            score += (self.index_threshold as i64 - index_count as i64).abs() as u32 / 32;
+
+            score += self.bigram_count[pair_to_index(k1, k2)]
                 + self.bigram_count[pair_to_index(k1, k3)]
                 + self.bigram_count[pair_to_index(k2, k3)]
                 + self.bigram_count[pair_to_index(k4, k5)]
@@ -90,7 +99,7 @@ impl Corpus {
                 + self.bigram_count[pair_to_index(k3, k6)];
         }
 
-        sfb
+        score
     }
 
     pub fn evaluate_layout(&self, layout: &Layout) -> Evaluation {
@@ -98,7 +107,9 @@ impl Corpus {
         let mut sfb: [u32; 8] = [0; 8];
 
         for (i, &c) in self.character_count.iter().enumerate() {
-            keypress[layout.translate_index(i)] = c;
+            if c > 0 {
+                keypress[layout.get_index((i as u8) as char)] = c;
+            }
         }
 
         // TODO: We make assumptions about NUM_KEYS here