#pragma once
#include <cstdint>
#include <string>
#include <unordered_map>

constexpr int TEST_VOCAB_SIZE       = 64;
constexpr int TEST_EMBEDDING_DIM    = 8;
constexpr int TEST_NUM_HEADS        = 2;
constexpr int TEST_HEAD_DIM         = TEST_EMBEDDING_DIM / TEST_NUM_HEADS;
constexpr int TEST_NUM_LAYERS       = 2;
constexpr int TEST_MAX_CTX          = 500;

inline const std::unordered_map<std::string, uint32_t> TEST_VOCAB = {
    // Специальные токены
    {"<s>",         0},
    {"</s>",        1},
    {"<think>",     2},
    {"</think>",    3},

    // Приветствия и базовые фразы
    {"привет",      4},
    {"здравствуй",  5},
    {"как",         6},
    {"дела",        7},
    {"ты",          8},
    {"я",           9},
    {"мы",          10},

    // Вопросы
    {"что",         11},
    {"где",         12},
    {"когда",       13},
    {"почему",      14},
    {"сколько",     15},

    // Существительные
    {"мир",         16},
    {"день",        17},
    {"время",       18},
    {"работа",      19},
    {"код",         20},
    {"программа",   21},
    {"компьютер",   22},
    {"система",     23},
    {"данные",      24},
    {"память",      25},

    // Глаголы
    {"делать",      26},
    {"идти",        27},
    {"работать",    28},
    {"думать",      29},
    {"писать",      30},
    {"читать",      31},
    {"понимать",    32},

    // Прилагательные
    {"хорошо",      33},
    {"плохо",       34},
    {"быстро",      35},
    {"медленно",    36},
    {"новый",       37},
    {"старый",      38},

    // Технические термины
    {"функция",     39},
    {"переменная",  40},
    {"класс",       41},
    {"объект",      42},
    {"массив",      43},
    {"указатель",   44},
    {"алгоритм",    45},

    // Союзы и предлоги
    {"и",           46},
    {"или",         47},
    {"но",          48},
    {"если",        49},
    {"то",          50},
    {"в",           51},
    {"на",          52},
    {"с",           53},
    {"из",          54},
    {"для",         55},

    // Частицы
    {"не",          56},
    {"ли",          57},
    {"же",          58},
    {"бы",          59},

    // Знаки препинания (опционально)
    {".",           60},
    {",",           61},
    {"?",           62},
    {"!",           63},
};

