use std::fmt;
use std::fs;

use crate::layout::*;

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_count: [u32; NUM_BIGRAMS],
    character_count: [u32; 256],
    total_count: u32,
    index_threshold: u32,
}

fn pair_to_index(x: char, y: char) -> 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_count(&self, c: char) -> u32 {
        self.character_count[c as usize]
    }

    pub fn load(path: &str) -> Result<Corpus, std::io::Error> {
        let contents = fs::read_to_string(path)?;

        let mut bigram_count = [0; NUM_BIGRAMS];
        let mut character_count = [0; 256];
        let mut total_count = 0;

        let mut last_char = None;
        for c in contents.chars() {
            if let Some(c) = canonicalise(c) {
                if let Some(lc) = last_char {
                    // We don't count these anyway
                    if lc != c {
                        bigram_count[pair_to_index(c, lc)] += 1;
                        bigram_count[pair_to_index(lc, c)] += 1;
                    }
                }
                last_char = Some(c);
                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 score: u32 = 0;

        for i in 0..6 {
            let keys = layout.get_standard_column(i);
            let k1 = keys[0];
            let k2 = keys[1];
            let k3 = keys[2];

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

        for i in 0..2 {
            let keys = layout.get_index_column(i);
            let k1 = keys[0];
            let k2 = keys[1];
            let k3 = keys[2];
            let k4 = keys[3];
            let k5 = keys[4];
            let k6 = keys[5];

            // 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 / 64;

            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)]
                + self.bigram_count[pair_to_index(k4, k6)]
                + self.bigram_count[pair_to_index(k5, k6)]
                + self.bigram_count[pair_to_index(k1, k4)]
                + self.bigram_count[pair_to_index(k1, k5)]
                + self.bigram_count[pair_to_index(k1, k6)]
                + self.bigram_count[pair_to_index(k2, k4)]
                + self.bigram_count[pair_to_index(k2, k5)]
                + self.bigram_count[pair_to_index(k2, k6)]
                + self.bigram_count[pair_to_index(k3, k4)]
                + self.bigram_count[pair_to_index(k3, k5)]
                + self.bigram_count[pair_to_index(k3, k6)];
        }

        score
    }

    pub fn evaluate_layout(&self, layout: &Layout) -> Evaluation {
        let mut keypress: [u32; NUM_KEYS] = [0; NUM_KEYS];
        let mut sfb: [u32; 8] = [0; 8];

        for (i, &c) in self.character_count.iter().enumerate() {
            if c > 0 {
                keypress[layout.get_index((i as u8) as char)] = c;
            }
        }

        // TODO: We make assumptions about NUM_KEYS here
        for i in 0..8 {
            let ind = if i < 4 { i } else { i + 2 };

            let k1 = layout.get_key(ind + 10 * 0);
            let k2 = layout.get_key(ind + 10 * 1);
            let k3 = layout.get_key(ind + 10 * 2);

            sfb[i] = self.bigram_count[pair_to_index(k1, k2)]
                + self.bigram_count[pair_to_index(k1, k3)]
                + self.bigram_count[pair_to_index(k2, k3)];

            if i == 3 {
                let k4 = layout.get_key(4 + 10 * 0);
                let k5 = layout.get_key(4 + 10 * 1);
                let k6 = layout.get_key(4 + 10 * 2);
                sfb[i] += self.bigram_count[pair_to_index(k4, k5)]
                    + self.bigram_count[pair_to_index(k4, k6)]
                    + self.bigram_count[pair_to_index(k5, k6)]
                    + self.bigram_count[pair_to_index(k1, k4)]
                    + self.bigram_count[pair_to_index(k1, k5)]
                    + self.bigram_count[pair_to_index(k1, k6)]
                    + self.bigram_count[pair_to_index(k2, k4)]
                    + self.bigram_count[pair_to_index(k2, k5)]
                    + self.bigram_count[pair_to_index(k2, k6)]
                    + self.bigram_count[pair_to_index(k3, k4)]
                    + self.bigram_count[pair_to_index(k3, k5)]
                    + self.bigram_count[pair_to_index(k3, k6)];
            } else if i == 4 {
                let k4 = layout.get_key(5 + 10 * 0);
                let k5 = layout.get_key(5 + 10 * 1);
                let k6 = layout.get_key(5 + 10 * 2);
                sfb[i] += self.bigram_count[pair_to_index(k4, k5)]
                    + self.bigram_count[pair_to_index(k5, k6)]
                    + self.bigram_count[pair_to_index(k4, k6)]
                    + self.bigram_count[pair_to_index(k1, k4)]
                    + self.bigram_count[pair_to_index(k1, k5)]
                    + self.bigram_count[pair_to_index(k1, k6)]
                    + self.bigram_count[pair_to_index(k2, k4)]
                    + self.bigram_count[pair_to_index(k2, k5)]
                    + self.bigram_count[pair_to_index(k2, k6)]
                    + self.bigram_count[pair_to_index(k3, k4)]
                    + self.bigram_count[pair_to_index(k3, k5)]
                    + self.bigram_count[pair_to_index(k3, k6)];
            }
        }

        return Evaluation::new(keypress, self.total_count, sfb);
    }
}

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

    for (i, &x) in KEY_CHARS.iter().enumerate() {
        for &y in &KEY_CHARS[i..] {
            let mut pair = String::from(x);
            pair.push(y);
            result.push((pair, corpus.bigram_count[pair_to_index(x, y)]));
        }
    }

    result.sort_by(|(_, c1), (_, c2)| c1.cmp(c2).reverse());

    return result;
}

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,
            &dump_bigrams(self)[0..5]
        ))
    }
}