#include <vector>
#include <complex>
#include <cmath>
#include <algorithm>
#include <cassert>
#include <fftw3.h>

std::vector<std::vector<double>> ModS2(const std::vector<std::vector<double>>& image, double dt, double du, double ERMv, double tau0, const std::vector<std::vector<double>>& ModTransd, int R_x = 1) {
    // Check R_x argument
    R_x = std::max(1, R_x);
    assert(R_x > 0 && "R_x must be positive");

    // Zero-padding
    int nt0 = image.size();
    int nu0 = image[0].size();
    int nt = ModTransd.size();
    int nu = ModTransd[0].size();
    if (nt < nt0) {
        nt = std::pow(2, std::ceil(std::log2(nt0)) + 1);
    }
    if (nu < nu0) {
        nu = 2 * nu0;
    }

    // Data and image grids
    std::vector<std::vector<double>> f(nt, std::vector<double>(nu));
    std::vector<std::vector<double>> ku(nt, std::vector<double>(nu));
    std::vector<std::vector<double>> fkz(nt, std::vector<double>(nu));

    for (int i = 0; i < nt; ++i) {
        for (int j = 0; j < nu; ++j) {
            f[i][j] = (i - nt/2) / (dt * nt);
            ku[i][j] = (j - nu/2) / (du * nu);
            fkz[i][j] = ERMv * std::copysign(1.0, f[i][j]) * std::sqrt(std::pow(ku[i][j], 2) + std::pow(f[i][j] / ERMv, 2));
        }
    }

    // Converting image to frequency domain
    fftw_complex *in, *out;
    fftw_plan p;
    in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * nt * nu);
    out = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * nt * nu);

    for (int i = 0; i < nt0; ++i) {
        for (int j = 0; j < nu0; ++j) {
            in[i*nu + j][0] = image[i][j];
            in[i*nu + j][1] = 0;
        }
    }

    p = fftw_plan_dft_2d(nt, nu, in, out, FFTW_FORWARD, FFTW_ESTIMATE);
    fftw_execute(p);

    std::vector<std::vector<std::complex<double>>> ftimage(nt, std::vector<std::complex<double>>(nu));
    for (int i = 0; i < nt; ++i) {
        for (int j = 0; j < nu; ++j) {
            ftimage[i][j] = std::complex<double>(out[i*nu + j][0], out[i*nu + j][1]);
        }
    }

    // Generates data spectrum (by adjoint of linear interpolation)
    // Note: linterpAdj function is not provided, so it's omitted here
    std::vector<std::vector<std::complex<double>>> ftdata = ftimage;  // Placeholder

    if (tau0 != 0) {
        for (int i = 0; i < nt; ++i) {
            for (int j = 0; j < nu; ++j) {
                ftdata[i][j] *= std::exp(std::complex<double>(0, 2 * M_PI * (fkz[i][j] - f[i][j]) * tau0));
            }
        }
    }

    // Inserts transducer model (if exists)
    if (!ModTransd.empty()) {
        int rowShift = (nt - ModTransd.size()) / 2;
        int colShift = (nu - ModTransd[0].size()) / 2;
        for (int i = 0; i < nt; ++i) {
            for (int j = 0; j < nu; ++j) {
                int ii = (i + rowShift) % nt;
                int jj = (j + colShift) % nu;
                if (ii < ModTransd.size() && jj < ModTransd[0].size()) {
                    ftdata[i][j] *= ModTransd[ii][jj];
                } else {
                    ftdata[i][j] = 0;
                }
            }
        }
    }

    // Converting data to time domain
    for (int i = 0; i < nt; ++i) {
        for (int j = 0; j < nu; ++j) {
            in[i*nu + j][0] = ftdata[i][j].real();
            in[i*nu + j][1] = ftdata[i][j].imag();
        }
    }

    p = fftw_plan_dft_2d(nt, nu, in, out, FFTW_BACKWARD, FFTW_ESTIMATE);
    fftw_execute(p);

    std::vector<std::vector<double>> data(nt0, std::vector<double>(nu0));
    for (int i = 0; i < nt0; ++i) {
        for (int j = 0; j < nu0; ++j) {
            data[i][j] = out[i*nu + j][0] / (nt * nu);
        }
    }

    // Applies sub-sampling
    if (R_x > 1) {
        for (int i = 0; i < nt0; ++i) {
            for (int j = 0; j < nu0; j += R_x) {
                data[i][j/R_x] = data[i][j] * R_x;
            }
        }
        data[0].resize(nu0/R_x);
    }

    fftw_destroy_plan(p);
    fftw_free(in);
    fftw_free(out);

    return data;
}