// Copyright (C) 2026 J. Nathanael Philipp (jnphilipp) <jnathanael@philipp.land>
//
// Bikkuri, calculate the surprisal of words in texts.
//
// This file is part of bikkuri.
//
// bikkuri 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.
//
// bikkuri 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 bikkuri. If not, see <http://www.gnu.org/licenses/>
//
// SPDX-License-Identifier: GPL-3.0-or-later

extern crate bikkuri;

use bikkuri::ngram::NGramSurprisal;

fn get_texts() -> Vec<Vec<&'static str>> {
    vec![
        vec![
            "lorem",
            "ipsum",
            "dolor",
            "sit",
            "amet",
            "consectetur",
            "adipiscing",
            "elit",
            "sed",
            "do",
            "eiusmod",
            "tempor",
            "incididunt",
            "ut",
            "labore",
            "et",
            "dolore",
            "magna",
            "aliqua",
            "lectus",
            "vitae",
            "mi",
            "mauris",
            "facilisis",
            "luctus",
            "posuere",
            "elementum",
            "cursus",
            "fermentum",
            "non",
            "massa",
            "litora",
            "ultricies",
            "donec",
            "consectetur",
            "lectus",
            "sodales",
            "elementum",
            "facilisis",
            "fermentum",
            "senectus",
            "et",
            "mi",
            "lobortis",
            "iaculis",
            "fermentum",
            "ante",
            "consequat",
            "aptent",
            "quam",
            "ultricies",
            "sed",
            "libero",
        ],
        vec![
            "convallis",
            "fringilla",
            "dignissim",
            "massa",
            "arcu",
            "posuere",
            "in",
            "imperdiet",
            "eu",
            "at",
            "aenean",
            "sit",
            "sollicitudin",
            "arcu",
            "consectetur",
            "bibendum",
            "a",
            "class",
            "consequat",
            "nulla",
            "fusce",
            "donec",
            "fringilla",
            "lobortis",
            "vehicula",
            "tristique",
            "hac",
            "vestibulum",
            "aenean",
            "fusce",
            "hac",
            "eros",
            "pulvinar",
            "cubilia",
            "ultricies",
            "bibendum",
            "tempor",
            "at",
            "gravida",
            "malesuada",
        ],
        vec![
            "quam",
            "nullam",
            "eget",
            "est",
            "ac",
            "fusce",
            "elit",
            "eros",
            "tellus",
            "facilisis",
            "malesuada",
            "hendrerit",
            "lectus",
            "dictumst",
            "praesent",
            "sodales",
            "augue",
            "netus",
            "nullam",
            "eget",
            "justo",
            "duis",
            "nulla",
            "eros",
            "potenti",
            "at",
        ],
        vec![
            "sagittis",
            "non",
            "nullam",
            "maximus",
            "purus",
            "nisi",
            "ante",
            "sollicitudin",
            "etiam",
            "urna",
            "volutpat",
            "mauris",
            "semper",
            "sollicitudin",
            "felis",
            "aliquet",
            "nullam",
            "praesent",
            "vel",
            "curabitur",
            "habitasse",
            "condimentum",
            "metus",
            "eu",
            "pretium",
            "maximus",
            "tristique",
            "litora",
            "nisl",
            "quam",
            "purus",
            "facilisis",
            "laoreet",
            "dignissim",
        ],
        vec![
            "dui",
            "proin",
            "adipiscing",
            "consequat",
            "felis",
            "amet",
            "laoreet",
            "blandit",
            "laoreet",
            "porttitor",
            "litora",
            "in",
            "orci",
            "quis",
            "himenaeos",
            "nam",
            "adipiscing",
            "nisi",
            "metus",
            "pulvinar",
            "ante",
            "dignissim",
            "dolor",
            "dapibus",
            "a",
            "vivamus",
            "nisi",
        ],
        vec![
            "sit",
            "dapibus",
            "neque",
            "sapien",
            "cubilia",
            "porttitor",
            "egestas",
            "proin",
            "nunc",
            "purus",
            "nisi",
            "dapibus",
            "nunc",
            "venenatis",
            "cras",
            "mauris",
            "suscipit",
            "pellentesque",
            "netus",
            "porta",
            "vitae",
            "fermentum",
            "turpis",
            "arcu",
            "odio",
            "duis",
            "erat",
            "proin",
            "tellus",
            "lectus",
            "etiam",
            "vitae",
            "pulvinar",
            "magna",
            "turpis",
            "velit",
            "curabitur",
            "mauris",
            "suscipit",
            "turpis",
            "congue",
            "dui",
            "suscipit",
            "neque",
            "quam",
            "amet",
            "vestibulum",
            "quam",
            "class",
            "suspendisse",
            "torquent",
            "risus",
            "nisi",
            "laoreet",
            "tellus",
            "lectus",
            "interdum",
            "habitant",
            "molestie",
            "commodo",
            "tristique",
            "scelerisque",
        ],
        vec![
            "malesuada",
            "curabitur",
            "erat",
            "mauris",
            "tempor",
            "vel",
            "leo",
            "condimentum",
            "class",
            "pharetra",
            "platea",
            "justo",
            "rutrum",
            "egestas",
            "magna",
            "interdum",
            "per",
            "dolor",
            "aliquet",
            "vitae",
            "hac",
            "eleifend",
            "torquent",
            "maximus",
            "felis",
            "nullam",
            "ante",
            "elit",
            "etiam",
            "nisl",
            "mattis",
            "consectetur",
            "ipsum",
            "nisl",
            "semper",
            "accumsan",
        ],
        vec![
            "diam",
            "pulvinar",
            "conubia",
            "purus",
            "libero",
            "consequat",
            "turpis",
            "vulputate",
            "hac",
            "donec",
            "porttitor",
            "eleifend",
            "rhoncus",
            "volutpat",
            "leo",
            "sapien",
            "condimentum",
            "conubia",
            "rhoncus",
            "torquent",
            "commodo",
            "non",
        ],
        vec![
            "morbi", "etiam", "porta", "ornare", "posuere", "integer", "etiam", "purus",
            "pulvinar", "nullam", "justo", "amet", "arcu", "dictum", "vitae", "quisque", "etiam",
        ],
        vec![
            "feugiat",
            "massa",
            "maximus",
            "ultrices",
            "ipsum",
            "nec",
            "dui",
            "enim",
            "orci",
            "convallis",
            "est",
            "rhoncus",
            "nisi",
            "congue",
            "felis",
            "nisi",
            "euismod",
        ],
    ]
}

