// 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 crate::config::*;
use crate::corpus::*;
use crate::layout::*;
use crate::output::*;

use std::fmt;

pub fn prelayout_fitness(corpus: &Corpus, layout: &Prelayout) -> f32 {
    let mut score: f32 = 0.0;
    let mut min1: f32 = std::f32::INFINITY;
    let mut min2: f32 = std::f32::INFINITY;

    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 += corpus.get_bigram_perc(k1, k2)
            + corpus.get_bigram_perc(k1, k3)
            + corpus.get_bigram_perc(k2, k3);

        let col_usage = corpus.get_character_perc(k1)
            + corpus.get_character_perc(k2)
            + corpus.get_character_perc(k3);

        if col_usage < min2 {
            if col_usage < min1 {
                min2 = min1;
                min1 = col_usage;
            } else {
                min2 = col_usage;
            }
        }

        score += (corpus.get_skipgram_perc(k1, k2)
            + corpus.get_skipgram_perc(k1, k3)
            + corpus.get_skipgram_perc(k2, k3))
            * STANDARD_DSFB_WEIGHT;
    }

    score += pinky_col_penalty(min1, min2);

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

        let index_percs = keys.iter().map(|&k| corpus.get_character_perc(k)).sum();
        score += index_usage_fitness(index_percs);

        score += corpus.get_bigram_perc(k1, k2)
            + corpus.get_bigram_perc(k1, k3)
            + corpus.get_bigram_perc(k2, k3)
            + corpus.get_bigram_perc(k4, k5)
            + corpus.get_bigram_perc(k4, k6)
            + corpus.get_bigram_perc(k5, k6)
            + corpus.get_bigram_perc(k1, k4)
            + corpus.get_bigram_perc(k1, k5)
            + corpus.get_bigram_perc(k1, k6)
            + corpus.get_bigram_perc(k2, k4)
            + corpus.get_bigram_perc(k2, k5)
            + corpus.get_bigram_perc(k2, k6)
            + corpus.get_bigram_perc(k3, k4)
            + corpus.get_bigram_perc(k3, k5)
            + corpus.get_bigram_perc(k3, k6);

        score += (corpus.get_skipgram_perc(k1, k2)
            + corpus.get_skipgram_perc(k1, k3)
            + corpus.get_skipgram_perc(k2, k3)
            + corpus.get_skipgram_perc(k4, k5)
            + corpus.get_skipgram_perc(k4, k6)
            + corpus.get_skipgram_perc(k5, k6)
            + corpus.get_skipgram_perc(k1, k4)
            + corpus.get_skipgram_perc(k1, k5)
            + corpus.get_skipgram_perc(k1, k6)
            + corpus.get_skipgram_perc(k2, k4)
            + corpus.get_skipgram_perc(k2, k5)
            + corpus.get_skipgram_perc(k2, k6)
            + corpus.get_skipgram_perc(k3, k4)
            + corpus.get_skipgram_perc(k3, k5)
            + corpus.get_skipgram_perc(k3, k6))
            * INDEX_DSFB_WEIGHT;
    }

    score
}

pub struct Evaluation {
    keypress_perc: [f32; NUM_KEYS],
    sfb: [f32; 8],
    dsfb: [f32; 8],
    lsb: f32,
}

impl Evaluation {
    pub fn evaluate_layout(corpus: &Corpus, layout: &Layout) -> Evaluation {
        let mut keypress_perc: [f32; NUM_KEYS] = [0.0; NUM_KEYS];
        let mut sfb: [f32; 8] = [0.0; 8];
        let mut dsfb: [f32; 8] = [0.0; 8];
        let mut lsb = 0.0;

        for (i, &c) in layout.get_keys().iter().enumerate() {
            keypress_perc[i] = corpus.get_character_perc(c as u8);
        }

        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] = corpus.get_bigram_perc(k1, k2)
                + corpus.get_bigram_perc(k1, k3)
                + corpus.get_bigram_perc(k2, k3);

