/**
 * AVX2 + FMA convolution kernels.
 *
 * AVX2 provides 256-bit vectors (8 x float32) with FMA (Fused Multiply-Add).
 * FMA reduces latency and improves throughput for multiply-accumulate operations.
 * This is the most common SIMD level on modern x86 CPUs (Intel Haswell+ / AMD Excavator+).
 *
 * Compile with: -mavx2 -mfma
 */

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

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

/**
 * Convolution for standard (time, channels) layout using AVX2+FMA.
 */
void convolve_avx2(
    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 8 channels at a time with AVX2 */
        unsigned int ch = 0;
        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);  /* FMA: acc += kval * inp */
            }

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

        /* Process remaining 4 channels with SSE+FMA */
        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 code */
        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 AVX2+FMA.
 */
void convolve_transposed_avx2(
    float * __restrict__ output,
    const float * __restrict__ kernel,
    const float * __restrict__ input,
    unsigned int kernel_len,
    unsigned int n_out)
{
    /* Process 8 output samples at a time with AVX2 */
    unsigned int s = 0;
    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+FMA */
    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 with scalar code */
    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 AVX2+FMA.
 *
 * 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_avx2(
    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 8 output samples at a time with AVX2, 4 filters */
            unsigned int s = 0;
            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);
            }

            /* Process 4 output samples with SSE+FMA */
            for (; s + 4 <= tile_size; s += 4) {
                __m128 acc0 = _mm_setzero_ps();
                __m128 acc1 = _mm_setzero_ps();
                __m128 acc2 = _mm_setzero_ps();
                __m128 acc3 = _mm_setzero_ps();

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

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

                _mm_storeu_ps(&out0[s], acc0);
                _mm_storeu_ps(&out1[s], acc1);
                _mm_storeu_ps(&out2[s], acc2);
                _mm_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 + 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 + 4 <= tile_size; s += 4) {
                __m128 acc = _mm_setzero_ps();
                for (unsigned int k = 0; k < filter_len; k++) {
                    acc = _mm_fmadd_ps(_mm_set1_ps(filter[k]),
                                       _mm_loadu_ps(&input[s_base + s + k]), acc);
                }
                _mm_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();
}

/**
 * Polyphase upsampling for transposed layout using AVX2+FMA.
 *
 * Processes all phases simultaneously per output position.
 * For factor=8, all phases fit in a single __m256 register.
 * For factor=4, uses __m128.
 * For other factors, falls back to scalar with FMA.
 *
 * kernel_T layout: [kernel_len][factor] (row-major, 64-byte aligned rows)
 */
