// 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::output::*;
use rand::prelude::*;
use std::fmt;

#[derive(Copy, Clone)]
pub struct Prelayout {
    standard_columns: [[u8; 3]; 6],
    index_columns: [[u8; 6]; 2],
}

impl Prelayout {
    pub fn get_standard_column(&self, index: usize) -> &[u8; 3] {
        &self.standard_columns[index]
    }

    pub fn get_index_column(&self, index: usize) -> &[u8; 6] {
        &self.index_columns[index]
    }

    pub fn new_random_from<R: RngCore>(pl: &Prelayout, rng: &mut R) -> Prelayout {
        let mut standard_columns = pl.standard_columns.clone();
        let mut index_columns = pl.index_columns.clone();

        let mut count = rng.gen_range(1..=MAX_NUM_TRANSPOSITIONS);
        while count > 0 {
            let source_index: bool = rng.gen();
            let target_index: bool = rng.gen();

            let saved;
            let target_col: usize;
            let target_idx: usize;
            if target_index {
                target_col = rng.gen_range(0..2);
                target_idx = rng.gen_range(0..6);
                saved = index_columns[target_col][target_idx];
            } else {
                target_col = rng.gen_range(0..6);
                target_idx = rng.gen_range(0..3);
                saved = standard_columns[target_col][target_idx];
            }

            let mut source_col: usize;
            let source_idx: usize;
            if source_index {
                source_col = rng.gen_range(0..2);
                source_idx = rng.gen_range(0..6);
                if target_index {
                    while source_col != target_col {
                        source_col = rng.gen_range(0..2)
                    }
                    index_columns[target_col][target_idx] = index_columns[source_col][source_idx];
                } else {
                    standard_columns[target_col][target_idx] =
                        index_columns[source_col][source_idx];
                }
                index_columns[source_col][source_idx] = saved;
            } else {
                source_col = rng.gen_range(0..6);
                source_idx = rng.gen_range(0..3);
                if target_index {
                    index_columns[target_col][target_idx] =
                        standard_columns[source_col][source_idx];
                } else {
                    while source_col != target_col {
                        source_col = rng.gen_range(0..6);
                    }
                    standard_columns[target_col][target_idx] =
                        standard_columns[source_col][source_idx];
                }
                standard_columns[source_col][source_idx] = saved;
            }
            count -= 1;
        }
        Prelayout {
            standard_columns,
            index_columns,
        }
    }

    fn from_char_array(ca: &[char; NUM_KEYS]) -> Prelayout {
        let mut standard_columns = [[0; 3]; 6];
        let mut index_columns = [[0; 6]; 2];

        for i in 0..6 {
            let ind = if i < 3 { i } else { i + 4 };
            for j in 0..3 {
                standard_columns[i][j] = ca[ind + j * ROW_LENGTH] as u8;
            }
        }

        for j in 0..3 {
            index_columns[0][j] = ca[3 + j * ROW_LENGTH] as u8;
            index_columns[1][j] = ca[6 + j * ROW_LENGTH] as u8;
            index_columns[0][j + 3] = ca[4 + j * ROW_LENGTH] as u8;
            index_columns[1][j + 3] = ca[5 + j * ROW_LENGTH] as u8;
        }

        Prelayout {
            standard_columns,
            index_columns,
        }
    }
}

#[derive(Copy, Clone)]
pub struct Layout {
    keys: [char; NUM_KEYS],
}

impl Layout {
    pub fn as_prelayout(&self) -> Prelayout {
        Prelayout::from_char_array(&self.keys)
    }

