// Pruebas de selección. Comprueban lo que afirman los apuntes.
#include <algorithm>
#include <cassert>
#include <cstdio>
#include <random>
#include <vector>

#include "select.h"

static long long count_at_most(const std::vector<int>& a, int pivot) {
    return std::count_if(a.begin(), a.end(), [&](int x) { return x <= pivot; });
}

static long long count_at_least(const std::vector<int>& a, int pivot) {
    return std::count_if(a.begin(), a.end(), [&](int x) { return x >= pivot; });
}

// El lema: al menos 3n/10 - 3 elementos a cada lado, y por lo tanto a lo sumo 7n/10 + 3 en cada parte.
static void check_pivot_bound(const std::vector<int>& a) {
    long long n = static_cast<long long>(a.size());
    long long comparisons = 0;
    int pivot = median_of_medians(a, comparisons);
    assert(10 * count_at_most(a, pivot) >= 3 * n - 30);
    assert(10 * count_at_least(a, pivot) >= 3 * n - 30);
    Parts parts = partition3(a, pivot, comparisons);
    assert(10 * static_cast<long long>(parts.less.size()) <= 7 * n + 30);
    assert(10 * static_cast<long long>(parts.greater.size()) <= 7 * n + 30);
    assert(!parts.equal.empty());
}

static void test_agree_with_sort() {
    std::mt19937 rng(1);
    for (int round = 0; round < 200; round++) {
        int n = 1 + static_cast<int>(rng() % 60);
        int range = round % 4 == 0 ? 3 : 1000;
        std::vector<int> a(n);
        for (int& x : a) {
            x = static_cast<int>(rng() % range);
        }
        std::vector<int> sorted = a;
        std::sort(sorted.begin(), sorted.end());
        for (int k = 0; k < n; k++) {
            long long c1 = 0;
            long long c2 = 0;
            long long c3 = 0;
            assert(quickselect_first(a, k, c1) == sorted[k]);
            assert(quickselect_random(a, k, rng, c2) == sorted[k]);
            assert(select_linear(a, k, c3) == sorted[k]);
            std::vector<int> copy = a;
            std::nth_element(copy.begin(), copy.begin() + k, copy.end());
            assert(copy[k] == sorted[k]);
        }
    }
}

static void test_bound() {
    std::mt19937 rng(2);
    for (int n = 1; n <= 400; n++) {
        std::vector<int> random_values(n);
        std::vector<int> few(n);
        std::vector<int> sorted(n);
        for (int i = 0; i < n; i++) {
            random_values[i] = static_cast<int>(rng() % 100000);
            few[i] = static_cast<int>(rng() % 3);
            sorted[i] = i;
        }
        std::vector<int> reversed(sorted.rbegin(), sorted.rend());
        check_pivot_bound(random_values);
        check_pivot_bound(few);
        check_pivot_bound(sorted);
        check_pivot_bound(reversed);
        check_pivot_bound(std::vector<int>(n, 5));
    }
}

// La entrada construida queda en exactamente 3·techo(m/2) del lado de los <=.
static void test_tight_input() {
    for (int n = 5; n <= 1000; n += 5) {
        std::vector<int> a = tight_input(n);
        long long comparisons = 0;
        int pivot = median_of_medians(a, comparisons);
        int groups = n / 5;
        assert(count_at_most(a, pivot) == 3 * ((groups + 1) / 2));
        check_pivot_bound(a);
    }
}

// Ordenado y k = n - 1: las partes miden n, n - 1, ..., 2, y luego queda un solo elemento.
static void test_worst_case_first_pivot() {
    for (long long n = 2; n <= 3000; n += 37) {
        std::vector<int> sorted(n);
        for (long long i = 0; i < n; i++) {
            sorted[i] = static_cast<int>(i);
        }
        long long count = 0;
        assert(quickselect_first(sorted, n - 1, count) == n - 1);
        assert(count == n * (n + 1) / 2 - 1);

        // Con k = 0 el primer pivote ya es la respuesta: una sola pasada.
        count = 0;
        assert(quickselect_first(sorted, 0, count) == 0);
        assert(count == n);

        // Al revés, k = 0 es el caso malo.
        std::vector<int> reversed(sorted.rbegin(), sorted.rend());
        count = 0;
        assert(quickselect_first(reversed, 0, count) == 0);
        assert(count == n * (n + 1) / 2 - 1);
    }
}

// Las comparaciones crecen como n. Cada vuelta paga a lo sumo 3n: 10 por grupo de 5 y 1 por elemento con el
// pivote. La demostración da entonces c = 3 y T(n) <= 20·c·n = 60n, si el rango pequeño no pide más.
static void test_linear_comparisons() {
    std::mt19937 rng(3);
    for (int n : {1000, 10000, 100000}) {
        std::vector<int> a(n);
        for (int& x : a) {
            x = static_cast<int>(rng());
        }
        long long count = 0;
        select_linear(a, a.size() / 2, count);
        assert(count <= 60LL * n);
    }
}

int main() {
    test_agree_with_sort();
    test_bound();
    test_tight_input();
    test_worst_case_first_pivot();
    test_linear_comparisons();
    std::puts("test_select: ok");
}
