#if defined(__ARM_NEON) || defined(__aarch64__)

#include "labneura/backends/cpu/neon.h"
#include "neon_internal.h"
#include "labneura/utils/alignment.hpp"
#include <algorithm>
#include <stdexcept>
#include <arm_neon.h>

namespace labneura {

// FP16 arithmetic requires proper IEEE 754 half-precision handling.
// FP16 is stored as int16_t bit patterns, so we must:
// 1. Reinterpret int16_t as float16_t (bit cast)
// 2. Convert FP16 → FP32 for arithmetic
// 3. Perform the operation in FP32
// 4. Convert FP32 → FP16
// 5. Reinterpret back to int16_t

void NEONBackend::operation_fp16(const TensorBackend& other, OperationType op_type) {
    const int16_t* other_data = other.data_fp16();
    int16_t* this_data = data_fp16();

    // Prefetch distance for FP16: 64 elements (~128 bytes)
    constexpr std::size_t PREFETCH_DISTANCE = 64;

    // Process 8 FP16 values at a time (converts to 2x float32x4)
    // Then handle tails: 4, then scalar
    const std::size_t len8 = (size_ / 8) * 8;
    const std::size_t len4 = len8 + ((size_ - len8) / 4) * 4;

    if (op_type == OperationType::ADD) {
        // Main loop: 8 fp16s at a time
        for (std::size_t i = 0; i < len8; i += 8) {
            if (i + PREFETCH_DISTANCE < size_) {
                __builtin_prefetch(&this_data[i + PREFETCH_DISTANCE], 1, 1);
                __builtin_prefetch(&other_data[i + PREFETCH_DISTANCE], 0, 1);
            }
            
            // Load 8 FP16 values as float16x8 (reinterpret from int16)
            const float16x8_t va = vreinterpretq_f16_s16(vld1q_s16(&this_data[i]));
            const float16x8_t vb = vreinterpretq_f16_s16(vld1q_s16(&other_data[i]));
            
            // Convert low and high halves to float32
            const float32x4_t fa_low = vcvt_f32_f16(vget_low_f16(va));
            const float32x4_t fa_high = vcvt_f32_f16(vget_high_f16(va));
            const float32x4_t fb_low = vcvt_f32_f16(vget_low_f16(vb));
            const float32x4_t fb_high = vcvt_f32_f16(vget_high_f16(vb));
            
            // Add in FP32
            const float32x4_t result_low = vaddq_f32(fa_low, fb_low);
            const float32x4_t result_high = vaddq_f32(fa_high, fb_high);
            
            // Convert back to FP16
            const float16x4_t r_low = vcvt_f16_f32(result_low);
            const float16x4_t r_high = vcvt_f16_f32(result_high);
            const float16x8_t result = vcombine_f16(r_low, r_high);
            
            // Store as int16
            vst1q_s16(&this_data[i], vreinterpretq_s16_f16(result));
        }
        
        // Tail: 4 fp16s
        if (len4 > len8) {
            const std::size_t i = len8;
            const float16x4_t va = vreinterpret_f16_s16(vld1_s16(&this_data[i]));
            const float16x4_t vb = vreinterpret_f16_s16(vld1_s16(&other_data[i]));
            
            const float32x4_t fa = vcvt_f32_f16(va);
            const float32x4_t fb = vcvt_f32_f16(vb);
            const float32x4_t result_f32 = vaddq_f32(fa, fb);
            const float16x4_t result = vcvt_f16_f32(result_f32);
            
            vst1_s16(&this_data[i], vreinterpret_s16_f16(result));
        }
        
        // Scalar tail: process remaining 1-3 elements
        for (std::size_t i = len4; i < size_; ++i) {
            // Use union for type punning (bit cast)
            union { int16_t i; __fp16 f; } a_u, b_u, r_u;
            a_u.i = this_data[i];
            b_u.i = other_data[i];
            r_u.f = static_cast<__fp16>(static_cast<float>(a_u.f) + static_cast<float>(b_u.f));
            this_data[i] = r_u.i;
        }
        return;
    }

    if (op_type == OperationType::SUB) {
        // Main loop: 8 fp16s at a time
        for (std::size_t i = 0; i < len8; i += 8) {
            if (i + PREFETCH_DISTANCE < size_) {
                __builtin_prefetch(&this_data[i + PREFETCH_DISTANCE], 1, 1);
                __builtin_prefetch(&other_data[i + PREFETCH_DISTANCE], 0, 1);
            }
            
            const float16x8_t va = vreinterpretq_f16_s16(vld1q_s16(&this_data[i]));
            const float16x8_t vb = vreinterpretq_f16_s16(vld1q_s16(&other_data[i]));
            
            const float32x4_t fa_low = vcvt_f32_f16(vget_low_f16(va));
            const float32x4_t fa_high = vcvt_f32_f16(vget_high_f16(va));
            const float32x4_t fb_low = vcvt_f32_f16(vget_low_f16(vb));
            const float32x4_t fb_high = vcvt_f32_f16(vget_high_f16(vb));
            
            const float32x4_t result_low = vsubq_f32(fa_low, fb_low);
            const float32x4_t result_high = vsubq_f32(fa_high, fb_high);
            
            const float16x4_t r_low = vcvt_f16_f32(result_low);
            const float16x4_t r_high = vcvt_f16_f32(result_high);
            const float16x8_t result = vcombine_f16(r_low, r_high);
            
            vst1q_s16(&this_data[i], vreinterpretq_s16_f16(result));
        }
        
        // Tail: 4 fp16s
        if (len4 > len8) {
            const std::size_t i = len8;
            const float16x4_t va = vreinterpret_f16_s16(vld1_s16(&this_data[i]));
            const float16x4_t vb = vreinterpret_f16_s16(vld1_s16(&other_data[i]));
            
            const float32x4_t fa = vcvt_f32_f16(va);
            const float32x4_t fb = vcvt_f32_f16(vb);
            const float32x4_t result_f32 = vsubq_f32(fa, fb);
            const float16x4_t result = vcvt_f16_f32(result_f32);
            
            vst1_s16(&this_data[i], vreinterpret_s16_f16(result));
        }
        
        // Scalar tail
        for (std::size_t i = len4; i < size_; ++i) {
            union { int16_t i; __fp16 f; } a_u, b_u, r_u;
            a_u.i = this_data[i];
            b_u.i = other_data[i];
            r_u.f = static_cast<__fp16>(static_cast<float>(a_u.f) - static_cast<float>(b_u.f));
            this_data[i] = r_u.i;
        }
        return;
    }

    if (op_type == OperationType::MUL) {
        // Main loop: 8 fp16s at a time
        for (std::size_t i = 0; i < len8; i += 8) {
            if (i + PREFETCH_DISTANCE < size_) {
                __builtin_prefetch(&this_data[i + PREFETCH_DISTANCE], 1, 1);
                __builtin_prefetch(&other_data[i + PREFETCH_DISTANCE], 0, 1);
            }
            
            const float16x8_t va = vreinterpretq_f16_s16(vld1q_s16(&this_data[i]));
            const float16x8_t vb = vreinterpretq_f16_s16(vld1q_s16(&other_data[i]));
            
            const float32x4_t fa_low = vcvt_f32_f16(vget_low_f16(va));
            const float32x4_t fa_high = vcvt_f32_f16(vget_high_f16(va));
            const float32x4_t fb_low = vcvt_f32_f16(vget_low_f16(vb));
            const float32x4_t fb_high = vcvt_f32_f16(vget_high_f16(vb));
            
            const float32x4_t result_low = vmulq_f32(fa_low, fb_low);
            const float32x4_t result_high = vmulq_f32(fa_high, fb_high);
            
            const float16x4_t r_low = vcvt_f16_f32(result_low);
            const float16x4_t r_high = vcvt_f16_f32(result_high);
            const float16x8_t result = vcombine_f16(r_low, r_high);
            
            vst1q_s16(&this_data[i], vreinterpretq_s16_f16(result));
        }
        
        // Tail: 4 fp16s
        if (len4 > len8) {
            const std::size_t i = len8;
            const float16x4_t va = vreinterpret_f16_s16(vld1_s16(&this_data[i]));
            const float16x4_t vb = vreinterpret_f16_s16(vld1_s16(&other_data[i]));
            
            const float32x4_t fa = vcvt_f32_f16(va);
            const float32x4_t fb = vcvt_f32_f16(vb);
            const float32x4_t result_f32 = vmulq_f32(fa, fb);
            const float16x4_t result = vcvt_f16_f32(result_f32);
            
            vst1_s16(&this_data[i], vreinterpret_s16_f16(result));
        }
        
        // Scalar tail
        for (std::size_t i = len4; i < size_; ++i) {
            union { int16_t i; __fp16 f; } a_u, b_u, r_u;
            a_u.i = this_data[i];
            b_u.i = other_data[i];
            r_u.f = static_cast<__fp16>(static_cast<float>(a_u.f) * static_cast<float>(b_u.f));
            this_data[i] = r_u.i;
        }
        return;
    }

    if (op_type == OperationType::DIV) {
        // Main loop: 8 fp16s at a time
        for (std::size_t i = 0; i < len8; i += 8) {
            if (i + PREFETCH_DISTANCE < size_) {
                __builtin_prefetch(&this_data[i + PREFETCH_DISTANCE], 1, 1);
                __builtin_prefetch(&other_data[i + PREFETCH_DISTANCE], 0, 1);
            }
            
            const float16x8_t va = vreinterpretq_f16_s16(vld1q_s16(&this_data[i]));
            const float16x8_t vb = vreinterpretq_f16_s16(vld1q_s16(&other_data[i]));
            
            const float32x4_t fa_low = vcvt_f32_f16(vget_low_f16(va));
            const float32x4_t fa_high = vcvt_f32_f16(vget_high_f16(va));
            const float32x4_t fb_low = vcvt_f32_f16(vget_low_f16(vb));
            const float32x4_t fb_high = vcvt_f32_f16(vget_high_f16(vb));
            
            const float32x4_t result_low = vdivq_f32(fa_low, fb_low);
            const float32x4_t result_high = vdivq_f32(fa_high, fb_high);
            
            const float16x4_t r_low = vcvt_f16_f32(result_low);
            const float16x4_t r_high = vcvt_f16_f32(result_high);
            const float16x8_t result = vcombine_f16(r_low, r_high);
            
            vst1q_s16(&this_data[i], vreinterpretq_s16_f16(result));
        }
        
        // Tail: 4 fp16s
        if (len4 > len8) {
            const std::size_t i = len8;
            const float16x4_t va = vreinterpret_f16_s16(vld1_s16(&this_data[i]));
            const float16x4_t vb = vreinterpret_f16_s16(vld1_s16(&other_data[i]));
            
            const float32x4_t fa = vcvt_f32_f16(va);
            const float32x4_t fb = vcvt_f32_f16(vb);
            const float32x4_t result_f32 = vdivq_f32(fa, fb);
            const float16x4_t result = vcvt_f16_f32(result_f32);
            
            vst1_s16(&this_data[i], vreinterpret_s16_f16(result));
        }
        
        // Scalar tail
        for (std::size_t i = len4; i < size_; ++i) {
            union { int16_t i; __fp16 f; } a_u, b_u, r_u;
            a_u.i = this_data[i];
            b_u.i = other_data[i];
            float a_f = static_cast<float>(a_u.f);
            float b_f = static_cast<float>(b_u.f);
            r_u.f = static_cast<__fp16>(a_f / b_f);
            this_data[i] = r_u.i;
        }
        return;
    }
}

} // namespace labneura

#endif // __ARM_NEON || __aarch64__