void upsample_polyphase_avx2(
    float * __restrict__ output,
    const float * __restrict__ kernel_T,
    const float * __restrict__ input,
    unsigned int kernel_len,
    unsigned int factor,
    unsigned int n_valid)
{
    if (factor == 8) {
        /* Process 4 output positions at once for ILP (4 independent FMA chains)
         * Use streaming stores to avoid write-allocate on 16MB output */
        unsigned int s = 0;
        for (; s + 4 <= n_valid; s += 4) {
            __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 < kernel_len; k++) {
                __m256 kern = _mm256_loadu_ps(&kernel_T[k * 8]);
                acc0 = _mm256_fmadd_ps(_mm256_set1_ps(input[s + 0 + k]), kern, acc0);
                acc1 = _mm256_fmadd_ps(_mm256_set1_ps(input[s + 1 + k]), kern, acc1);
                acc2 = _mm256_fmadd_ps(_mm256_set1_ps(input[s + 2 + k]), kern, acc2);
                acc3 = _mm256_fmadd_ps(_mm256_set1_ps(input[s + 3 + k]), kern, acc3);
            }

            _mm256_storeu_ps(&output[s * 8 +  0], acc0);
            _mm256_storeu_ps(&output[s * 8 +  8], acc1);
            _mm256_storeu_ps(&output[s * 8 + 16], acc2);
            _mm256_storeu_ps(&output[s * 8 + 24], acc3);
        }

        /* Handle remaining positions */
        for (; s < n_valid; s++) {
            __m256 acc = _mm256_setzero_ps();
            for (unsigned int k = 0; k < kernel_len; k++) {
                acc = _mm256_fmadd_ps(_mm256_set1_ps(input[s + k]),
                                      _mm256_loadu_ps(&kernel_T[k * 8]), acc);
            }
            _mm256_storeu_ps(&output[s * 8], acc);
        }
    } else if (factor == 4) {
        /* Process 4 positions at once, each producing 4 outputs = 4 x __m128 */
        unsigned int s = 0;
        for (; s + 4 <= n_valid; s += 4) {
            __m128 acc0 = _mm_setzero_ps();
            __m128 acc1 = _mm_setzero_ps();
            __m128 acc2 = _mm_setzero_ps();
            __m128 acc3 = _mm_setzero_ps();

            for (unsigned int k = 0; k < kernel_len; k++) {
                __m128 kern = _mm_loadu_ps(&kernel_T[k * 4]);
                acc0 = _mm_fmadd_ps(_mm_set1_ps(input[s + 0 + k]), kern, acc0);
                acc1 = _mm_fmadd_ps(_mm_set1_ps(input[s + 1 + k]), kern, acc1);
                acc2 = _mm_fmadd_ps(_mm_set1_ps(input[s + 2 + k]), kern, acc2);
                acc3 = _mm_fmadd_ps(_mm_set1_ps(input[s + 3 + k]), kern, acc3);
            }

            _mm_storeu_ps(&output[s * 4 +  0], acc0);
            _mm_storeu_ps(&output[s * 4 +  4], acc1);
            _mm_storeu_ps(&output[s * 4 +  8], acc2);
            _mm_storeu_ps(&output[s * 4 + 12], acc3);
        }

        for (; s < n_valid; s++) {
            __m128 acc = _mm_setzero_ps();
            for (unsigned int k = 0; k < kernel_len; k++) {
                acc = _mm_fmadd_ps(_mm_set1_ps(input[s + k]),
                                    _mm_loadu_ps(&kernel_T[k * 4]), acc);
            }
            _mm_storeu_ps(&output[s * 4], acc);
        }
    } else if (factor == 2) {
        /* factor=2: process 8 positions at once per phase, then interleave.
         * Two __m256 accumulators (one per phase) process 8 contiguous input positions.
         * Inner loop: 1 load + 2 broadcasts + 2 FMAs per tap — no expensive _mm256_set_ps.
         * Output interleave: unpacklo/hi + permute2f128. */
        unsigned int s = 0;
        for (; s + 8 <= n_valid; s += 8) {
            __m256 acc_p0 = _mm256_setzero_ps();
            __m256 acc_p1 = _mm256_setzero_ps();

            for (unsigned int k = 0; k < kernel_len; k++) {
                __m256 in_val = _mm256_loadu_ps(&input[s + k]);
                __m256 k0 = _mm256_set1_ps(kernel_T[k * 2 + 0]);
                __m256 k1 = _mm256_set1_ps(kernel_T[k * 2 + 1]);
                acc_p0 = _mm256_fmadd_ps(in_val, k0, acc_p0);
                acc_p1 = _mm256_fmadd_ps(in_val, k1, acc_p1);
            }

            /* Interleave: acc_p0=[a0..a7], acc_p1=[b0..b7]
             * Want output: [a0,b0,a1,b1,a2,b2,a3,b3, a4,b4,a5,b5,a6,b6,a7,b7] */
            __m256 lo = _mm256_unpacklo_ps(acc_p0, acc_p1); /* [a0,b0,a1,b1, a4,b4,a5,b5] */
            __m256 hi = _mm256_unpackhi_ps(acc_p0, acc_p1); /* [a2,b2,a3,b3, a6,b6,a7,b7] */
            __m256 out0 = _mm256_permute2f128_ps(lo, hi, 0x20); /* [a0,b0,a1,b1,a2,b2,a3,b3] */
            __m256 out1 = _mm256_permute2f128_ps(lo, hi, 0x31); /* [a4,b4,a5,b5,a6,b6,a7,b7] */

            _mm256_storeu_ps(&output[s * 2 +  0], out0);
            _mm256_storeu_ps(&output[s * 2 +  8], out1);
        }

        /* Handle remaining positions (< 8) with SSE: 4 at a time */
        for (; s + 4 <= n_valid; s += 4) {
            __m128 acc_p0 = _mm_setzero_ps();
            __m128 acc_p1 = _mm_setzero_ps();

            for (unsigned int k = 0; k < kernel_len; k++) {
                __m128 in_val = _mm_loadu_ps(&input[s + k]);
                acc_p0 = _mm_fmadd_ps(in_val, _mm_set1_ps(kernel_T[k * 2 + 0]), acc_p0);
                acc_p1 = _mm_fmadd_ps(in_val, _mm_set1_ps(kernel_T[k * 2 + 1]), acc_p1);
            }

            __m128 lo = _mm_unpacklo_ps(acc_p0, acc_p1); /* [a0,b0,a1,b1] */
            __m128 hi = _mm_unpackhi_ps(acc_p0, acc_p1); /* [a2,b2,a3,b3] */
            _mm_storeu_ps(&output[s * 2 + 0], lo);
            _mm_storeu_ps(&output[s * 2 + 4], hi);
        }

        /* Scalar remainder */
        for (; s < n_valid; s++) {
            float acc0 = 0.0f, acc1 = 0.0f;
            for (unsigned int k = 0; k < kernel_len; k++) {
                float in_val = input[s + k];
                acc0 += in_val * kernel_T[k * 2 + 0];
                acc1 += in_val * kernel_T[k * 2 + 1];
            }
            output[s * 2 + 0] = acc0;
            output[s * 2 + 1] = acc1;
        }
    } else if (factor == 16) {
        /* 16 phases = 2 x __m256 */
        for (unsigned int s = 0; s < n_valid; s++) {
            __m256 acc_lo = _mm256_setzero_ps();
            __m256 acc_hi = _mm256_setzero_ps();

            for (unsigned int k = 0; k < kernel_len; k++) {
                __m256 in_val = _mm256_set1_ps(input[s + k]);
                __m256 kern_lo = _mm256_loadu_ps(&kernel_T[k * factor]);
                __m256 kern_hi = _mm256_loadu_ps(&kernel_T[k * factor + 8]);
                acc_lo = _mm256_fmadd_ps(in_val, kern_lo, acc_lo);
                acc_hi = _mm256_fmadd_ps(in_val, kern_hi, acc_hi);
            }

            _mm256_storeu_ps(&output[s * factor], acc_lo);
            _mm256_storeu_ps(&output[s * factor + 8], acc_hi);
        }
    } else {
        /* General fallback: process factor values with AVX2 in chunks of 8 */
        for (unsigned int s = 0; s < n_valid; s++) {
            float *out_ptr = output + s * factor;

            /* Zero output with AVX2 */
            unsigned int p = 0;
            for (; p + 8 <= factor; p += 8) {
                _mm256_storeu_ps(&out_ptr[p], _mm256_setzero_ps());
            }
            for (; p < factor; p++) {
                out_ptr[p] = 0.0f;
            }

            for (unsigned int k = 0; k < kernel_len; k++) {
                __m256 in_val = _mm256_set1_ps(input[s + k]);
                const float *kern_row = kernel_T + k * factor;

                p = 0;
                for (; p + 8 <= factor; p += 8) {
                    __m256 acc = _mm256_loadu_ps(&out_ptr[p]);
                    __m256 kern = _mm256_loadu_ps(&kern_row[p]);
                    acc = _mm256_fmadd_ps(in_val, kern, acc);
                    _mm256_storeu_ps(&out_ptr[p], acc);
                }
                for (; p < factor; p++) {
                    out_ptr[p] += input[s + k] * kern_row[p];
                }
            }
        }
    }

    _mm256_zeroupper();
}

#endif /* x86 */
