#include "audio_buffer.h"

#include <cstring>
#include "audio_buffer_pool.h"
#include "misc/cpu_simd.h"
#include "misc/likely.h"

#if CPU_AMD64
#include <experimental/simd>
using namespace std::experimental;
#endif

void(*audio_buffer::add_func)(audio_buffer& in_buffer, audio_buffer& from_buffer, float percent);
void(*audio_buffer::multiple_func)(audio_buffer& in_buffer, float percent);

#if CPU_AMD64
#define DEFINE_SIMD_FUNC(simd_max) \
    constexpr size_t simd_size      = (simd_max) / sizeof(sample_t) / 8; \
    add_func = &add_simd<simd_size>; \
    multiple_func = &multiple_simd<simd_size>;

template<int simd_size>
void add_simd(audio_buffer& in_buffer, audio_buffer& from_buffer, float percent) {
    using simd_type = simd_abi::fixed_size<simd_size>;

    simd<sample_t, simd_type> percent_simd(percent);
    for (uint32_t channel_index = 0; channel_index < in_buffer.get_num_channels(); channel_index++) {
        sample_t* channel    = in_buffer.get_headers()[channel_index];
        sample_t* in_channel = from_buffer.get_headers()[channel_index];
        int i = 0;
        for (; i < in_buffer.get_num_samples(); i += simd_size) {
            simd<sample_t, simd_type> a(channel, element_aligned);
            simd<sample_t, simd_type> b(in_channel, element_aligned);
            a += b * percent_simd;
            a.copy_to(channel, element_aligned);

            channel += simd_size;
            in_channel += simd_size;
        }
        // if the number of samples is not a multiple of simd_size
        for (; i < in_buffer.get_num_samples(); ++i) {
            channel[i] += in_channel[i] * percent;
        }
    }
}

template<int simd_size>
void multiple_simd(audio_buffer& in_buffer, float percent) {
    using simd_type = simd_abi::fixed_size<simd_size>;

    simd<sample_t, simd_type> percent_simd(percent);
    for (auto channel : in_buffer.get_headers_vector()) {
        int i = 0;
        for (; i < in_buffer.get_num_samples(); i += simd_size) {
            simd<sample_t, simd_type> a(channel, element_aligned);
            a *= percent_simd;
            a.copy_to(channel, element_aligned);
            channel += simd_size;
        }
        // if the number of samples is not a multiple of simd_size
        for (; i < in_buffer.get_num_samples(); ++i) {
            channel[i] *= percent;
        }
    }
}
#endif

void add_no_simd(audio_buffer& in_buffer, audio_buffer& from_buffer, float percent) {
    for (uint32_t channel_index = 0; channel_index < in_buffer.get_num_channels(); channel_index++) {
        sample_t* channel    = in_buffer.get_headers()[channel_index];
        sample_t* in_channel = from_buffer.get_headers()[channel_index];
        for (int i = 0; i < in_buffer.get_num_samples(); ++i) {
            channel[i] += in_channel[i] * percent;
        }
    }
}

void multiple_no_simd(audio_buffer& in_buffer, float percent) {
    for (auto channel : in_buffer.get_headers_vector()) {
        for (int i = 0; i < in_buffer.get_num_samples(); ++i) {
            channel[i] *= percent;
        }
    }
}

audio_buffer::audio_buffer() {
    static bool func_initialized = false;
    if (UNLIKELY(!func_initialized)) {
		cpuid const cpu;
#if CPU_AMD64
        if (cpu.support_avx512()) {
            DEFINE_SIMD_FUNC(512)
        } else if (cpu.support_avx() || cpu.support_avx2()) {
            DEFINE_SIMD_FUNC(256)
        } else if (cpu.support_sse()) {
            DEFINE_SIMD_FUNC(128)
        }
#elif CPU_ARM
//        if (cpu.support_neon()) {
//            DEFINE_SIMD_FUNC(128)
//        }
#endif
        if (!add_func) {
            add_func = &add_no_simd;
            multiple_func = &multiple_no_simd;
        }
        func_initialized = true;
    }
}

audio_buffer::~audio_buffer() {
    free();
}

void audio_buffer::resize(uint32_t channel_num, uint32_t frame_size) {
    if (channel_num == headers_.size() && frame_size == frame_size_)
        return;
    frame_size_ = frame_size;
    free();
    audio_buffer_pool* pool = get_audio_buffer_pool();
    for (int i = 0; i < channel_num; ++i) {
        sample_t* block = pool->alloc(frame_size);
        headers_.push_back(block);
    }
    clear();
}

void audio_buffer::clear() {
    for (sample_t* channel : headers_) {
        std::memset(channel, 0, frame_size_ * sizeof(sample_t));
    }
}


void audio_buffer::add(audio_buffer& from_buffer, float percent) {
    std::scoped_lock lock(lock_);
    add_func(*this, from_buffer, percent);
}

void audio_buffer::multiple(float percent) {
    std::scoped_lock lock(lock_);
    multiple_func(*this, percent);
}

std::vector<sample_t> audio_buffer::get_interleaved_buffer() const {
    std::vector<sample_t> result;
    result.reserve(headers_.size() * frame_size_);
    for (int i = 0; i < frame_size_; ++i) {
        for (const sample_t* channel : headers_) {
            result.push_back(channel[i]);
        }
    }
    return result;
}

void audio_buffer::free() {
    for (sample_t* header : headers_)
        get_audio_buffer_pool()->free(header, frame_size_);
    headers_.clear();
}

#undef DEFINE_SIMD_FUNC