#include <riscv_vector.h>
#include <algorithm>

void MNNC3ToYUV(const unsigned char* source, unsigned char* dest, 
                    size_t count, bool bgr, bool yuv) {
    static const int coeffs[] = {
        // Y
         4899,    9617,    1868,
        // Cr
         8192,   -6860,   -1332,
        // Cb
        -2765,   -5427,    8192,
        // U
        -2412,   -4734,    7146,
        // V
        10076,  -8438,   -1638
    };
    
    int r0 = 0, r1 = 3, r2 = 6,
        g0 = 1, g1 = 4, g2 = 7,
        b0 = 2, b1 = 5, b2 = 8;
    if (yuv) {
        r1 = 9,  r2 = 12;
        g1 = 10, g2 = 13;
        b1 = 11, b2 = 14;
    }
    if (bgr) {
        std::swap(r0, b0);
        std::swap(r1, b1);
        std::swap(r2, b2);
    }
    
    int16_t C0 = coeffs[r0], C1 = coeffs[g0], C2 = coeffs[b0],
            C3 = coeffs[r1], C4 = coeffs[g1], C5 = coeffs[b1],
            C6 = coeffs[r2], C7 = coeffs[g2], C8 = coeffs[b2];
    
    size_t i = 0;
    const int32_t rounding = 1 << 13;
    
    while (i < count) {
        size_t vl = __riscv_vsetvl_e8m2(count - i);
        vuint8m2_t vrU8 = __riscv_vlse8_v_u8m2(source + 3 * i + 0, 3, vl);
        vuint8m2_t vgU8 = __riscv_vlse8_v_u8m2(source + 3 * i + 1, 3, vl);
        vuint8m2_t vbU8 = __riscv_vlse8_v_u8m2(source + 3 * i + 2, 3, vl);
        
        vint16m4_t vr = __riscv_vreinterpret_v_u16m4_i16m4(__riscv_vzext_vf2_u16m4(vrU8, vl));
        vint16m4_t vg = __riscv_vreinterpret_v_u16m4_i16m4(__riscv_vzext_vf2_u16m4(vgU8, vl));
        vint16m4_t vb = __riscv_vreinterpret_v_u16m4_i16m4(__riscv_vzext_vf2_u16m4(vbU8, vl));
        
        vint32m8_t sum = __riscv_vwmul_vx_i32m8(vr, C0, vl);
        sum = __riscv_vwmacc_vx_i32m8(sum, C1, vg, vl);
        sum = __riscv_vwmacc_vx_i32m8(sum, C2, vb, vl);
        sum = __riscv_vadd_vx_i32m8(sum, rounding, vl);
        sum = __riscv_vsra_vx_i32m8(sum, 14, vl);
        sum = __riscv_vmax_vx_i32m8(sum, 0, vl);
        sum = __riscv_vmin_vx_i32m8(sum, 255, vl);
        vint16m4_t sum16 = __riscv_vnsra_wx_i16m4(sum, 0, vl);
        vuint8m2_t result = __riscv_vnsrl_wx_u8m2(__riscv_vreinterpret_v_i16m4_u16m4(sum16), 0, vl);
        __riscv_vsse8_v_u8m2(dest + 3 * i + 0, 3, result, vl);
        
        sum = __riscv_vwmul_vx_i32m8(vr, C3, vl);
        sum = __riscv_vwmacc_vx_i32m8(sum, C4, vg, vl);
        sum = __riscv_vwmacc_vx_i32m8(sum, C5, vb, vl);
        sum = __riscv_vadd_vx_i32m8(sum, rounding, vl);
        sum = __riscv_vsra_vx_i32m8(sum, 14, vl);
        sum = __riscv_vadd_vx_i32m8(sum, 128, vl);
        sum = __riscv_vmax_vx_i32m8(sum, 0, vl);
        sum = __riscv_vmin_vx_i32m8(sum, 255, vl);
        sum16 = __riscv_vnsra_wx_i16m4(sum, 0, vl);
        result = __riscv_vnsrl_wx_u8m2(__riscv_vreinterpret_v_i16m4_u16m4(sum16), 0, vl);
        __riscv_vsse8_v_u8m2(dest + 3 * i + 1, 3, result, vl);
        
        sum = __riscv_vwmul_vx_i32m8(vr, C6, vl);
        sum = __riscv_vwmacc_vx_i32m8(sum, C7, vg, vl);
        sum = __riscv_vwmacc_vx_i32m8(sum, C8, vb, vl);
        sum = __riscv_vadd_vx_i32m8(sum, rounding, vl);
        sum = __riscv_vsra_vx_i32m8(sum, 14, vl);
        sum = __riscv_vadd_vx_i32m8(sum, 128, vl);
        sum = __riscv_vmax_vx_i32m8(sum, 0, vl);
        sum = __riscv_vmin_vx_i32m8(sum, 255, vl);
        sum16 = __riscv_vnsra_wx_i16m4(sum, 0, vl);
        result = __riscv_vnsrl_wx_u8m2(__riscv_vreinterpret_v_i16m4_u16m4(sum16), 0, vl);
        __riscv_vsse8_v_u8m2(dest + 3 * i + 2, 3, result, vl);
        
        i += vl;
    }
}