// Copyright © 2024 Mikhail Hogrefe
//
// This file is part of Malachite.
//
// Malachite is free software: you can redistribute it and/or modify it under the terms of the GNU
// Lesser General Public License (LGPL) as published by the Free Software Foundation; either version
// 3 of the License, or (at your option) any later version. See <https://www.gnu.org/licenses/>.

use itertools::Itertools;
use malachite_base::num::basic::unsigneds::PrimitiveUnsigned;
use malachite_base::num::random::geometric::geometric_random_unsigneds;
use malachite_base::num::random::striped::striped_random_unsigned_vecs_from_length_iterator;
use malachite_base::random::{Seed, EXAMPLE_SEED};
use malachite_base::strings::ToBinaryString;
use malachite_base::test_util::stats::common_values_map::common_values_map_debug;
use malachite_base::test_util::stats::median;
use malachite_base::vecs::random_values_from_vec;
use std::iter::repeat;

fn striped_random_unsigned_vecs_from_length_iterator_helper<
    T: PrimitiveUnsigned,
    I: Clone + Iterator<Item = u64>,
>(
    lengths_gen: &dyn Fn(Seed) -> I,
    mean_stripe_numerator: u64,
    mean_stripe_denominator: u64,
    expected_values: &[&[&str]],
    expected_common_values: &[(&[&str], usize)],
    expected_median: (&[&str], Option<&[&str]>),
) {
    let xss = striped_random_unsigned_vecs_from_length_iterator::<T, _>(
        EXAMPLE_SEED,
        lengths_gen,
        mean_stripe_numerator,
        mean_stripe_denominator,
    );
    let values = xss
        .clone()
        .take(20)
        .map(|xs| {
            xs.into_iter()
                .map(|x: T| x.to_binary_string())
                .collect_vec()
        })
        .collect_vec();
    let common_values = common_values_map_debug(1000000, 10, xss.clone())
        .into_iter()
        .map(|(xs, freq)| {
            (
                xs.into_iter()
                    .map(|x: T| x.to_binary_string())
                    .collect_vec(),
                freq,
            )
        })
        .collect_vec();
    let (median_lo, median_hi) = median(xss.take(1000000));
    let median_lo = median_lo
        .into_iter()
        .map(|x: T| x.to_binary_string())
        .collect_vec();
    let median_hi = median_hi.map(|xs| {
        xs.into_iter()
            .map(|x: T| x.to_binary_string())
            .collect_vec()
    });

    let values = values
        .iter()
        .map(|xs| xs.iter().map(String::as_str).collect_vec())
        .collect_vec();
    let common_values = common_values
        .iter()
        .map(|(xs, freq)| (xs.iter().map(String::as_str).collect_vec(), *freq))
        .collect_vec();
    let median_lo = median_lo.iter().map(String::as_str).collect_vec();
    let median_hi = median_hi
        .as_ref()
        .map(|xs| xs.iter().map(String::as_str).collect_vec());
    assert_eq!(
        (
            values.iter().map(Vec::as_slice).collect_vec().as_slice(),
            common_values
                .iter()
                .map(|(xs, f)| (xs.as_slice(), *f))
                .collect_vec()
                .as_slice(),
            (median_lo.as_slice(), median_hi.as_deref())
        ),
        (expected_values, expected_common_values, expected_median)
    );
}