inline constexpr float TEST_W_EMB[TEST_VOCAB_SIZE][TEST_EMBEDDING_DIM] = {
    // <s>
    {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f},
    // </s>
    {0.5f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f},
    // <think>
    {0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f},
    // </think>
    {0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f},
    // привет
    {0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f},
    // здравствуй
    {0.85f, 0.75f, 0.65f, 0.55f, 0.45f, 0.35f, 0.25f, 0.15f},
    // как
    {0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 0.1f},
    // дела
    {0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 0.1f, 0.2f},
    // ты
    {0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 0.1f, 0.2f, 0.3f},
    // я
    {0.6f, 0.7f, 0.8f, 0.9f, 0.1f, 0.2f, 0.3f, 0.4f},
    // мы
    {0.7f, 0.8f, 0.9f, 0.1f, 0.2f, 0.3f, 0.4f, 0.5f},
    // что
    {0.15f, 0.25f, 0.35f, 0.45f, 0.55f, 0.65f, 0.75f, 0.85f},
    // где
    {0.25f, 0.35f, 0.45f, 0.55f, 0.65f, 0.75f, 0.85f, 0.95f},
    // когда
    {0.35f, 0.45f, 0.55f, 0.65f, 0.75f, 0.85f, 0.95f, 0.05f},
    // почему
    {0.45f, 0.55f, 0.65f, 0.75f, 0.85f, 0.95f, 0.05f, 0.15f},
    // сколько
    {0.55f, 0.65f, 0.75f, 0.85f, 0.95f, 0.05f, 0.15f, 0.25f},
    // мир
    {0.1f, 0.9f, 0.2f, 0.8f, 0.3f, 0.7f, 0.4f, 0.6f},
    // день
    {0.2f, 0.8f, 0.3f, 0.7f, 0.4f, 0.6f, 0.5f, 0.5f},
    // время
    {0.3f, 0.7f, 0.4f, 0.6f, 0.5f, 0.5f, 0.6f, 0.4f},
    // работа
    {0.4f, 0.6f, 0.5f, 0.5f, 0.6f, 0.4f, 0.7f, 0.3f},
    // код
    {0.5f, 0.5f, 0.6f, 0.4f, 0.7f, 0.3f, 0.8f, 0.2f},
    // программа
    {0.6f, 0.4f, 0.7f, 0.3f, 0.8f, 0.2f, 0.9f, 0.1f},
    // компьютер
    {0.7f, 0.3f, 0.8f, 0.2f, 0.9f, 0.1f, 0.0f, 0.0f},
    // система
    {0.8f, 0.2f, 0.9f, 0.1f, 0.0f, 0.0f, 0.1f, 0.9f},
    // данные
    {0.9f, 0.1f, 0.0f, 0.0f, 0.1f, 0.9f, 0.2f, 0.8f},
    // память
    {0.0f, 0.0f, 0.1f, 0.9f, 0.2f, 0.8f, 0.3f, 0.7f},
    // делать
    {0.12f, 0.23f, 0.34f, 0.45f, 0.56f, 0.67f, 0.78f, 0.89f},
    // идти
    {0.23f, 0.34f, 0.45f, 0.56f, 0.67f, 0.78f, 0.89f, 0.12f},
    // работать
    {0.34f, 0.45f, 0.56f, 0.67f, 0.78f, 0.89f, 0.12f, 0.23f},
    // думать
    {0.45f, 0.56f, 0.67f, 0.78f, 0.89f, 0.12f, 0.23f, 0.34f},
    // писать
    {0.56f, 0.67f, 0.78f, 0.89f, 0.12f, 0.23f, 0.34f, 0.45f},
    // читать
    {0.67f, 0.78f, 0.89f, 0.12f, 0.23f, 0.34f, 0.45f, 0.56f},
    // понимать
    {0.78f, 0.89f, 0.12f, 0.23f, 0.34f, 0.45f, 0.56f, 0.67f},
    // хорошо
    {0.89f, 0.12f, 0.23f, 0.34f, 0.45f, 0.56f, 0.67f, 0.78f},
    // плохо
    {0.11f, 0.22f, 0.33f, 0.44f, 0.55f, 0.66f, 0.77f, 0.88f},
    // быстро
    {0.22f, 0.33f, 0.44f, 0.55f, 0.66f, 0.77f, 0.88f, 0.99f},
    // медленно
    {0.33f, 0.44f, 0.55f, 0.66f, 0.77f, 0.88f, 0.99f, 0.11f},
    // новый
    {0.44f, 0.55f, 0.66f, 0.77f, 0.88f, 0.99f, 0.11f, 0.22f},
    // старый
    {0.55f, 0.66f, 0.77f, 0.88f, 0.99f, 0.11f, 0.22f, 0.33f},
    // функция
    {0.66f, 0.77f, 0.88f, 0.99f, 0.11f, 0.22f, 0.33f, 0.44f},
    // переменная
    {0.77f, 0.88f, 0.99f, 0.11f, 0.22f, 0.33f, 0.44f, 0.55f},
    // класс
    {0.88f, 0.99f, 0.11f, 0.22f, 0.33f, 0.44f, 0.55f, 0.66f},
    // объект
    {0.99f, 0.11f, 0.22f, 0.33f, 0.44f, 0.55f, 0.66f, 0.77f},
    // массив
    {0.13f, 0.24f, 0.35f, 0.46f, 0.57f, 0.68f, 0.79f, 0.80f},
    // указатель
    {0.24f, 0.35f, 0.46f, 0.57f, 0.68f, 0.79f, 0.80f, 0.13f},
    // алгоритм
    {0.35f, 0.46f, 0.57f, 0.68f, 0.79f, 0.80f, 0.13f, 0.24f},
    // и
    {0.46f, 0.57f, 0.68f, 0.79f, 0.80f, 0.13f, 0.24f, 0.35f},
    // или
    {0.57f, 0.68f, 0.79f, 0.80f, 0.13f, 0.24f, 0.35f, 0.46f},
    // но
    {0.68f, 0.79f, 0.80f, 0.13f, 0.24f, 0.35f, 0.46f, 0.57f},
    // если
    {0.79f, 0.80f, 0.13f, 0.24f, 0.35f, 0.46f, 0.57f, 0.68f},
    // то
    {0.80f, 0.13f, 0.24f, 0.35f, 0.46f, 0.57f, 0.68f, 0.79f},
    // в
    {0.14f, 0.25f, 0.36f, 0.47f, 0.58f, 0.69f, 0.70f, 0.81f},
    // на
    {0.25f, 0.36f, 0.47f, 0.58f, 0.69f, 0.70f, 0.81f, 0.14f},
    // с
    {0.36f, 0.47f, 0.58f, 0.69f, 0.70f, 0.81f, 0.14f, 0.25f},
    // из
    {0.47f, 0.58f, 0.69f, 0.70f, 0.81f, 0.14f, 0.25f, 0.36f},
    // для
    {0.58f, 0.69f, 0.70f, 0.81f, 0.14f, 0.25f, 0.36f, 0.47f},
    // не
    {0.69f, 0.70f, 0.81f, 0.14f, 0.25f, 0.36f, 0.47f, 0.58f},
    // ли
    {0.70f, 0.81f, 0.14f, 0.25f, 0.36f, 0.47f, 0.58f, 0.69f},
    // же
    {0.81f, 0.14f, 0.25f, 0.36f, 0.47f, 0.58f, 0.69f, 0.70f},
    // бы
    {0.16f, 0.27f, 0.38f, 0.49f, 0.50f, 0.61f, 0.72f, 0.83f},
    // .
    {0.27f, 0.38f, 0.49f, 0.50f, 0.61f, 0.72f, 0.83f, 0.16f},
    // ,
    {0.38f, 0.49f, 0.50f, 0.61f, 0.72f, 0.83f, 0.16f, 0.27f},
    // ?
    {0.49f, 0.50f, 0.61f, 0.72f, 0.83f, 0.16f, 0.27f, 0.38f},
    // !
    {0.50f, 0.61f, 0.72f, 0.83f, 0.16f, 0.27f, 0.38f, 0.49f},
};

