/**
 * SSE2 convolution kernels.
 * Baseline SIMD for x86_64 (SSE2 is mandatory on x86_64).
 *
 * SSE2 provides 128-bit vectors (4 x float32).
 * Compile with: -msse2
 */

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

#include <emmintrin.h>  /* SSE2 */
#include "../dispatch.h"

/**
 * Convolution for standard (time, channels) layout using SSE2.
 */
void convolve_sse2(
    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 4 channels at a time with SSE2 */
        unsigned int ch = 0;
        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]);
                __m128 prod = _mm_mul_ps(kval, inp);
                acc = _mm_add_ps(acc, prod);
            }

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

        /* Handle remaining channels */
        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;
        }
    }
}

/**
 * Convolution for transposed (channels, time) layout using SSE2.
 */
void convolve_transposed_sse2(
    float * __restrict__ output,
    const float * __restrict__ kernel,
    const float * __restrict__ input,
    unsigned int kernel_len,
    unsigned int n_out)
{
    unsigned int s = 0;
    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]);
            __m128 prod = _mm_mul_ps(kval, inp);
            acc = _mm_add_ps(acc, prod);
        }

        _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;
    }
}

/**
 * Multi-filter convolution for groups=1 case using SSE2.
 *
 * 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_sse2(
    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 4 output samples at a time with SSE2, 4 filters */
            unsigned int s = 0;
            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_add_ps(acc0, _mm_mul_ps(_mm_set1_ps(filt0[k]), in_vec));
                    acc1 = _mm_add_ps(acc1, _mm_mul_ps(_mm_set1_ps(filt1[k]), in_vec));
                    acc2 = _mm_add_ps(acc2, _mm_mul_ps(_mm_set1_ps(filt2[k]), in_vec));
                    acc3 = _mm_add_ps(acc3, _mm_mul_ps(_mm_set1_ps(filt3[k]), in_vec));
                }

                _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 + 4 <= tile_size; s += 4) {
                __m128 acc = _mm_setzero_ps();
                for (unsigned int k = 0; k < filter_len; k++) {
                    __m128 kval = _mm_set1_ps(filter[k]);
                    __m128 inp = _mm_loadu_ps(&input[s_base + s + k]);
                    acc = _mm_add_ps(acc, _mm_mul_ps(kval, inp));
                }
                _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;
            }
        }
    }
}

/**
 * Polyphase upsampling for transposed layout using SSE2.
 * 128-bit (4 floats), no FMA.
 */
