1 files changed, 63 insertions, 12 deletions

diff --git a/srchr/src/corpus.rs b/srchr/src/corpus.rs
index 306018c..102aabd 100644
--- a/srchr/src/corpus.rs
+++ b/srchr/src/corpus.rs
@@ -2,6 +2,7 @@ use std::fmt;
 use std::fs;
 
 use crate::config::*;
+extern crate json;
 
 const fn build_lookup_table() -> [usize; 128] {
     let mut result = [0; 128];
@@ -20,8 +21,6 @@ pub struct Corpus {
     bigram_percs: [f32; NUM_BIGRAMS],
     skipgram_percs: [f32; NUM_BIGRAMS],
     character_percs: [f32; 128],
-    total_count: u32,
-    total_bigrams: u32,
 }
 
 fn pair_to_index(x: u8, y: u8) -> usize {
@@ -43,11 +42,63 @@ impl Corpus {
         self.skipgram_percs[pair_to_index(x, y)]
     }
 
-    pub fn get_total_bigrams(&self) -> u32 {
-        self.total_bigrams
+    pub fn load_from_json_file(path: &str) -> Result<Corpus, std::io::Error> {
+        let contents = fs::read_to_string(path)?;
+
+        let json = json::parse(&contents).unwrap();
+
+        let mut character_percs = [0.0; 128];
+        let mut bigram_percs = [0.0; NUM_BIGRAMS];
+        let mut skipgram_percs = [0.0; NUM_BIGRAMS];
+
+        for (key, value) in json.entries() {
+            if key == "characters" {
+                for (c, p) in value.entries() {
+                    if let Some(c) = canonicalise(c.as_bytes()[0] as char) {
+                        character_percs[c as usize] = p.as_f32().unwrap();
+                    }
+                }
+            } else if key == "bigrams" {
+                for (b, p) in value.entries() {
+                    let canon: Vec<Option<u8>> = b
+                        .as_bytes()
+                        .iter()
+                        .map(|&b| canonicalise(b as char))
+                        .collect();
+                    if let Some(c1) = canon[0] {
+                        if let Some(c2) = canon[1] {
+                            bigram_percs[pair_to_index(c1, c2)] += p.as_f32().unwrap();
+                            bigram_percs[pair_to_index(c2, c1)] += p.as_f32().unwrap();
+                        }
+                    }
+                }
+            } else if key == "skipgrams" {
+                for (t, p) in value.entries() {
+                    let canon: Vec<Option<u8>> = t
+                        .as_bytes()
+                        .iter()
+                        .map(|&b| canonicalise(b as char))
+                        .collect();
+                    if let Some(c1) = canon[0] {
+                        if let Some(c3) = canon[1] {
+                            skipgram_percs[pair_to_index(c1, c3)] += p.as_f32().unwrap();
+                            skipgram_percs[pair_to_index(c3, c1)] += p.as_f32().unwrap();
+                        }
+                    }
+                }
+            }
+        }
+
+        println!("Chars: {:?}\nBigs:{:?}", character_percs, bigram_percs);
+
+        return Ok(Corpus {
+            skipgram_percs,
+            bigram_percs,
+            character_percs,
+        });
     }
 
-    pub fn load(path: &str) -> Result<Corpus, std::io::Error> {
+    pub fn load_from_text_file(path: &str) -> Result<Corpus, std::io::Error> {
         let contents = fs::read_to_string(path)?;
 
         let mut bigram_count = [0; NUM_BIGRAMS];
@@ -100,8 +151,6 @@ impl Corpus {
             skipgram_percs,
             bigram_percs,
             character_percs,
-            total_count,
-            total_bigrams,
         });
     }
 }
@@ -109,11 +158,14 @@ impl Corpus {
 fn dump_bigrams(corpus: &Corpus) -> Vec<(String, f32)> {
     let mut result: Vec<(String, f32)> = Vec::new();
 
-    for (i, &x) in KEY_CHARS.iter().enumerate() {
-        for &y in &KEY_CHARS[i..] {
+    for &x in KEY_CHARS.iter() {
+        for &y in KEY_CHARS.iter() {
             let mut pair = String::from(x);
             pair.push(y);
-            result.push((pair, corpus.bigram_percs[pair_to_index(x as u8, y as u8)]));
+            let perc = corpus.bigram_percs[pair_to_index(x as u8, y as u8)];
+            if perc > 0.0 {
+                result.push((pair, perc));
+            }
         }
     }
 
@@ -131,8 +183,7 @@ fn dump_bigrams(corpus: &Corpus) -> Vec<(String, f32)> {
 impl fmt::Display for Corpus {
     fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
         formatter.write_str(&format!(
-            "Counted {} characters, top 5 bigrams {:?}",
-            self.total_count,
+            "Corpus has top 5 bigrams {:?}",
             &dump_bigrams(self)[0..5]
         ))
     }