/**
 * AVX-512 convolution kernels.
 *
 * AVX-512 provides 512-bit vectors (16 x float32) with mask registers for
 * efficient handling of non-aligned data.
 *
 * Compile with: -mavx512f -mfma
 */

#if defined(__x86_64__) || defined(_M_X64)

#include <immintrin.h>
#include "../dispatch.h"

/**
 * Convolution for standard (time, channels) layout using AVX-512.
 */
void convolve_avx512(
    float * __restrict__ output,
    const float * __restrict__ kernel,
    const float * __restrict__ input,
    unsigned int kernel_len,
    unsigned int channels,
    unsigned int n_out)
{
    for (unsigned int s = 0; s < n_out; s++) {
        float *out_row = output + s * channels;

        /* Process 16 channels at a time with AVX-512 */
        unsigned int ch = 0;
        for (; ch + 16 <= channels; ch += 16) {
            __m512 acc = _mm512_setzero_ps();

            for (unsigned int k = 0; k < kernel_len; k++) {
                __m512 kval = _mm512_set1_ps(kernel[k]);
                __m512 inp = _mm512_loadu_ps(&input[(s + k) * channels + ch]);
                acc = _mm512_fmadd_ps(kval, inp, acc);
            }

            _mm512_storeu_ps(&out_row[ch], acc);
        }

        /* Process remaining 8 channels with AVX2 */
        for (; ch + 8 <= channels; ch += 8) {
            __m256 acc = _mm256_setzero_ps();

            for (unsigned int k = 0; k < kernel_len; k++) {
                __m256 kval = _mm256_set1_ps(kernel[k]);
                __m256 inp = _mm256_loadu_ps(&input[(s + k) * channels + ch]);
                acc = _mm256_fmadd_ps(kval, inp, acc);
            }

            _mm256_storeu_ps(&out_row[ch], acc);
        }

        /* Process remaining 4 channels with SSE */
        for (; ch + 4 <= channels; ch += 4) {
            __m128 acc = _mm_setzero_ps();

            for (unsigned int k = 0; k < kernel_len; k++) {
                __m128 kval = _mm_set1_ps(kernel[k]);
                __m128 inp = _mm_loadu_ps(&input[(s + k) * channels + ch]);
                acc = _mm_fmadd_ps(kval, inp, acc);
            }

            _mm_storeu_ps(&out_row[ch], acc);
        }

        /* Handle remaining channels with scalar */
        for (; ch < channels; ch++) {
            float sum = 0.0f;
            for (unsigned int k = 0; k < kernel_len; k++) {
                sum += kernel[k] * input[(s + k) * channels + ch];
            }
            out_row[ch] = sum;
        }
    }

    _mm256_zeroupper();
}

/**
 * Convolution for transposed (channels, time) layout using AVX-512.
 */
void convolve_transposed_avx512(
    float * __restrict__ output,
    const float * __restrict__ kernel,
    const float * __restrict__ input,
    unsigned int kernel_len,
    unsigned int n_out)
{
    /* Process 16 output samples at a time with AVX-512 */
    unsigned int s = 0;
    for (; s + 16 <= n_out; s += 16) {
        __m512 acc = _mm512_setzero_ps();

        for (unsigned int k = 0; k < kernel_len; k++) {
            __m512 kval = _mm512_set1_ps(kernel[k]);
            __m512 inp = _mm512_loadu_ps(&input[s + k]);
            acc = _mm512_fmadd_ps(kval, inp, acc);
        }

        _mm512_storeu_ps(&output[s], acc);
    }

    /* Process remaining 8 samples with AVX2 */
    for (; s + 8 <= n_out; s += 8) {
        __m256 acc = _mm256_setzero_ps();

        for (unsigned int k = 0; k < kernel_len; k++) {
            __m256 kval = _mm256_set1_ps(kernel[k]);
            __m256 inp = _mm256_loadu_ps(&input[s + k]);
            acc = _mm256_fmadd_ps(kval, inp, acc);
        }

        _mm256_storeu_ps(&output[s], acc);
    }

    /* Process remaining 4 samples with SSE */
    for (; s + 4 <= n_out; s += 4) {
        __m128 acc = _mm_setzero_ps();

        for (unsigned int k = 0; k < kernel_len; k++) {
            __m128 kval = _mm_set1_ps(kernel[k]);
            __m128 inp = _mm_loadu_ps(&input[s + k]);
            acc = _mm_fmadd_ps(kval, inp, acc);
        }

        _mm_storeu_ps(&output[s], acc);
    }

    /* Handle remaining samples */
    for (; s < n_out; s++) {
        float sum = 0.0f;
        for (unsigned int k = 0; k < kernel_len; k++) {
            sum += kernel[k] * input[s + k];
        }
        output[s] = sum;
    }

    _mm256_zeroupper();
}

/**
 * Multi-filter convolution for groups=1 case using AVX-512.
 *
 * Applies multiple filters to a single input channel.
 * Strategy: Tiled approach - process blocks of filters and output samples together
 * to maximize data reuse in cache.
 *
 * filters: (n_filters, filter_len) - row-major
 * input: (n_samples,) - single channel
 * output: (n_filters, n_out) - row-major
 */
