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

use core::cmp::Ordering;
use rayon::prelude::*;

use rand::prelude::*;
use rand_pcg::*;

const NUM_KEYS: usize = 30;

#[derive(Copy, Clone)]
struct Keyboard {
    mappings: [u8; NUM_KEYS],
}

#[derive(Copy, Clone)]
struct KeyboardEvaluation {
    keypress: [u32; NUM_KEYS],
    total_keypress: u32,
    sfb: [u32; 8],
}

impl Keyboard {
    const KEY_TO_FINGER: [usize; NUM_KEYS] = [
        0, 1, 2, 3, 3, 4, 4, 5, 6, 7, 0, 1, 2, 3, 3, 4, 4, 5, 6, 7, 0, 1, 2, 3, 3, 4, 4, 5, 6, 7,
    ];

    fn to_lookup(&self) -> [Option<usize>; 256] {
        let mut result = [None; 256];
        for (i, m) in self.mappings.into_iter().enumerate() {
            let index = Some(i);
            if m.is_ascii_alphabetic() {
                result[m.to_ascii_lowercase() as usize] = index;
                result[m.to_ascii_uppercase() as usize] = index;
            } else {
                // For now we just hardcode this
                match m {
                    b'.' => {
                        result[b'>' as usize] = index;
                    }
                    b',' => {
                        result[b'<' as usize] = index;
                    }
                    b'/' => {
                        result[b'?' as usize] = index;
                    }
                    b'\'' => {
                        result[b'"' as usize] = index;
                    }
                    _ => (),
                }
                result[m as usize] = index;
            }
        }
        return result;
    }

    fn from_verbose(inp: &str) -> Option<Keyboard> {
        let mut mappings: [u8; NUM_KEYS] = [0; NUM_KEYS];
        let mut k: usize = 0;
        for c in inp.chars() {
            let valid = (c.is_uppercase() && c.is_alphabetic())
                || c == '.'
                || c == '/'
                || c == ','
                || c == '\'';
            if valid {
                mappings[k] = c as u8;
                k += 1;
            }
            if k > NUM_KEYS {
                return None;
            }
        }
        return Some(Keyboard { mappings });
    }

    fn new_random_from<R: RngCore>(kbd: &Keyboard, rng: &mut R) -> Keyboard {
        let mut mappings = kbd.mappings;
        mappings.shuffle(rng);
        Keyboard { mappings }
    }
}

impl KeyboardEvaluation {
    fn new() -> KeyboardEvaluation {
        KeyboardEvaluation {
            keypress: [0; NUM_KEYS],
            sfb: [0; 8],
            total_keypress: 0,
        }
    }

    fn evaluate_keyboard(kbd: &Keyboard, corpus: &str) -> KeyboardEvaluation {
        let mut result = KeyboardEvaluation::new();
        let lookup_key = kbd.to_lookup();

        let mut last_finger: Option<usize> = None;
        for c in corpus.chars() {
            let key = lookup_key[c as usize];
            if let Some(key_idx) = key {
                result.total_keypress += 1;
                let finger = Keyboard::KEY_TO_FINGER[key_idx];
                result.keypress[key_idx] += 1;
                if let Some(last_finger) = last_finger {
                    if last_finger == finger {
                        result.sfb[last_finger] += 1;
                    }
                }
                last_finger = Some(finger);
            }
        }
        return result;
    }

    fn output_eval(&self) -> String {
        let mut result = String::new();
        let tot = self.total_keypress as f32;

        result += &format!(
            "Total keypresses: {}\nPercent per key:\n",
            self.total_keypress
        );
        result += &format_block_output(self.keypress.into_iter().map(|k| 100.0 * k as f32 / tot));

        let mut finger_usages = [0; 8];
        for (i, &c) in self.keypress.iter().enumerate() {
            finger_usages[Keyboard::KEY_TO_FINGER[i]] += c;
        }

        result += "finger usage: ";
        for (i, &u) in finger_usages.iter().enumerate() {
            result += &format!(
                "{:>5.2}%{}",
                u as f32 / tot * 100.0,
                if i < 7 { ", " } else { "" }
            );
        }

        result += "\nsame finger bigrams: ";
        for (i, &u) in self.sfb.iter().enumerate() {
            result += &format!(
                "{:>5.2}%{}",
                u as f32 / finger_usages[i] as f32 * 50.0,
                if i < 7 { ", " } else { "" }
            );
        }

        result += &format!("\ntotal sfb: {}", self.sfb.iter().sum::<u32>());

        return result;
    }