void upsample_polyphase_sse2(
    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) {
        /* factor=8: 8 phases = 2 x __m128 per position, 4 positions for ILP */
        unsigned int s = 0;
        for (; s + 4 <= n_valid; s += 4) {
            __m128 a0_lo = _mm_setzero_ps(), a0_hi = _mm_setzero_ps();
            __m128 a1_lo = _mm_setzero_ps(), a1_hi = _mm_setzero_ps();
            __m128 a2_lo = _mm_setzero_ps(), a2_hi = _mm_setzero_ps();
            __m128 a3_lo = _mm_setzero_ps(), a3_hi = _mm_setzero_ps();

            for (unsigned int k = 0; k < kernel_len; k++) {
                __m128 kern_lo = _mm_loadu_ps(&kernel_T[k * 8]);
                __m128 kern_hi = _mm_loadu_ps(&kernel_T[k * 8 + 4]);
                __m128 i0 = _mm_set1_ps(input[s + 0 + k]);
                __m128 i1 = _mm_set1_ps(input[s + 1 + k]);
                __m128 i2 = _mm_set1_ps(input[s + 2 + k]);
                __m128 i3 = _mm_set1_ps(input[s + 3 + k]);
                a0_lo = _mm_add_ps(a0_lo, _mm_mul_ps(i0, kern_lo));
                a0_hi = _mm_add_ps(a0_hi, _mm_mul_ps(i0, kern_hi));
                a1_lo = _mm_add_ps(a1_lo, _mm_mul_ps(i1, kern_lo));
                a1_hi = _mm_add_ps(a1_hi, _mm_mul_ps(i1, kern_hi));
                a2_lo = _mm_add_ps(a2_lo, _mm_mul_ps(i2, kern_lo));
                a2_hi = _mm_add_ps(a2_hi, _mm_mul_ps(i2, kern_hi));
                a3_lo = _mm_add_ps(a3_lo, _mm_mul_ps(i3, kern_lo));
                a3_hi = _mm_add_ps(a3_hi, _mm_mul_ps(i3, kern_hi));
            }

            _mm_storeu_ps(&output[s * 8 +  0], a0_lo);
            _mm_storeu_ps(&output[s * 8 +  4], a0_hi);
            _mm_storeu_ps(&output[s * 8 +  8], a1_lo);
            _mm_storeu_ps(&output[s * 8 + 12], a1_hi);
            _mm_storeu_ps(&output[s * 8 + 16], a2_lo);
            _mm_storeu_ps(&output[s * 8 + 20], a2_hi);
            _mm_storeu_ps(&output[s * 8 + 24], a3_lo);
            _mm_storeu_ps(&output[s * 8 + 28], a3_hi);
        }

        for (; s < n_valid; s++) {
            __m128 acc_lo = _mm_setzero_ps(), acc_hi = _mm_setzero_ps();
            for (unsigned int k = 0; k < kernel_len; k++) {
                __m128 iv = _mm_set1_ps(input[s + k]);
                acc_lo = _mm_add_ps(acc_lo, _mm_mul_ps(iv, _mm_loadu_ps(&kernel_T[k * 8])));
                acc_hi = _mm_add_ps(acc_hi, _mm_mul_ps(iv, _mm_loadu_ps(&kernel_T[k * 8 + 4])));
            }
            _mm_storeu_ps(&output[s * 8], acc_lo);
            _mm_storeu_ps(&output[s * 8 + 4], acc_hi);
        }
    } else if (factor == 2) {
        /* factor=2: process 4 positions at once per phase, then interleave */
        unsigned int s = 0;
        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]);
                __m128 k0 = _mm_set1_ps(kernel_T[k * 2 + 0]);
                __m128 k1 = _mm_set1_ps(kernel_T[k * 2 + 1]);
                acc_p0 = _mm_add_ps(acc_p0, _mm_mul_ps(in_val, k0));
                acc_p1 = _mm_add_ps(acc_p1, _mm_mul_ps(in_val, k1));
            }

            __m128 lo = _mm_unpacklo_ps(acc_p0, acc_p1);
            __m128 hi = _mm_unpackhi_ps(acc_p0, acc_p1);
            _mm_storeu_ps(&output[s * 2 + 0], lo);
            _mm_storeu_ps(&output[s * 2 + 4], hi);
        }

        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 == 4) {
        /* factor=4: 4 phases = 1 x __m128 per position */
        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_add_ps(acc0, _mm_mul_ps(_mm_set1_ps(input[s + 0 + k]), kern));
                acc1 = _mm_add_ps(acc1, _mm_mul_ps(_mm_set1_ps(input[s + 1 + k]), kern));
                acc2 = _mm_add_ps(acc2, _mm_mul_ps(_mm_set1_ps(input[s + 2 + k]), kern));
                acc3 = _mm_add_ps(acc3, _mm_mul_ps(_mm_set1_ps(input[s + 3 + k]), kern));
            }

            _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_add_ps(acc, _mm_mul_ps(_mm_set1_ps(input[s + k]),
                                                  _mm_loadu_ps(&kernel_T[k * 4])));
            }
            _mm_storeu_ps(&output[s * 4], acc);
        }
    } else {
        /* General fallback */
        for (unsigned int s = 0; s < n_valid; s++) {
            float *out_ptr = output + s * factor;
            unsigned int p = 0;
            for (; p + 4 <= factor; p += 4)
                _mm_storeu_ps(&out_ptr[p], _mm_setzero_ps());
            for (; p < factor; p++)
                out_ptr[p] = 0.0f;
            for (unsigned int k = 0; k < kernel_len; k++) {
                __m128 iv = _mm_set1_ps(input[s + k]);
                const float *kr = kernel_T + k * factor;
                p = 0;
                for (; p + 4 <= factor; p += 4) {
                    __m128 a = _mm_loadu_ps(&out_ptr[p]);
                    a = _mm_add_ps(a, _mm_mul_ps(iv, _mm_loadu_ps(&kr[p])));
                    _mm_storeu_ps(&out_ptr[p], a);
                }
                for (; p < factor; p++)
                    out_ptr[p] += input[s + k] * kr[p];
            }
        }
    }
}

#endif /* x86 */