            dsfb[i] = corpus.get_skipgram_perc(k1, k2)
                + corpus.get_skipgram_perc(k1, k3)
                + corpus.get_skipgram_perc(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] += corpus.get_bigram_perc(k4, k5)
                    + corpus.get_bigram_perc(k4, k6)
                    + corpus.get_bigram_perc(k5, k6)
                    + corpus.get_bigram_perc(k1, k4)
                    + corpus.get_bigram_perc(k1, k5)
                    + corpus.get_bigram_perc(k1, k6)
                    + corpus.get_bigram_perc(k2, k4)
                    + corpus.get_bigram_perc(k2, k5)
                    + corpus.get_bigram_perc(k2, k6)
                    + corpus.get_bigram_perc(k3, k4)
                    + corpus.get_bigram_perc(k3, k5)
                    + corpus.get_bigram_perc(k3, k6);
                dsfb[i] += corpus.get_skipgram_perc(k4, k5)
                    + corpus.get_skipgram_perc(k4, k6)
                    + corpus.get_skipgram_perc(k5, k6)
                    + corpus.get_skipgram_perc(k1, k4)
                    + corpus.get_skipgram_perc(k1, k5)
                    + corpus.get_skipgram_perc(k1, k6)
                    + corpus.get_skipgram_perc(k2, k4)
                    + corpus.get_skipgram_perc(k2, k5)
                    + corpus.get_skipgram_perc(k2, k6)
                    + corpus.get_skipgram_perc(k3, k4)
                    + corpus.get_skipgram_perc(k3, k5)
                    + corpus.get_skipgram_perc(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] += corpus.get_bigram_perc(k4, k5)
                    + corpus.get_bigram_perc(k5, k6)
                    + corpus.get_bigram_perc(k4, k6)
                    + corpus.get_bigram_perc(k1, k4)
                    + corpus.get_bigram_perc(k1, k5)
                    + corpus.get_bigram_perc(k1, k6)
                    + corpus.get_bigram_perc(k2, k4)
                    + corpus.get_bigram_perc(k2, k5)
                    + corpus.get_bigram_perc(k2, k6)
                    + corpus.get_bigram_perc(k3, k4)
                    + corpus.get_bigram_perc(k3, k5)
                    + corpus.get_bigram_perc(k3, k6);
                dsfb[i] += corpus.get_skipgram_perc(k4, k5)
                    + corpus.get_skipgram_perc(k4, k6)
                    + corpus.get_skipgram_perc(k5, k6)
                    + corpus.get_skipgram_perc(k1, k4)
                    + corpus.get_skipgram_perc(k1, k5)
                    + corpus.get_skipgram_perc(k1, k6)
                    + corpus.get_skipgram_perc(k2, k4)
                    + corpus.get_skipgram_perc(k2, k5)
                    + corpus.get_skipgram_perc(k2, k6)
                    + corpus.get_skipgram_perc(k3, k4)
                    + corpus.get_skipgram_perc(k3, k5)
                    + corpus.get_skipgram_perc(k3, k6);
            }
        }

        // LSB for index-middle only
        for left in [true, false] {
            let mind = if left { 2 } else { 7 };
            let iind = if left { 4 } else { 5 };
            let m1 = layout.get_key(mind + 10 * 0);
            let m2 = layout.get_key(mind + 10 * 1);
            let m3 = layout.get_key(mind + 10 * 2);
            let i1 = layout.get_key(iind + 10 * 0);
            let i2 = layout.get_key(iind + 10 * 1);
            let i3 = layout.get_key(iind + 10 * 2);
            lsb += corpus.get_bigram_perc(m1, i1)
                + corpus.get_bigram_perc(m1, i2)
                + corpus.get_bigram_perc(m1, i3)
                + corpus.get_bigram_perc(m2, i1)
                + corpus.get_bigram_perc(m2, i2)
                + corpus.get_bigram_perc(m2, i3)
                + corpus.get_bigram_perc(m3, i1)
                + corpus.get_bigram_perc(m3, i2)
                + corpus.get_bigram_perc(m3, i3);
        }

        return Evaluation {
            keypress_perc,
            sfb,
            dsfb,
            lsb,
        };
    }

    fn output_eval(&self) -> String {
        let mut result = String::new();

        result += &"Percent per key:\n";
        result += &format_block_output(self.keypress_perc.into_iter().map(|u| u * 100.0));

        let mut finger_percs = [0.0; 8];
        for (i, &c) in self.keypress_perc.iter().enumerate() {
            finger_percs[KEY_TO_FINGER[i]] += c;
        }

        result += "\nFinger usage: ";
        let mut lh: f32 = 0.0;
        let mut rh: f32 = 0.0;
        for (i, &u) in finger_percs.iter().enumerate() {
            result += &format!("{:>5.2}%{}", u * 100.0, if i < 7 { ", " } else { "" });
            if i % 10 < 4 {
                lh += u;
            } else {
                rh += u;
            }
        }

        result += &format!("\nHand usage: {:.2}% vs {:.2}%", lh * 100.0, rh * 100.0);

        result += "\nSFB distribution : ";
        for (i, &u) in self.sfb.iter().enumerate() {
            result += &format!("{:>6.3}%{}", u * 100.0, if i < 7 { ", " } else { "" });
        }

        result += "\nDSFB distribution: ";
        for (i, &u) in self.dsfb.iter().enumerate() {
            result += &format!("{:>6.3}%{}", u * 100.0, if i < 7 { ", " } else { "" });
        }

        let sfb = self.sfb.iter().sum::<f32>();
        let dsfb = self.dsfb.iter().sum::<f32>();
        result += &format!("\nTotal sfb: {:.3}%", sfb * 100.0);
        result += &format!("\nTotal dsfb: {:.3}%", dsfb * 100.0);
        result += &format!("\nTotal lsb: {:.3}%", self.lsb * 100.0);

        return result;
    }
}

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