#pragma once
#include <array>
#include <stdexcept>
#include <string>
#include <string_view>

namespace ling3 {
inline constexpr std::array<std::size_t, 7> kContexts {4096, 8192, 16384, 32768, 65536, 131072, 262144};
// 6 MLA layers, 16 heads, FP16 K192/V128 + five FP32 score arrays.
inline constexpr std::size_t kContextBytesPerToken = 6 * (16 * (192 + 128) * 2 + 5 * 4);
inline double EstimateContextMiB(std::size_t context) {
    return 10034.0 + (static_cast<double>(context) - 4096) * kContextBytesPerToken / (1024 * 1024);
}
inline std::size_t ParseContext(std::string_view value) {
    for (auto context : kContexts) {
        const auto k = std::to_string(context / 1024);
        if (value == k + "K" || value == k + "k" || value == std::to_string(context)) return context;
    }
    throw std::invalid_argument("context must be 4K, 8K, 16K, 32K, 64K, 128K or 256K");
}
} // namespace ling3