use std::fmt;
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,
        });
    }
}

fn dump_bigrams(corpus: &Corpus) -> Vec<(String, f32)> {
    let mut result: Vec<(String, f32)> = Vec::new();

    for &x in KEY_CHARS.iter() {
        for &y in KEY_CHARS.iter() {
            let mut pair = String::from(x);
            pair.push(y);
            let perc = corpus.bigram_percs[pair_to_index(x as u8, y as u8)];
            if perc > 0.0 {
                result.push((pair, perc));
            }
        }
    }

    result.sort_by(|(_, c1), (_, c2)| {
        if c2 < c1 {
            std::cmp::Ordering::Less
        } else {
            std::cmp::Ordering::Greater
        }
    });

    return result;
}

impl fmt::Display for Corpus {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str(&format!(
            "Corpus has top 5 bigrams {:?}",
            &dump_bigrams(self)[0..5]
        ))
    }
}