#include <balancer/kernel/rpslimiter/rate.h>

#include <library/cpp/testing/unittest/registar.h>

using namespace NSrvKernel::NRpsLimiter;

Y_UNIT_TEST_SUITE(TRateTest) {

    void DoTestStorage(const TLocalRate& rate, TLocalRate o, TSourceLocation src) {
        {
            TLocalRate other = o;
            other.Unpack(rate.Pack());
            UNIT_ASSERT_VALUES_EQUAL_C(other, rate, rate << " at " << src);
            UNIT_ASSERT_VALUES_EQUAL_C(other.Time(), rate.Time(), rate << " at " << src);
            UNIT_ASSERT_VALUES_EQUAL_C(other.Window(), rate.Window(), rate << " at " << src);
            UNIT_ASSERT_VALUES_EQUAL_C(other.Rate(), rate.Rate(), rate << " at " << src);
        }

        {
            TAtomicStorage<TUi128> s;
            s.Update([=](TUi128) -> TUi128 {
                return rate.Pack();
            });
            UNIT_ASSERT_EQUAL_C(s.Get(), rate.Pack(), rate << " at " << src);
            s.Update([=](TUi128 s) -> TUi128 {
                UNIT_ASSERT_EQUAL_C(rate.Pack(), s, rate << " at " << src);
                return s;
            });
        }
    }

    struct TLocalRateCase {
        TDuration TimeInc;
        float CntInc = 0;
        double ExpectedRate = 0;
        double RateError = 0.0001;
    };

    using TLocalRateCases = TVector<TLocalRateCase>;

    void DoTestLocalRate(TLocalRate rate, TLocalRateCases incs, TSourceLocation src) {
        for (size_t i = 0; i < incs.size(); ++i) {
            rate.IncTime(incs[i].TimeInc);
            UNIT_ASSERT(rate.IncCnt(incs[i].CntInc));
            DoTestStorage(rate, *TLocalRate::Create(rate.Window(), rate.Start()), src);
            UNIT_ASSERT_DOUBLES_EQUAL_C(
                rate.Rate(),
                incs[i].ExpectedRate,
                incs[i].RateError * std::max({incs[i].ExpectedRate, rate.Rate(), 1.}),
                rate << " at " << src
            );
        }
    }

#define Y_BALANCER_LOCAL_RATE_TEST(rate, incs) DoTestLocalRate(rate, incs, __LOCATION__)

    Y_UNIT_TEST(TestLocalRate) {
        UNIT_ASSERT(!TLocalRate::Create({}));
        UNIT_ASSERT(TLocalRate::Create(TDuration::MicroSeconds(1)));

        const auto now = TInstant::Seconds(1600000000);
        {
            const auto window = TDuration::Seconds(1);
            auto lr = *TLocalRate::Create(window, now);
            UNIT_ASSERT(!lr.IncCnt(std::numeric_limits<float>::quiet_NaN()));
            UNIT_ASSERT(!lr.IncCnt(std::numeric_limits<float>::signaling_NaN()));
            UNIT_ASSERT(!lr.IncCnt(std::numeric_limits<float>::infinity()));
            UNIT_ASSERT(!lr.IncCnt(-std::numeric_limits<float>::infinity()));
            UNIT_ASSERT_VALUES_EQUAL(lr.Rate(), 0);
            UNIT_ASSERT(lr.IncCnt(Max<float>()));
            UNIT_ASSERT_VALUES_EQUAL(lr.Rate(), double(Max<float>()));
            UNIT_ASSERT(lr.IncCnt(Max<float>()));
            UNIT_ASSERT_VALUES_EQUAL(lr.Rate(), double(Max<float>()));
            UNIT_ASSERT(lr.IncCnt(-Max<float>()));
            UNIT_ASSERT_VALUES_EQUAL(lr.Rate(), 0);
            UNIT_ASSERT(lr.IncCnt(-Max<float>()));
            UNIT_ASSERT_VALUES_EQUAL(lr.Rate(), -Max<float>());
            UNIT_ASSERT(lr.IncCnt(-Max<float>()));
            UNIT_ASSERT_VALUES_EQUAL(lr.Rate(), -Max<float>());
        }
        {
            const auto window = TDuration::Seconds(1);

            Y_BALANCER_LOCAL_RATE_TEST(
                *TLocalRate::Create(window, now), {});
            Y_BALANCER_LOCAL_RATE_TEST(
                *TLocalRate::Create(window, now), (TLocalRateCases{{.CntInc=1, .ExpectedRate=1}}));
            Y_BALANCER_LOCAL_RATE_TEST(*TLocalRate::Create(window, now), (TLocalRateCases{
                {.CntInc=1, .ExpectedRate=1},
                {.TimeInc=window, .CntInc=1, .ExpectedRate=1},
                {.TimeInc=2*window, .CntInc=1, .ExpectedRate=1},
                {.TimeInc=3*window, .CntInc=1, .ExpectedRate=1},
            }));
            const auto us = TDuration::MicroSeconds(1);
            const auto ms = TDuration::MilliSeconds(1);
            Y_BALANCER_LOCAL_RATE_TEST(*TLocalRate::Create(window, now), (TLocalRateCases{
                {.TimeInc=us, .CntInc=1, .ExpectedRate=1},              // [_|_______][.|_______]
                {.TimeInc=window + us, .CntInc=1, .ExpectedRate=2},     // [_|______.][.|_______]
                {.TimeInc=10*ms, .CntInc=0, .ExpectedRate=1.99},        // [_|______.][.|_______]
                {.TimeInc=10*ms, .CntInc=0, .ExpectedRate=1.98},        // [_|______.][.|_______]
                {.TimeInc=10*ms, .CntInc=0, .ExpectedRate=1.97},        // [_|______.][.|_______]
                {.TimeInc=10*ms, .CntInc=0, .ExpectedRate=1.96},        // [_|______.][.|_______]
                {.TimeInc=window + us, .CntInc=1, .ExpectedRate=1.96},  // [_|______.][.|_______]
                {.TimeInc=2*window + us, .CntInc=1, .ExpectedRate=1},   // [_|_______][.|_______]
                {.TimeInc=2*window + us, .CntInc=1, .ExpectedRate=1},   // [_|_______][.|_______]
            }));
        }
        for (auto start : std::initializer_list<TInstant>{
            {},
            TInstant::MicroSeconds(1),
            now,
            now + TDuration::Days(100 * 365)
        }) {
            for (auto window : std::initializer_list<TDuration>{
                TDuration::MilliSeconds(1),
                TDuration::MicroSeconds(1500),
                TDuration::Days(100 * 365),
            }) {
                // How the rate interpolation works
                Y_BALANCER_LOCAL_RATE_TEST(*TLocalRate::Create(window, start), (TLocalRateCases{
                    {.TimeInc=window / 4, .CntInc=1, .ExpectedRate=1},  // [_|___][.|___]
                    {.TimeInc={}, .CntInc=1, .ExpectedRate=2},          // [_|___][:|___]
                    {.TimeInc=window / 2, .CntInc=2, .ExpectedRate=4},  // [___|_][:_:|_]
                    {.TimeInc=window / 2, .CntInc=2, .ExpectedRate=5},  // [.|...][:|___]
                    {.TimeInc=window / 2, .CntInc=2, .ExpectedRate=5},  // [...|.][:_:|_]
                    {.TimeInc=window / 4, .CntInc=0, .ExpectedRate=4},  // [|....][|____]
                    {.TimeInc=window / 4, .CntInc=2, .ExpectedRate=5},  // [.|...][:|___]
                    {.TimeInc=window / 4, .CntInc=0, .ExpectedRate=4},  // [..|..][:_|__]
                    {.TimeInc=window / 4, .CntInc=1, .ExpectedRate=4},  // [...|.][:_.|_]
                    {.TimeInc=window / 4, .CntInc=1, .ExpectedRate=4},  // [|....][|____]
                    {.TimeInc=window / 4, .CntInc=2, .ExpectedRate=5},  // [.|...][:|___]
                    {.TimeInc=window / 2, .CntInc=2, .ExpectedRate=5},  // [...|.][:_:|_]
                    {.TimeInc=window / 2, .CntInc=0, .ExpectedRate=3},  // [.|...][_|___]
                    {.TimeInc=window / 2, .CntInc=0, .ExpectedRate=1},  // [...|.][___|_]
                    {.TimeInc=window / 2, .CntInc=0, .ExpectedRate=0},  // [_|___][_|___]
                }));
                // Deterministically forgetting about the huge burst
                Y_BALANCER_LOCAL_RATE_TEST(*TLocalRate::Create(window, start), (TLocalRateCases{
                    {.TimeInc=window / 4, .CntInc=Max<ui32>(), .ExpectedRate=float(Max<ui32>())},   // [_|___][@|___]
                    {.TimeInc=window / 4, .CntInc=0, .ExpectedRate=float(Max<ui32>())},             // [__|__][@_|__]
                    {.TimeInc=window / 4, .CntInc=0, .ExpectedRate=float(Max<ui32>())},             // [___|_][@__|_]
                    {.TimeInc=window / 4, .CntInc=0, .ExpectedRate=float(Max<ui32>())},             // [|****][|____]
                    {.TimeInc=window / 4, .CntInc=0, .ExpectedRate=float(Max<ui32>() / 4 * 3)},     // [*|***][_|___]
                    {.TimeInc=window / 4, .CntInc=0, .ExpectedRate=float(Max<ui32>() / 2)},         // [**|**][__|__]
                    {.TimeInc=window / 4, .CntInc=0, .ExpectedRate=float(Max<ui32>() / 4)},         // [***|*][___|_]
                    {.TimeInc=window / 4, .CntInc=0, .ExpectedRate=0},                              // [|____][|____]
                }));
            }
        }
    }

    void DoTestStorage(const TPeerRate& rate, TPeerRate o, TSourceLocation src) {
        {
            TPeerRate other = o;
            other.Unpack(rate.Pack());
            UNIT_ASSERT_VALUES_EQUAL_C(other, rate, rate << " at " << src);
            UNIT_ASSERT_VALUES_EQUAL_C(other.Window(), rate.Window(), rate << " at " << src);
            UNIT_ASSERT_VALUES_EQUAL_C(other.Rate(), rate.Rate(), rate << " at " << src);
        }

        {
            TAtomicStorage<ui64> s;
            s.Update([=](ui64) -> ui64 {
                return rate.Pack();
            });
            UNIT_ASSERT_EQUAL_C(s.Get(), rate.Pack(), rate << " at " << src);
            s.Update([=](ui64 s) -> ui64 {
                UNIT_ASSERT_EQUAL_C(rate.Pack(), s, rate << " at " << src);
                return s;
            });
        }
    }

#define Y_BALANCER_PACKING_TEST(rate, other) DoTestStorage(rate, other, __LOCATION__)

    void DoTestPeerRateCreate(TDuration w, double r, TMaybe<double> expected, TSourceLocation src) {
        const auto res = TPeerRate::Create(w, r);
        if (expected) {
            UNIT_ASSERT_C(res, " at " << src);
            UNIT_ASSERT_VALUES_EQUAL_C(res->Window(), w, " at " << src);
            UNIT_ASSERT_VALUES_EQUAL_C(res->Rate(), expected, " at " << src);
            DoTestStorage(*res, *TPeerRate::Create(w), src);
        } else {
            UNIT_ASSERT_C(!res, " at " << src);
        }
    }

#define Y_BALANCER_REMOTE_RATE_CREATE_TEST(w, r, e) DoTestPeerRateCreate(w, r, e, __LOCATION__)

    Y_UNIT_TEST(TestPeerRateCreate) {
        const auto w1ms = TDuration::MilliSeconds(1);
        const auto w100y = TDuration::Days(100 * 365);
        Y_BALANCER_REMOTE_RATE_CREATE_TEST({}, 0, Nothing());
        Y_BALANCER_REMOTE_RATE_CREATE_TEST(w1ms, std::numeric_limits<double>::quiet_NaN(), Nothing());
        Y_BALANCER_REMOTE_RATE_CREATE_TEST(w1ms, std::numeric_limits<double>::signaling_NaN(), Nothing());
        Y_BALANCER_REMOTE_RATE_CREATE_TEST(w1ms, std::numeric_limits<double>::infinity(), Nothing());
        Y_BALANCER_REMOTE_RATE_CREATE_TEST(w1ms, -std::numeric_limits<double>::infinity(), Nothing());

        Y_BALANCER_REMOTE_RATE_CREATE_TEST(TDuration::MicroSeconds(1), 0, 0);

        for (auto w : std::initializer_list<TDuration>{w1ms, w100y}) {
            Y_BALANCER_REMOTE_RATE_CREATE_TEST(w, 1, 1);
            Y_BALANCER_REMOTE_RATE_CREATE_TEST(w, Max<double>(), Max<double>());
            Y_BALANCER_REMOTE_RATE_CREATE_TEST(w, Min<double>(), Min<double>());
            Y_BALANCER_REMOTE_RATE_CREATE_TEST(w, -1, -1);
            Y_BALANCER_REMOTE_RATE_CREATE_TEST(w, -Max<double>(), -Max<double>());
            Y_BALANCER_REMOTE_RATE_CREATE_TEST(w, -Min<double>(), -Min<double>());
        }
    }

    Y_UNIT_TEST(TestPeerRateUpdate) {
        const auto window = TDuration::Seconds(1);

        const auto rr = *TPeerRate::Create(window, 100.5);
        {
            // Matching windows
            auto rr1 = rr;
            rr1.Update(*TPeerRate::Create(window, 200.1));
            UNIT_ASSERT_DOUBLES_EQUAL(rr1.Rate(), 200.1, 0.000001);
            UNIT_ASSERT_VALUES_EQUAL(rr1.Window(), window);
        }
        {
            // Windows do not match, rate needs rescale
            auto rr1 = rr;
            rr1.Update(*TPeerRate::Create(window / 2, 200.1));
            UNIT_ASSERT_DOUBLES_EQUAL(rr1.Rate(), 400.2, 0.000001);
            UNIT_ASSERT_VALUES_EQUAL(rr1.Window(), window);
        }
    }

    Y_UNIT_TEST(TestPeerTime) {
        TPeerTime pt{
            .RemoteTime=TInstant::MicroSeconds(1607455072'654321),
            .LocalTime=TInstant::MicroSeconds(1607455072'123456),
        };

        TPeerTime other;
        other.Unpack(pt.Pack());
        UNIT_ASSERT_VALUES_EQUAL(other, pt);
    }
}
