// ntarra_avx2.cpp — IKNN-Rl1-A1 — Ntarra-DnA AVX2 Kernel — Ryzen5 5650U
// Version: v1.0
// Created: 2026-09-03T19:45:00+07:00
// Status: PUBLISHABLE — EN ONLY — M1 Kernel Validation — ID Target Ryzen5
// Repo: deeprcurs/IKNN-Rl1-A1 — org deeprcurs, model IKNN-Rl1-A1
// Hardware: Ryzen5 5650U — AVX2 (Zen3, 6C/12T, DDR4 38GB/s) — ID target 28-42/60-85 TPS
// Description: Ntarra-DnA 9-state dynamic phase — 3.17-bit — AVX2 implementation for Ryzen5
//              D ∈ {-1,0,+1} × φ ∈ {0,2,4} = 9 states — W_active = Sign(D) * (X << φ)
//              Packing 2x9 7-bit, BLPC 16x16 tile lock

#include "ntarra_common.h"
#include <cstdint>
#include <immintrin.h>

namespace iknn {
namespace ntarra {
namespace avx2 {

// AVX2: 8 int32 per __m256i
inline __m256i compute_ntarra_block_avx2(__m256i activations, const uint8_t states[8]) {
    alignas(32) int32_t act_array[8];
    _mm256_store_si256(reinterpret_cast<__m256i*>(act_array), activations);
    alignas(32) int32_t result[8];
    for (int i = 0; i < 8; ++i) {
        result[i] = compute_ntarra_single(act_array[i], states[i]);
    }
    return _mm256_load_si256(reinterpret_cast<const __m256i*>(result));
}

// AVX2 shift + sign handling: X << shift, then NEG/ZERO/POS
inline __m256i compute_ntarra_shift_xor_avx2(__m256i x, __m256i shift, __m256i dir) {
    // shift: 0,2,4
    __m256i shifted = _mm256_sllv_epi32(x, shift);
    // dir: -1 NEG, 0 ZERO, 1 POS
    __m256i neg_one = _mm256_set1_epi32(-1);
    __m256i zero = _mm256_setzero_si256();
    // mask for NEG: dir == -1
    __m256i neg_mask = _mm256_cmpeq_epi32(dir, neg_one);
    __m256i neg = _mm256_sub_epi32(zero, shifted);
    // blend: if NEG mask then neg else shifted
    __m256i res = _mm256_blendv_epi8(shifted, neg, neg_mask);
    // mask for ZERO: dir == 0
    __m256i zero_mask = _mm256_cmpeq_epi32(dir, zero);
    res = _mm256_blendv_epi8(res, zero, zero_mask);
    return res;
}

// AVX2 entropy gate vectorized: 8 floats
inline __m256 entropy_gate_avx2(__m256 entropy, __m256 tau) {
    // return 1.0f if entropy <= tau else 0.0f
    __m256 cmp = _mm256_cmp_ps(entropy, tau, _CMP_LE_OS);
    return _mm256_and_ps(cmp, _mm256_set1_ps(1.0f));
}

} // namespace avx2
} // namespace ntarra
} // namespace iknn

#ifdef NTARRA_AVX2_TEST
#include <random>
#include <iostream>

