// Copyright (C) 2022 tslil clingman
//
// This file is part of srchr.
//
// srchr is free software: you can redistribute it and/or modify it under the
// terms of the GNU General Public License as published by the Free Software
// Foundation, either version 3 of the License, or (at your option) any later
// version.
//
// srchr is distributed in the hope that it will be useful, but WITHOUT ANY
// WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
// A PARTICULAR PURPOSE. See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along with
// srchr. If not, see <https://www.gnu.org/licenses/>.

use std::fs;

use crate::config::*;
extern crate json;

const fn build_lookup_table() -> [usize; 128] {
    let mut result = [0; 128];
    let mut i = 0;

    while i < NUM_KEYS {
        result[KEY_CHARS[i] as usize] = i;
        i += 1;
    }
    return result;
}
const CHAR_TO_INDEX: [usize; 128] = build_lookup_table();

const NUM_BIGRAMS: usize = NUM_KEYS * NUM_KEYS;
pub struct Corpus {
    bigram_percs: [f32; NUM_BIGRAMS],
    skipgram_percs: [f32; NUM_BIGRAMS],
    character_percs: [f32; 128],
}

fn pair_to_index(x: u8, y: u8) -> usize {
    let xi = CHAR_TO_INDEX[x as usize];
    let yi = CHAR_TO_INDEX[y as usize];
    xi + NUM_KEYS * yi
}

impl Corpus {
    pub fn get_character_perc(&self, c: u8) -> f32 {
        self.character_percs[c as usize]
    }

    pub fn get_bigram_perc(&self, x: u8, y: u8) -> f32 {
        self.bigram_percs[pair_to_index(x, y)]
    }

    pub fn get_skipgram_perc(&self, x: u8, y: u8) -> f32 {
        self.skipgram_percs[pair_to_index(x, y)]
    }

    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();
                        }
                    }
                }
            }
        }

        return Ok(Corpus {
            skipgram_percs,
            bigram_percs,
            character_percs,
        });
    }

    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];
        let mut skipgram_count = [0; NUM_BIGRAMS];
        let mut character_count = [0; 128];
        let mut total_count = 0;
        let mut total_bigrams = 0;
        let mut total_trigrams = 0;

        let mut last_char: Option<u8> = None;
        let mut last_last_char: Option<u8> = None;
        for c in contents.chars() {
            let current_char = canonicalise(c);
            if let Some(c) = current_char {
                if let Some(lc) = last_char {
                    if let Some(llc) = last_last_char {
                        if llc != lc && lc != c && llc != c {
                            skipgram_count[pair_to_index(llc, c)] += 1;
                            skipgram_count[pair_to_index(c, llc)] += 1;
                        }
                        total_trigrams += 1;
                    }
                    if lc != c {
                        bigram_count[pair_to_index(c, lc)] += 1;
                        bigram_count[pair_to_index(lc, c)] += 1;
                    }
                    total_bigrams += 1;
                }
                character_count[c as usize] += 1;
                total_count += 1;
            }
            last_last_char = last_char;
            last_char = current_char;
        }

        let mut bigram_percs = [0.0; NUM_BIGRAMS];
        let mut skipgram_percs = [0.0; NUM_BIGRAMS];
        let mut character_percs = [0.0; 128];

        for i in 0..NUM_BIGRAMS {
            bigram_percs[i] = bigram_count[i] as f32 / total_bigrams as f32;
            skipgram_percs[i] = skipgram_count[i] as f32 / total_trigrams as f32;
        }

        for i in 0..128 {
            character_percs[i] = character_count[i] as f32 / total_count as f32;
        }

        return Ok(Corpus {
            skipgram_percs,
            bigram_percs,
            character_percs,
        });
    }
}