void convolve_multi_filter_avx512(
    float * __restrict__ output,
    const float * __restrict__ filters,
    const float * __restrict__ input,
    unsigned int n_filters,
    unsigned int filter_len,
    unsigned int n_out)
{
    /* Tile size for output samples - chosen for L1 cache efficiency */
    const unsigned int TILE_S = 256;

    /* Process tiles of output samples */
    for (unsigned int s_base = 0; s_base < n_out; s_base += TILE_S) {
        unsigned int s_end = (s_base + TILE_S < n_out) ? (s_base + TILE_S) : n_out;
        unsigned int tile_size = s_end - s_base;

        /* Process 4 filters at a time for this output tile */
        unsigned int f = 0;
        for (; f + 4 <= n_filters; f += 4) {
            const float *filt0 = &filters[(f + 0) * filter_len];
            const float *filt1 = &filters[(f + 1) * filter_len];
            const float *filt2 = &filters[(f + 2) * filter_len];
            const float *filt3 = &filters[(f + 3) * filter_len];

            float *out0 = &output[(f + 0) * n_out + s_base];
            float *out1 = &output[(f + 1) * n_out + s_base];
            float *out2 = &output[(f + 2) * n_out + s_base];
            float *out3 = &output[(f + 3) * n_out + s_base];

            /* Process 16 output samples at a time with AVX-512, 4 filters */
            unsigned int s = 0;
            for (; s + 16 <= tile_size; s += 16) {
                __m512 acc0 = _mm512_setzero_ps();
                __m512 acc1 = _mm512_setzero_ps();
                __m512 acc2 = _mm512_setzero_ps();
                __m512 acc3 = _mm512_setzero_ps();

                for (unsigned int k = 0; k < filter_len; k++) {
                    __m512 in_vec = _mm512_loadu_ps(&input[s_base + s + k]);

                    acc0 = _mm512_fmadd_ps(_mm512_set1_ps(filt0[k]), in_vec, acc0);
                    acc1 = _mm512_fmadd_ps(_mm512_set1_ps(filt1[k]), in_vec, acc1);
                    acc2 = _mm512_fmadd_ps(_mm512_set1_ps(filt2[k]), in_vec, acc2);
                    acc3 = _mm512_fmadd_ps(_mm512_set1_ps(filt3[k]), in_vec, acc3);
                }

                _mm512_storeu_ps(&out0[s], acc0);
                _mm512_storeu_ps(&out1[s], acc1);
                _mm512_storeu_ps(&out2[s], acc2);
                _mm512_storeu_ps(&out3[s], acc3);
            }

            /* Process 8 output samples with AVX2 */
            for (; s + 8 <= tile_size; s += 8) {
                __m256 acc0 = _mm256_setzero_ps();
                __m256 acc1 = _mm256_setzero_ps();
                __m256 acc2 = _mm256_setzero_ps();
                __m256 acc3 = _mm256_setzero_ps();

                for (unsigned int k = 0; k < filter_len; k++) {
                    __m256 in_vec = _mm256_loadu_ps(&input[s_base + s + k]);

                    acc0 = _mm256_fmadd_ps(_mm256_set1_ps(filt0[k]), in_vec, acc0);
                    acc1 = _mm256_fmadd_ps(_mm256_set1_ps(filt1[k]), in_vec, acc1);
                    acc2 = _mm256_fmadd_ps(_mm256_set1_ps(filt2[k]), in_vec, acc2);
                    acc3 = _mm256_fmadd_ps(_mm256_set1_ps(filt3[k]), in_vec, acc3);
                }

                _mm256_storeu_ps(&out0[s], acc0);
                _mm256_storeu_ps(&out1[s], acc1);
                _mm256_storeu_ps(&out2[s], acc2);
                _mm256_storeu_ps(&out3[s], acc3);
            }

            /* Handle remaining samples (scalar) */
            for (; s < tile_size; s++) {
                float sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f, sum3 = 0.0f;
                for (unsigned int k = 0; k < filter_len; k++) {
                    float in_val = input[s_base + s + k];
                    sum0 += filt0[k] * in_val;
                    sum1 += filt1[k] * in_val;
                    sum2 += filt2[k] * in_val;
                    sum3 += filt3[k] * in_val;
                }
                out0[s] = sum0;
                out1[s] = sum1;
                out2[s] = sum2;
                out3[s] = sum3;
            }
        }

        /* Handle remaining filters (not multiple of 4) */
        for (; f < n_filters; f++) {
            const float *filter = &filters[f * filter_len];
            float *out_row = &output[f * n_out + s_base];

            unsigned int s = 0;
            for (; s + 16 <= tile_size; s += 16) {
                __m512 acc = _mm512_setzero_ps();
                for (unsigned int k = 0; k < filter_len; k++) {
                    acc = _mm512_fmadd_ps(_mm512_set1_ps(filter[k]),
                                          _mm512_loadu_ps(&input[s_base + s + k]), acc);
                }
                _mm512_storeu_ps(&out_row[s], acc);
            }

            for (; s + 8 <= tile_size; s += 8) {
                __m256 acc = _mm256_setzero_ps();
                for (unsigned int k = 0; k < filter_len; k++) {
                    acc = _mm256_fmadd_ps(_mm256_set1_ps(filter[k]),
                                          _mm256_loadu_ps(&input[s_base + s + k]), acc);
                }
                _mm256_storeu_ps(&out_row[s], acc);
            }

            for (; s < tile_size; s++) {
                float sum = 0.0f;
                for (unsigned int k = 0; k < filter_len; k++) {
                    sum += filter[k] * input[s_base + s + k];
                }
                out_row[s] = sum;
            }
        }
    }

    _mm256_zeroupper();
}

#endif /* x86_64 */