    pub fn from_prelayout(pl: &Prelayout, corpus: &Corpus) -> Layout {
        fn my_f32_compare(r: f32, l: f32) -> std::cmp::Ordering {
            if r < l {
                std::cmp::Ordering::Less
            } else {
                std::cmp::Ordering::Greater
            }
        }
        fn weight_function<const N: usize, const M: usize>(
            columns: &[[u8; N]; M],
            corpus: &Corpus,
        ) -> Vec<(Vec<u8>, f32)> {
            let mut result = columns
                .iter()
                .map(|col| {
                    let mut weight = 0.0;
                    let mut wcol: Vec<(u8, f32)> = col
                        .iter()
                        .map(|&c| {
                            let w = corpus.get_character_perc(c);
                            weight += w;
                            (c, w)
                        })
                        .collect();
                    wcol.sort_by(|(_, l), (_, r)| my_f32_compare(*r, *l));
                    wcol.swap(0, 1);
                    if N == 6 {
                        wcol.swap(0, 2);
                        wcol.swap(3, 4);
                        wcol.swap(3, 5);
                    }
                    (wcol.into_iter().map(|(k, _)| k as u8).collect(), weight)
                })
                .collect::<Vec<(Vec<u8>, f32)>>();
            result.sort_by(|(_, l), (_, r)| my_f32_compare(*l, *r));
            result
        }

        // Derive column orderings
        let w_index_columns = weight_function(&pl.index_columns, corpus);
        let w_standard_columns = weight_function(&pl.standard_columns, corpus);

        let mut standard_column_ordering = [0; 6];
        let mut left_ind;
        let mut right_ind;

        match PRELAYOUT_TO_LAYOUT_ALGORITHM {
            Algorithm::MinimiseLSBThenBalance => {
                // first we still want left pinky to have lowest load
                let mut balance: f32 = -w_index_columns[0].1 + w_index_columns[1].1;
                standard_column_ordering[0] = 0;
                standard_column_ordering[5] = 1;

                // next find minimum lsb pairing
                let mut lsbs = [0.0; 2];
                for swap in [false, true] {
                    for first_pairing in [true, false] {
                        let mind = match (swap, first_pairing) {
                            (true, true) => 5,
                            (true, false) => 4,
                            (false, true) => 4,
                            (false, false) => 5,
                        };
                        let m1 = w_standard_columns[mind].0[0];
                        let m2 = w_standard_columns[mind].0[1];
                        let m3 = w_standard_columns[mind].0[2];
                        let iind = if first_pairing { 0 } else { 1 };
                        let i1 = w_index_columns[iind].0[3];
                        let i2 = w_index_columns[iind].0[4];
                        let i3 = w_index_columns[iind].0[5];
                        lsbs[if swap { 1 } else { 0 }] += 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);
                    }
                }

                // if lsbs[0] < lsbs[1] then (4, 0) and (5, 1) are paired, else (5, 0) and (4, 1)
                // record weights of pairings and whether we swapped 4 <-> 5 for min lsb
                let (swap, w1, w2) = if lsbs[0] < lsbs[1] {
                    (
                        false,
                        w_standard_columns[4].1 + w_index_columns[0].1,
                        w_standard_columns[5].1 + w_index_columns[1].1,
                    )
                } else {
                    (
                        true,
                        w_standard_columns[5].1 + w_index_columns[0].1,
                        w_standard_columns[4].1 + w_index_columns[1].1,
                    )
                };

                // now we try to balance hands according to combined weight, place
                // on left and check if we need to swap
                let w: f32;
                if w1 < w2 {
                    (
                        (standard_column_ordering[2], left_ind),
                        (standard_column_ordering[3], right_ind),
                    ) = if swap {
                        ((5, 0), (4, 1))
                    } else {
                        ((4, 0), (5, 1))
                    };
                    w = -w1 + w2;
                } else {
                    (
                        (standard_column_ordering[2], left_ind),
                        (standard_column_ordering[3], right_ind),
                    ) = if swap {
                        ((4, 1), (5, 0))
                    } else {
                        ((5, 1), (4, 0))
                    };
                    w = -w2 + w1;
                }
                if (balance + w).abs() > (balance - w).abs() {
                    // swapping would be better
                    (left_ind, right_ind) = (right_ind, left_ind);
                    (standard_column_ordering[2], standard_column_ordering[3]) =
                        (standard_column_ordering[3], standard_column_ordering[2]);
                    balance -= w;
                } else {
                    balance += w;
                }

                // finally we do ring finger
                (standard_column_ordering[1], standard_column_ordering[4]) =
                    if balance >= 0.0 { (3, 2) } else { (2, 3) };
            }
            Algorithm::GreedyBalance => {
                // greedy on weight, not provably correct in *all* cases but good
                // enough in practice
                let mut balance: f32 = 0.0;
                let mut left_col: usize = 0;
                let mut right_col: usize = 9;

                for (src, (_, weight)) in w_standard_columns.iter().enumerate() {
                    let left: bool = ((balance >= 0.0) && (left_col <= 2)) || (right_col <= 6);
                    let ind = if left { left_col } else { right_col };
                    standard_column_ordering[ind] = src;
                    if left {
                        left_col += 1;
                    } else {
                        right_col -= 1;
                    }
                    balance += if left { -weight } else { *weight };
                }
                (left_ind, right_ind) = if balance >= 0.0 { (0, 1) } else { (1, 0) }
            }
        }

        // write keys
        let mut keys = ['x'; NUM_KEYS];
        // standard
        for (ind, src) in standard_column_ordering.into_iter().enumerate() {
            let target = if ind < 3 { ind } else { 4 + ind };
            for i in 0..3 {
                keys[target + i * ROW_LENGTH] = w_standard_columns[src].0[i] as char;
            }
        }
        // index
        for j in 0..3 {
            keys[3 + j * ROW_LENGTH] = w_index_columns[left_ind].0[j] as char;
            keys[6 + j * ROW_LENGTH] = w_index_columns[right_ind].0[j] as char;
            keys[4 + j * ROW_LENGTH] = w_index_columns[left_ind].0[j + 3] as char;
            keys[5 + j * ROW_LENGTH] = w_index_columns[right_ind].0[j + 3] as char;
        }

        Layout { keys }
    }

    pub fn get_keys(&self) -> &[char] {
        &self.keys
    }

    fn char_array_to_layout(keys: [char; NUM_KEYS]) -> Layout {
        return Layout { keys };
    }

    pub fn get_key(&self, index: usize) -> u8 {
        self.keys[index] as u8
    }

    pub fn from_verbose(inp: &str) -> Option<Layout> {
        let mut layout: [char; NUM_KEYS] = ['x'; NUM_KEYS];

        let mut k: usize = 0;
        for c in inp.chars() {
            if let Some(c) = canonicalise(c) {
                layout[k] = c as char;
                k += 1;
            }
            if k > NUM_KEYS {
                return None;
            }
        }

        return Some(Layout::char_array_to_layout(layout));
    }
}

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