#include "labneura/backends/cpu/generic.h"
#include <algorithm>
#include <stdexcept>
#include <cstdint>
#include <cmath>

namespace labneura {

// Helper functions for IEEE 754 half-precision (FP16) conversion
// FP16 format: 1 sign bit, 5 exponent bits, 10 mantissa bits

static inline float fp16_to_fp32(int16_t h) {
    uint16_t hu = static_cast<uint16_t>(h);
    uint32_t sign = (hu & 0x8000) << 16;
    uint32_t exp = (hu >> 10) & 0x1F;
    uint32_t mant = hu & 0x03FF;
    
    uint32_t f;
    if (exp == 0) {
        if (mant == 0) {
            // Zero
            f = sign;
        } else {
            // Subnormal - convert to normalized
            exp = 1;
            while ((mant & 0x0400) == 0) {
                mant <<= 1;
                exp--;
            }
            mant &= 0x03FF;
            f = sign | ((exp + 127 - 15) << 23) | (mant << 13);
        }
    } else if (exp == 31) {
        // Inf or NaN
        f = sign | 0x7F800000 | (mant << 13);
    } else {
        // Normalized
        f = sign | ((exp + 127 - 15) << 23) | (mant << 13);
    }
    
    float result;
    std::memcpy(&result, &f, sizeof(float));
    return result;
}

static inline int16_t fp32_to_fp16(float f) {
    uint32_t fu;
    std::memcpy(&fu, &f, sizeof(uint32_t));
    
    uint32_t sign = (fu >> 16) & 0x8000;
    int32_t exp = ((fu >> 23) & 0xFF) - 127 + 15;
    uint32_t mant = fu & 0x007FFFFF;
    
    uint16_t h;
    if (exp <= 0) {
        if (exp < -10) {
            // Too small - underflow to zero
            h = static_cast<uint16_t>(sign);
        } else {
            // Subnormal
            mant |= 0x00800000;
            int shift = 14 - exp;
            h = static_cast<uint16_t>(sign | (mant >> shift));
        }
    } else if (exp >= 31) {
        // Overflow to infinity
        h = static_cast<uint16_t>(sign | 0x7C00);
    } else {
        // Normalized
        h = static_cast<uint16_t>(sign | (exp << 10) | (mant >> 13));
    }
    
    return static_cast<int16_t>(h);
}

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

    // FP16 arithmetic: convert to FP32, operate, convert back
    
    if (op_type == OperationType::ADD) {
        for (std::size_t i = 0; i < size_; ++i) {
            float a = fp16_to_fp32(this_data[i]);
            float b = fp16_to_fp32(other_data[i]);
            this_data[i] = fp32_to_fp16(a + b);
        }
        return;
    }
    if (op_type == OperationType::SUB) {
        for (std::size_t i = 0; i < size_; ++i) {
            float a = fp16_to_fp32(this_data[i]);
            float b = fp16_to_fp32(other_data[i]);
            this_data[i] = fp32_to_fp16(a - b);
        }
        return;
    }
    if (op_type == OperationType::MUL) {
        for (std::size_t i = 0; i < size_; ++i) {
            float a = fp16_to_fp32(this_data[i]);
            float b = fp16_to_fp32(other_data[i]);
            this_data[i] = fp32_to_fp16(a * b);
        }
        return;
    }
    if (op_type == OperationType::DIV) {
        for (std::size_t i = 0; i < size_; ++i) {
            float a = fp16_to_fp32(this_data[i]);
            float b = fp16_to_fp32(other_data[i]);
            this_data[i] = fp32_to_fp16(a / b);
        }
        return;
    }
}

} // namespace labneura