inline constexpr float TEST_W_QKV[TEST_NUM_LAYERS][TEST_EMBEDDING_DIM][3][TEST_NUM_HEADS] = {
    {
        {{0.1f, 0.2f}, {0.3f, 0.4f}, {0.5f, 0.6f}},
        {{0.7f, 0.8f}, {0.9f, 1.0f}, {1.1f, 1.2f}},
        {{0.1f, 0.1f}, {0.2f, 0.2f}, {0.3f, 0.3f}},
        {{0.4f, 0.4f}, {0.5f, 0.5f}, {0.6f, 0.6f}},
        {{0.7f, 0.7f}, {0.8f, 0.8f}, {0.9f, 0.9f}},
        {{1.0f, 1.0f}, {1.1f, 1.1f}, {1.2f, 1.2f}},
        {{0.1f, 0.3f}, {0.5f, 0.7f}, {0.9f, 1.1f}},
        {{0.2f, 0.4f}, {0.6f, 0.8f}, {1.0f, 1.2f}}
    },
    {
        {{0.9f, 0.8f}, {0.7f, 0.6f}, {0.5f, 0.4f}},
        {{0.3f, 0.2f}, {0.1f, 0.0f}, {-0.1f, -0.2f}},
        {{0.5f, 0.5f}, {0.6f, 0.6f}, {0.7f, 0.7f}},
        {{0.8f, 0.8f}, {0.9f, 0.9f}, {1.0f, 1.0f}},
        {{0.2f, 0.4f}, {0.6f, 0.8f}, {1.0f, 1.2f}},
        {{0.3f, 0.5f}, {0.7f, 0.9f}, {1.1f, 1.3f}},
        {{0.4f, 0.6f}, {0.8f, 1.0f}, {1.2f, 1.4f}},
        {{0.5f, 0.7f}, {0.9f, 1.1f}, {1.3f, 1.5f}}
    }
};
inline constexpr float TEST_W_O[TEST_NUM_LAYERS][TEST_NUM_HEADS * TEST_HEAD_DIM][TEST_EMBEDDING_DIM] = {
    {
        {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f},
        {0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f},
        {0.1f, 0.9f, 0.2f, 0.8f, 0.3f, 0.7f, 0.4f, 0.6f},
        {0.5f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f},
        {0.2f, 0.4f, 0.6f, 0.8f, 1.0f, 1.2f, 1.4f, 1.6f},
        {0.3f, 0.5f, 0.7f, 0.9f, 1.1f, 1.3f, 1.5f, 1.7f},
        {0.4f, 0.6f, 0.8f, 1.0f, 1.2f, 1.4f, 1.6f, 1.8f},
        {0.5f, 0.7f, 0.9f, 1.1f, 1.3f, 1.5f, 1.7f, 1.9f}
    },
    {
        {0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f, 0.1f},
        {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f},
        {0.9f, 0.1f, 0.8f, 0.2f, 0.7f, 0.3f, 0.6f, 0.4f},
        {0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 1.0f, 1.1f, 1.2f},
        {0.1f, 0.3f, 0.5f, 0.7f, 0.9f, 1.1f, 1.3f, 1.5f},
        {0.2f, 0.4f, 0.6f, 0.8f, 1.0f, 1.2f, 1.4f, 1.6f},
        {0.3f, 0.5f, 0.7f, 0.9f, 1.1f, 1.3f, 1.5f, 1.7f},
        {0.4f, 0.6f, 0.8f, 1.0f, 1.2f, 1.4f, 1.6f, 1.8f}
    }
};