#[test]
fn test_unigram_surprisal() {
    let mut unigram_surprisal = NGramSurprisal::new(1);
    assert_eq!(
        None,
        unigram_surprisal.surprisal(&[[
            "lorem",
            "ipsum",
            "dolor",
            "sit",
            "amet",
            "consectetur",
            "adipiscing",
        ]])
    );
    let texts = get_texts();
    unigram_surprisal.fit(&texts);
    assert_eq!(
        Some(vec![
            vec![
                ("lorem".to_string(), 8.383704292474052),
                ("ipsum".to_string(), 6.798741791752896),
                ("dolor".to_string(), 6.798741791752896),
                ("sit".to_string(), 6.798741791752896),
                ("amet".to_string(), 6.383704292474052),
                ("consectetur".to_string(), 6.383704292474052),
                ("adipiscing".to_string(), 6.798741791752896),
            ],
            vec![
                ("feugiat".to_string(), 8.383704292474052),
                ("massa".to_string(), 6.798741791752896),
                ("maximus".to_string(), 6.383704292474052),
                ("ultrices".to_string(), 8.383704292474052),
                ("ipsum".to_string(), 6.798741791752896),
                ("nec".to_string(), 8.383704292474052),
            ],
        ]),
        unigram_surprisal.surprisal(&vec![
            vec![
                "lorem",
                "ipsum",
                "dolor",
                "sit",
                "amet",
                "consectetur",
                "adipiscing",
            ],
            vec!["feugiat", "massa", "maximus", "ultrices", "ipsum", "nec",],
        ])
    );
    assert_eq!(
        Some(vec![vec![
            ("lorem".to_string(), 8.383704292474052),
            ("ipsum".to_string(), 6.798741791752896),
            ("dolor".to_string(), 6.798741791752896),
            ("sit".to_string(), 6.798741791752896),
            ("foo".to_string(), 8.388017285345136),
            ("bar".to_string(), 8.388017285345136),
        ]]),
        unigram_surprisal.surprisal(&[["lorem", "ipsum", "dolor", "sit", "foo", "bar"]])
    );
}

#[test]
fn test_bigram_surprisal() {
    let mut bigram_surprisal = NGramSurprisal::new(2);
    assert_eq!(
        None,
        bigram_surprisal.surprisal(&[[
            "lorem",
            "ipsum",
            "dolor",
            "sit",
            "amet",
            "consectetur",
            "adipiscing",
        ]])
    );
    let texts = get_texts();
    bigram_surprisal.fit(&texts);
    assert_eq!(
        Some(vec![
            vec![
                ("lorem".to_string(), 3.321928094887362),
                ("ipsum".to_string(), 0.0),
                ("dolor".to_string(), 1.5849625007211563),
                ("sit".to_string(), 1.5849625007211563),
                ("amet".to_string(), 1.5849625007211563),
                ("consectetur".to_string(), 2.0),
                ("adipiscing".to_string(), 2.0),
            ],
            vec![
                ("feugiat".to_string(), 3.321928094887362),
                ("massa".to_string(), 0.0),
                ("maximus".to_string(), 1.5849625007211563),
                ("ultrices".to_string(), 2.0),
                ("ipsum".to_string(), 0.0),
                ("nec".to_string(), 1.5849625007211563),
            ],
        ]),
        bigram_surprisal.surprisal(&vec![
            vec![
                "lorem",
                "ipsum",
                "dolor",
                "sit",
                "amet",
                "consectetur",
                "adipiscing",
            ],
            vec!["feugiat", "massa", "maximus", "ultrices", "ipsum", "nec",],
        ])
    );
    assert_eq!(
        Some(vec![vec![
            ("lorem".to_string(), 3.321928094887362),
            ("ipsum".to_string(), -0.0),
            ("dolor".to_string(), 1.5849625007211563),
            ("sit".to_string(), 1.5849625007211563),
            ("amet".to_string(), 1.5849625007211563),
            ("foo".to_string(), 2.321928094887362),
            ("bar".to_string(), 8.388017285345136),
        ]]),
        bigram_surprisal.surprisal(&[["lorem", "ipsum", "dolor", "sit", "amet", "foo", "bar"]])
    );
}