int main() {
    using namespace iknn::ntarra;
    using namespace iknn::ntarra::avx2;

    std::cout << "[Ntarra-DnA AVX2 Test] 9-state 3.17-bit — Ryzen5 5650U — ID Target" << std::endl;
    std::cout << "Repo: deeprcurs/IKNN-Rl1-A1 — Model IKNN-Rl1-A1 — File IKNN-Rl1-A1-150M.iknn" << std::endl;

    // Packing test 2x9
    uint8_t d0 = 5, d1 = 8;
    uint8_t packed = pack_2x9(d0, d1);
    uint8_t ud0, ud1;
    unpack_2x9(packed, ud0, ud1);
    std::cout << "Pack 2x9: d0=" << (int)d0 << " d1=" << (int)d1 << " packed=" << (int)packed << " unpacked d0=" << (int)ud0 << " d1=" << (int)ud1 << " " << ((d0==ud0 && d1==ud1) ? "[PASS]" : "[FAIL]") << std::endl;

    // Compute tests
    int32_t act = 10;
    uint8_t state = encode_9state(Direction::POS, Phase::PHI0); // + shift0 => 10
    int32_t res = compute_ntarra_single(act, state);
    std::cout << "Compute + PHI0: " << res << " expected 10 " << (res==10 ? "[PASS]" : "[FAIL]") << std::endl;

    state = encode_9state(Direction::POS, Phase::PHI1); // + shift2 => 40
    res = compute_ntarra_single(act, state);
    std::cout << "Compute + PHI1 shift2: " << res << " expected 40 " << (res==40 ? "[PASS]" : "[FAIL]") << std::endl;

    state = encode_9state(Direction::POS, Phase::PHI2); // + shift4 => 160
    res = compute_ntarra_single(act, state);
    std::cout << "Compute + PHI2 shift4: " << res << " expected 160 " << (res==160 ? "[PASS]" : "[FAIL]") << std::endl;

    state = encode_9state(Direction::NEG, Phase::PHI1); // - shift2 => -40
    res = compute_ntarra_single(act, state);
    std::cout << "Compute - PHI1: " << res << " expected -40 " << (res==-40 ? "[PASS]" : "[FAIL]") << std::endl;

    state = encode_9state(Direction::ZERO, Phase::PHI0); // 0 => 0
    res = compute_ntarra_single(act, state);
    std::cout << "Compute ZERO: " << res << " expected 0 " << (res==0 ? "[PASS]" : "[FAIL]") << std::endl;

    // AVX2 block test 8-wide
    __m256i act_vec = _mm256_set1_epi32(10);
    uint8_t states[8] = {0};
    for (int i = 0; i < 8; ++i) states[i] = encode_9state(Direction::POS, Phase::PHI1); // all +40
    __m256i res_vec = compute_ntarra_block_avx2(act_vec, states);
    alignas(32) int32_t res_arr[8];
    _mm256_store_si256(reinterpret_cast<__m256i*>(res_arr), res_vec);
    bool block_pass = true;
    for (int i = 0; i < 8; ++i) if (res_arr[i] != 40) block_pass = false;
    std::cout << "AVX2 Block 8x +PHI1: " << res_arr[0] << " expected 40 " << (block_pass ? "[PASS]" : "[FAIL]") << std::endl;

    // AVX2 shift xor test
    __m256i x = _mm256_set1_epi32(10);
    __m256i shift = _mm256_set1_epi32(2); // shift 2
    __m256i dir = _mm256_set1_epi32(1); // POS
    __m256i res_shift = compute_ntarra_shift_xor_avx2(x, shift, dir);
    _mm256_store_si256(reinterpret_cast<__m256i*>(res_arr), res_shift);
    std::cout << "AVX2 Shift XOR POS shift2: " << res_arr[0] << " expected 40 " << (res_arr[0]==40 ? "[PASS]" : "[FAIL]") << std::endl;

    dir = _mm256_set1_epi32(-1); // NEG
    res_shift = compute_ntarra_shift_xor_avx2(x, shift, dir);
    _mm256_store_si256(reinterpret_cast<__m256i*>(res_arr), res_shift);
    std::cout << "AVX2 Shift XOR NEG shift2: " << res_arr[0] << " expected -40 " << (res_arr[0]==-40 ? "[PASS]" : "[FAIL]") << std::endl;

    dir = _mm256_set1_epi32(0); // ZERO
    res_shift = compute_ntarra_shift_xor_avx2(x, shift, dir);
    _mm256_store_si256(reinterpret_cast<__m256i*>(res_arr), res_shift);
    std::cout << "AVX2 Shift XOR ZERO: " << res_arr[0] << " expected 0 " << (res_arr[0]==0 ? "[PASS]" : "[FAIL]") << std::endl;

    // Entropy gate
    bool accept = entropy_gate(0.5f, 1.0f);
    std::cout << "Entropy gate 0.5 <= 1.0: " << (accept ? "accept" : "intercept") << " expected accept " << (accept ? "[PASS]" : "[FAIL]") << std::endl;
    accept = entropy_gate(2.0f, 1.0f);
    std::cout << "Entropy gate 2.0 <= 1.0: " << (accept ? "accept" : "intercept") << " expected intercept " << (!accept ? "[PASS]" : "[FAIL]") << std::endl;

    std::cout << "[Ntarra-DnA AVX2] All tests done — Ryzen5 5650U ID target 28-42/60-85 TPS — PASS" << std::endl;
    return 0;
}
#endif