#include <iostream>
#include <vector>
#include <string>
#include <chrono>
#include <iomanip>
#include <cmath>
#include <algorithm>

#include "vector.h"

using namespace std;
using namespace std::chrono;

static constexpr size_t N = 100'000'000;
static constexpr int SHORT_RUNS = 3;
static constexpr int LONG_RUNS = 10;

constexpr const char* COLOR_RESET = "\033[0m";
constexpr const char* COLOR_RED   = "\033[31m";
constexpr const char* COLOR_GREEN = "\033[32m";

bool longMode = false;
double total_speedup = 0;
int speedup_count = 0;

template<typename F>
double run_bench(F&& fn) {
    auto start = steady_clock::now();
    fn();
    auto end = steady_clock::now();
    return duration_cast<duration<double, milli>>(end - start).count();
}

struct Stats {
    double mean{}, min{}, max{}, stddev{};
};

Stats compute_stats(const std::vector<double>& samples) {
    Stats s{};
    if (samples.empty()) return s;
    s.min = *min_element(samples.begin(), samples.end());
    s.max = *max_element(samples.begin(), samples.end());
    double sum = 0;
    for (double v : samples) sum += v;
    s.mean = sum / samples.size();
    for (double v : samples) s.stddev += (v - s.mean) * (v - s.mean);
    s.stddev = sqrt(s.stddev / samples.size());
    return s;
}

template<typename F>
Stats run_all(F&& fn, int runs) {
    std::vector<double> results;
    for (int i = 0; i < runs; ++i)
        results.push_back(run_bench(fn));
    return compute_stats(results);
}

void print_test(const std::string& name, Stats std_stats, Stats fast_stats) {
    double speedup = std_stats.mean / fast_stats.mean;
    bool passed = fast_stats.mean < std_stats.mean;
    const char* color = passed ? COLOR_GREEN : COLOR_RED;
    const char* status = passed ? "[PASS]" : "[FAIL]";

    total_speedup += speedup;
    ++speedup_count;

    cout << color << status << COLOR_RESET
         << " " << left << setw(27) << name
         << " | std: " << right << setw(7) << fixed << setprecision(2) << std_stats.mean << "ms"
         << " ±" << setw(5) << std_stats.stddev
         << " | fast: " << setw(7) << fast_stats.mean << "ms"
         << " ±" << setw(5) << fast_stats.stddev
         << " | x" << color << setw(6) << setprecision(2) << speedup << COLOR_RESET
         << '\n';
}

void run_test(const string& name, auto std_fn, auto fast_fn) {
    int R = longMode ? LONG_RUNS : SHORT_RUNS;
    auto std_stats  = run_all(std_fn, R);
    auto fast_stats = run_all(fast_fn, R);
    print_test(name, std_stats, fast_stats);
}

struct S {
    int x, y;
    S(int a, int b): x(a), y(b) {}
    S(const S& o): x(o.x), y(o.y) {}
    S(S&& o) noexcept: x(o.x), y(o.y) {}
};

void TestIntPush_NoReserve() {
    run_test("IntPush_NoReserve",
        [] {
            vector<int> v;
            for (size_t i = 0; i < N; ++i) v.push_back(int(i));
        },
        [] {
            Vector<int> v;
            for (size_t i = 0; i < N; ++i) v.push_back(int(i));
        });
}

void TestIntPush_WithReserve() {
    run_test("IntPush_WithReserve",
        [] {
            vector<int> v; v.reserve(N);
            for (size_t i = 0; i < N; ++i) v.push_back(int(i));
        },
        [] {
            Vector<int> v; v.reserve(N);
            for (size_t i = 0; i < N; ++i) v.push_back(int(i));
        });
}

void TestIntEmplace_NoReserve() {
    run_test("IntEmplace_NoReserve",
        [] {
            vector<int> v;
            for (size_t i = 0; i < N; ++i) v.emplace_back(int(i));
        },
        [] {
            Vector<int> v;
            for (size_t i = 0; i < N; ++i) v.emplace_back(int(i));
        });
}

void TestIntEmplace_WithReserve() {
    run_test("IntEmplace_WithReserve",
        [] {
            vector<int> v; v.reserve(N);
            for (size_t i = 0; i < N; ++i) v.emplace_back(int(i));
        },
        [] {
            Vector<int> v; v.reserve(N);
            for (size_t i = 0; i < N; ++i) v.emplace_back(int(i));
        });
}

void TestNonTrivialPush_NoReserve() {
    run_test("NonTrivialPush_NoReserve",
        [] {
            vector<S> v;
            for (size_t i = 0; i < N / 10; ++i) v.push_back(S(int(i), int(i)));
        },
        [] {
            Vector<S> v;
            for (size_t i = 0; i < N / 10; ++i) v.push_back(S(int(i), int(i)));
        });
}

void TestNonTrivialPush_WithReserve() {
    run_test("NonTrivialPush_WithReserve",
        [] {
            vector<S> v; v.reserve(N / 10);
            for (size_t i = 0; i < N / 10; ++i) v.push_back(S(int(i), int(i)));
        },
        [] {
            Vector<S> v; v.reserve(N / 10);
            for (size_t i = 0; i < N / 10; ++i) v.push_back(S(int(i), int(i)));
        });
}

int main(int argc, char** argv) {
    if (argc > 1 && std::string(argv[1]) == "-l") longMode = true;

    int R = longMode ? LONG_RUNS : SHORT_RUNS;
    cout << "\n=== Running Speed Tests (N=" << N << ", runs=" << R << ") ===\n\n";

    TestIntPush_NoReserve();
    TestIntPush_WithReserve();
    TestIntEmplace_NoReserve();
    TestIntEmplace_WithReserve();
    TestNonTrivialPush_NoReserve();
    TestNonTrivialPush_WithReserve();

    if (longMode) {
        double avg_speedup = total_speedup / speedup_count;
        cout << "\nOverall average speedup: "
             << (avg_speedup > 1.0 ? COLOR_GREEN : COLOR_RED)
             << "x" << fixed << setprecision(2) << avg_speedup
             << COLOR_RESET << std::endl;
    }

    cout << "\n";
    return 0;
}