    // TODO
    fn fitness(&self) -> f32 {
        let mut finger_usages = [0; 8];
        // let tot = self.total_keypress as f32;
        for (i, &c) in self.keypress.iter().enumerate() {
            finger_usages[Keyboard::KEY_TO_FINGER[i]] += c;
        }
        // const MAX_FINGER_USAGES: [f32; 8] = [8.0, 11.0, 21.0, 21.0, 21.0, 21.0, 11.0, 8.0];
        // const MAX_FINGER_USAGES: [f32; 8] =
        //     [100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0];
        // if finger_usages
        //     .iter()
        //     .enumerate()
        //     .any(|(i, &c)| c as f32 / tot * 100.00 > MAX_FINGER_USAGES[i])
        // {
        //     std::u32::MAX
        // } else {
        self.sfb
            .iter()
            .enumerate()
            .map(|(i, &s)| s as f32 / finger_usages[i] as f32 * 50.0)
            .sum()
        // }
    }
}

// This is a silly amount of work to genericise the below...
trait ToMyString {
    const BLANK_STRING: &'static str;
    fn to_my_string(&self) -> String;
}

impl ToMyString for u8 {
    const BLANK_STRING: &'static str = " ";
    fn to_my_string(&self) -> String {
        (*self as char).to_string()
    }
}

impl ToMyString for u32 {
    const BLANK_STRING: &'static str = "     ";
    fn to_my_string(&self) -> String {
        format!("{:5}", *self)
    }
}

impl ToMyString for f32 {
    const BLANK_STRING: &'static str = "    ";
    fn to_my_string(&self) -> String {
        format!("{:>4.1}", *self)
    }
}

fn format_block_output<T: Iterator<Item = S>, S: ToMyString>(things: T) -> String {
    let mut result = String::new();
    for (i, t) in things.enumerate() {
        result += &t.to_my_string();
        if i < NUM_KEYS - 1 {
            result.push(' ');
        }
        if (i + 1) % 10 == 0 {
            result.push('\n');
            if NUM_KEYS < 30 && i == 19 {
                result += S::BLANK_STRING;
                result.push(' ')
            }
        } else if (i < 20 && (i + 1) % 5 == 0) || (i == (NUM_KEYS - 20) / 2 + 19) {
            result.push(' ');
        }
        if i + 1 >= NUM_KEYS {
            break;
        }
    }
    return result;
}

impl fmt::Display for Keyboard {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str(&format_block_output(self.mappings.into_iter()))
    }
}

impl fmt::Display for KeyboardEvaluation {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str(&self.output_eval())
    }
}

fn main() {
    let starting_kbd = Keyboard::from_verbose(
        "
        Y W F L M  K P O , Q
        U R S N H  D T E A I
        Z X C V J  B G ' . /
        ",
    )
    .unwrap();

    // let contents = fs::read_to_string("books.txt").unwrap();
    let contents = fs::read_to_string("chained_english_bigrams_10k.txt").unwrap();

    let mut rng = Pcg64::from_entropy();

    type Fun = Vec<(Keyboard, KeyboardEvaluation, f32)>;

    let mut best = std::f32::INFINITY;
    let mut kbds: Fun = (0..10)
        .into_iter()
        .map(|i| {
            let kbd;
            if i > 0 {
                kbd = Keyboard::new_random_from(&starting_kbd, &mut rng);
            } else {
                kbd = starting_kbd;
            }

            let evl = KeyboardEvaluation::evaluate_keyboard(&kbd, &contents);
            let fitness = evl.fitness();
            return (kbd, evl, fitness);
        })
        .collect();

    let mut counter = 0;
    while counter < 1000000000 {
        let new_layouts = kbds
            .iter()
            .enumerate()
            .map(|(i, (kbd, _, _))| {
                if i > 0 {
                    Keyboard::new_random_from(kbd, &mut rng)
                } else {
                    kbd.clone()
                }
            })
            .collect::<Vec<Keyboard>>();

        new_layouts
            .par_iter()
            .map(|kbd| {
                let evl = KeyboardEvaluation::evaluate_keyboard(kbd, &contents);
                let fitness = evl.fitness();
                return (*kbd, evl, fitness);
            })
            .collect_into_vec(&mut kbds);

        kbds.sort_by(|(_, _, lfit), (_, _, rfit)| {
            if lfit < rfit {
                Ordering::Less
            } else {
                Ordering::Greater
            }
        });

        if kbds[0].2 < best {
            best = kbds[0].2;
            println!("");
            println!("================================================================================\nLayout:\n{}",
                         kbds[0].0
                    );

            println!("{}\n{}", kbds[0].1, kbds[0].2);
        }

        counter += 1;
    }
}