#[test]
fn test_striped_random_unsigned_vecs_from_length_iterator() {
    striped_random_unsigned_vecs_from_length_iterator_helper::<u8, _>(
        &|seed| random_values_from_vec(seed, vec![0, 2, 4]),
        2,
        1,
        &[
            &["11001100", "11000"],
            &["110010", "1111010", "1111101", "1100001"],
            &["1001100", "1111"],
            &["10110000", "11011", "1111", "101000"],
            &["11100110", "1111011", "10001101", "111110"],
            &["100111", "11101011"],
            &[],
            &["1100010", "101011"],
            &[],
            &["10100111", "100000", "1110011", "11111111"],
            &["1011010", "100010", "1100100", "1110110"],
            &["10000111", "11101111"],
            &["1100011", "11100101", "10001010", "11011011"],
            &["1001000", "11011101"],
            &["10100000", "101110", "11101110", "111101"],
            &["1100010", "1001111", "1100001", "100"],
            &["100111", "11100010", "110011", "11011110"],
            &[],
            &["10100", "11001011", "11100000", "1001101"],
            &[],
        ],
        &[
            (&[], 333820),
            (&["11110010", "1100111"], 17),
            (&["10010101", "101"], 16),
            (&["11010100", "1100001"], 16),
            (&["1111101", "11100000"], 16),
            (&["11111111", "1111000"], 16),
            (&["101100", "110101"], 15),
            (&["1010111", "1010000"], 15),
            (&["10101", "1111110"], 15),
            (&["11110111", "10000"], 15),
        ],
        (
            &["111111", "11110001", "11111010", "10100110"],
            Some(&["111111", "11110010"]),
        ),
    );
    striped_random_unsigned_vecs_from_length_iterator_helper::<u8, _>(
        &|seed| random_values_from_vec(seed, vec![0, 2, 4]),
        10,
        1,
        &[
            &["0", "0"],
            &["1110000", "0", "11111100", "11"],
            &["11111110", "1111"],
            &["0", "0", "0", "11111000"],
            &["0", "0", "1111110", "0"],
            &["11011111", "11111111"],
            &[],
            &["11110000", "11111111"],
            &[],
            &["11111111", "11000011", "11111", "0"],
            &["0", "10000000", "11111001", "1111111"],
            &["11111111", "11111111"],
            &["11111111", "11111111", "11111111", "11111111"],
            &["0", "10000"],
            &["0", "0", "0", "11111100"],
            &["0", "0", "0", "0"],
            &["11111111", "0", "11110001", "11111111"],
            &[],
            &["10111100", "11", "10110100", "11101111"],
            &[],
        ],
        &[
            (&[], 333820),
            (&["11111111", "11111111"], 34254),
            (&["0", "0"], 34215),
            (&["11111111", "11111111", "11111111", "11111111"], 6413),
            (&["0", "0", "0", "0"], 6336),
            (&["11111111", "11111"], 3951),
            (&["111111", "0"], 3891),
            (&["11111110", "11111111"], 3890),
            (&["11111111", "0"], 3889),
            (&["11111111", "1111"], 3878),
        ],
        (&["1", "0", "11111110", "11111111"], None),
    );
    striped_random_unsigned_vecs_from_length_iterator_helper::<u8, _>(
        &|seed| random_values_from_vec(seed, vec![0, 2, 4]),
        11,
        10,
        &[
            &["1011010", "11010101"],
            &["1101010", "10010101", "1010010", "10101010"],
            &["10101010", "10101010"],
            &["11010100", "10101010", "101010", "1101011"],
            &["10101010", "10101010", "10101010", "1001010"],
            &["10110101", "1010010"],
            &[],
            &["1001010", "1010101"],
            &[],
            &["1010101", "10110101", "10101010", "1010110"],
            &["1010100", "1010101", "1010101", "10010101"],
            &["1001001", "1010101"],
            &["1010101", "100101", "1010101", "1010101"],
            &["10110100", "10101010"],
            &["10101010", "10101010", "10111100", "11100"],
            &["10101010", "10001010", "1010110", "10100101"],
            &["10101011", "1101010", "1010101", "10110101"],
            &[],
            &["10101010", "1010100", "10110101", "10101010"],
            &[],
        ],
        &[
            (&[], 333820),
            (&["1010101", "1010101"], 39945),
            (&["10101010", "10101010"], 39772),
            (&["1010101", "1010101", "1010101", "1010101"], 8708),
            (&["10101010", "10101010", "10101010", "10101010"], 8700),
            (&["1010010", "1010101"], 4100),
            (&["1010101", "10110101"], 4090),
            (&["1010101", "11010101"], 4089),
            (&["1010101", "10101101"], 4080),
            (&["1010110", "1010101"], 4060),
        ],
        (&["1010101", "1010101", "1010101", "10110101"], None),
    );
    striped_random_unsigned_vecs_from_length_iterator_helper::<u8, _>(
        &|seed| geometric_random_unsigneds::<u64>(seed, 2, 1).map(|x| x << 1),
        2,
        1,
        &[
            &[
                "11001100", "11000", "11001", "10111101", "10111110", "110000", "11010011", "11",
                "1110110", "11100011", "1", "1100101",
            ],
            &["1111100", "10101111"],
            &[
                "1011100", "11110000", "1100100", "11011101", "1001001", "111101", "101",
                "10011001", "11111011", "10010111", "1110110", "1101111", "100110", "1110",
                "11011111", "1100011",
            ],
            &["11001010", "10101"],
            &[
                "1101110", "1101111", "1000101", "10100000", "101110", "11101110", "111101",
                "10110001", "10100111", "110000", "10", "1110110", "111", "1110011", "10001000",
                "1100010", "11001", "10111100", "10101001", "11001000", "110", "1101", "1100101",
                "1111011", "11011", "10101101", "11111111", "10011110",
            ],
            &[],
            &[
                "11110001", "1100011", "11001000", "1111111", "11001100", "11001101", "11100011",
                "11011000", "110101", "10000111",
            ],
            &[
                "11011100", "10111010", "10000101", "0", "10110010", "10010010", "10111000",
                "10010100",
            ],
            &["11010001", "1111000"],
            &[],
            &[
                "1010100", "10101110", "1000001", "1001011", "1011100", "10110100", "11001000",
                "1101", "10010101", "10111011", "1111000", "1011",
            ],
            &["1010011", "111001", "10101111", "1100101"],
            &[],
            &[],
            &["10010001", "11010001"],
            &["10111110", "11101111", "11101011", "10010", "1101101", "1010110"],
            &[],
            &["100110", "10001010"],
            &[],
            &["1101100", "11000"],
        ],
        &[
            (&[], 333981),
            (&["1", "10011100"], 13),
            (&["11111011", "110100"], 13),
            (&["110000", "10110111"], 13),
            (&["1100110", "10011000"], 13),
            (&["11101100", "11011100"], 13),
            (&["1", "100101"], 12),
            (&["100010", "101100"], 12),
            (&["1010", "11000011"], 12),
            (&["1100101", "1001101"], 12),
        ],
        (&["111111", "10101111"], None),
    );
    striped_random_unsigned_vecs_from_length_iterator_helper::<u8, _>(
        &|seed| geometric_random_unsigneds::<u64>(seed, 2, 1).map(|x| x << 1),
        10,
        1,
        &[
            &[
                "0", "0", "111000", "0", "11111110", "10000001", "11111111", "11", "0", "0", "0",
                "11111111",
            ],
            &["0", "11000000"],
            &[
                "11100000", "11111111", "11111011", "11111111", "0", "0", "11100000", "1",
                "11111111", "11111111", "11111111", "11001", "0", "11111110", "11111111",
                "11111111",
            ],
            &["0", "0"],
            &[
                "0", "0", "100000", "0", "0", "0", "11111100", "11111111", "11111111", "11111111",
                "11111111", "111111", "1000000", "11111100", "1111111", "10001000", "1111111",
                "1001", "10", "0", "0", "0", "11111111", "11100001", "10111111", "1", "11111000",
                "1111111",
            ],
            &[],
            &[
                "11111111", "11011111", "11111111", "11111111", "11", "11000000", "11111111",
                "11111111", "11110000", "11111111",
            ],
            &["11111110", "0", "0", "11100000", "11111111", "1111", "0", "11111111"],
            &["11110011", "11111111"],
            &[],
            &[
                "10001000", "11111111", "11", "11111110", "0", "11111100", "11111111", "11111111",
                "11111111", "11", "11111111", "11111",
            ],
            &["11111111", "1100001", "0", "11100000"],
            &[],
            &[],
            &["11111111", "11111111"],
            &["11111000", "1111111", "0", "10000000", "111", "1111"],
            &[],
            &["0", "0"],
            &[],
            &["0", "0"],
        ],
        &[
            (&[], 333981),
            (&["0", "0"], 22998),
            (&["11111111", "11111111"], 22759),
            (&["11111111", "11111111", "11111111", "11111111"], 2851),
            (&["0", "0", "0", "0"], 2788),
            (&["11111111", "111111"], 2626),
            (&["1111", "0"], 2620),
            (&["111", "0"], 2615),
            (&["111111", "0"], 2608),
            (&["11111000", "11111111"], 2608),
        ],
        (
            &["1", "0", "11111000", "1", "1111100", "0", "0", "100000"],
            Some(&["1", "0", "11111000", "1", "10000000", "11111111", "11111111", "11111111"]),
        ),
    );
    striped_random_unsigned_vecs_from_length_iterator_helper::<u8, _>(
        &|seed| geometric_random_unsigneds::<u64>(seed, 2, 1).map(|x| x << 1),
        11,
        10,
        &[
            &[
                "1011010", "11010101", "1001010", "10110101", "11010110", "10101010", "10101010",
                "1101010", "10100101", "10101010", "10011010", "1010010",
            ],
            &["10101010", "10101010"],
            &[
                "10101010", "1010010", "10101001", "1010101", "10101011", "10101010", "10101010",
                "1010101", "10110101", "1010010", "1010101", "1010101", "1010101", "10010010",
                "10101010", "10101010",
            ],
            &["1001010", "10101010"],
            &[
                "10101010", "10010110", "10101010", "10101010", "10101010", "10111100", "11100",
                "1010101", "1000101", "10101011", "11010010", "10101010", "1011010", "1010101",
                "10101101", "1101010", "1010101", "10101001", "10101010", "10101010", "10101010",
                "10101010", "100101", "10100101", "1010101", "1010101", "1010101", "10010101",
            ],
            &[],
            &[
                "1010101", "1010101", "10101101", "1010110", "1101010", "11010101", "10100110",
                "10101010", "10101010", "11010100",
            ],
            &[
                "10101010", "10110110", "10101010", "101010", "1010101", "1010101", "10101010",
                "10101010",
            ],
            &["10010101", "10101010"],
            &[],
            &[
                "10101010", "1001010", "1010101", "1011101", "10010101", "10101010", "10101010",
                "10101010", "10101010", "10110101", "10010100", "1010101",
            ],
            &["1010101", "1010101", "1010101", "1010101"],
            &[],
            &[],
            &["1010101", "1001010"],
            &["1001010", "10101101", "11010010", "10101010", "1101010", "1010101"],
            &[],
            &["10000100", "10101010"],
            &[],
            &["10101010", "101010"],
        ],
        &[
            (&[], 333981),
            (&["10101010", "10101010"], 26752),
            (&["1010101", "1010101"], 26461),
            (&["10101010", "10101010", "10101010", "10101010"], 3857),
            (&["1010101", "1010101", "1010101", "1010101"], 3852),
            (&["10101010", "1010010"], 2802),
            (&["1010110", "1010101"], 2776),
            (&["1010101", "10100101"], 2768),
            (&["1010101", "10010101"], 2739),
            (&["10101011", "10101010"], 2730),
        ],
        (&["1010101", "1010101", "10010101", "10101010"], None),
    );
}

#[test]
#[should_panic]
fn striped_random_unsigned_vecs_from_length_iterator_fail_1() {
    striped_random_unsigned_vecs_from_length_iterator::<u8, _>(EXAMPLE_SEED, &|_| repeat(1), 1, 0);
}

#[test]
#[should_panic]
fn striped_random_unsigned_vecs_from_length_iterator_fail_2() {
    striped_random_unsigned_vecs_from_length_iterator::<u8, _>(EXAMPLE_SEED, &|_| repeat(1), 2, 3);
}