#ifndef ANACAL_FPFS_BASE_H
#define ANACAL_FPFS_BASE_H

#include "../image.h"
#include "../math.h"
#include "../mask.h"
#include "../stdafx.h"

#include <vector>

namespace anacal {
namespace fpfs {

inline constexpr double fpfs_det_sigma2 = 0.04;
inline constexpr double fpfs_cut_sigma_ratio = 1.6;
namespace detail {
    inline std::vector<double>
    make_rfftfreq(int n) {
        const int n_r = n / 2 + 1;
        std::vector<double> freq(n_r);
        const double factor = 2.0 * M_PI / static_cast<double>(n);
        for (int k = 0; k < n_r; ++k) {
            freq[k] = factor * static_cast<double>(k);
        }
        return freq;
    }

    inline std::vector<double>
    make_fftfreq(int n) {
        std::vector<double> freq(n);
        const double factor = 2.0 * M_PI / static_cast<double>(n);
        const int n_half = n / 2;
        for (int k = 0; k <= n_half; ++k) {
            freq[k] = factor * static_cast<double>(k);
        }
        for (int k = n_half + 1; k < n; ++k) {
            freq[k] = factor * static_cast<double>(k - n);
        }
        return freq;
    }

    inline std::complex<double>
    pow_i(int n) {
        switch (n % 4) {
            case 0: return {1.0, 0.0};
            case 1: return {0.0, 1.0};
            case 2: return {-1.0, 0.0};
            default: return {0.0, -1.0};
        }
    }

    inline size_t
    idx4(int n, int m, int y, int x, int norder, int mord, int ny, int nx) {
        return (
            (
                static_cast<size_t>(n) * (mord + 1)
                + static_cast<size_t>(m)
            ) * ny + static_cast<size_t>(y)
        ) * nx + static_cast<size_t>(x);
    }

    struct ShapeletResult {
        int ny;
        int nx;
        int norder;
        std::vector<std::complex<double>> chi;
    };

    inline ShapeletResult
    compute_shapelet_modes(int npix, int norder, double sigma, double kmax) {
        const int ny = npix;
        const int nx = npix / 2 + 1;
        const int mord = norder;
        const int nmodes = (norder + 1) * (mord + 1);
        const size_t total_size = static_cast<size_t>(nmodes) * ny * nx;

        std::vector<double> xfreq = make_rfftfreq(npix);
        std::vector<double> yfreq = make_fftfreq(npix);

        std::vector<double> gauss(static_cast<size_t>(ny) * nx, 0.0);
        std::vector<double> r2_over_sigma2(static_cast<size_t>(ny) * nx, 0.0);
        std::vector<std::complex<double>> eulfunc(static_cast<size_t>(ny) * nx);

        const double kmax2 = kmax * kmax;
        const double sigma2 = sigma * sigma;
        for (int j = 0; j < ny; ++j) {
            const double ky = yfreq[j];
            for (int i = 0; i < nx; ++i) {
                const double kx = xfreq[i];
                const double r2 = kx * kx + ky * ky;
                const bool inside = r2 <= kmax2;
                const size_t idx = static_cast<size_t>(j) * nx + static_cast<size_t>(i);
                if (inside) {
                    gauss[idx] = std::exp(-r2 / (2.0 * sigma2));
                } else {
                    gauss[idx] = 0.0;
                }
                if (sigma2 > 0.0) {
                    r2_over_sigma2[idx] = r2 / sigma2;
                } else {
                    r2_over_sigma2[idx] = 0.0;
                }
                if (r2 > 0.0) {
                    const double inv_r = 1.0 / std::sqrt(r2);
                    const double cos_phi = kx * inv_r;
                    const double sin_phi = ky * inv_r;
                    eulfunc[idx] = {cos_phi, sin_phi};
                } else {
                    eulfunc[idx] = {0.0, 0.0};
                }
            }
        }

        std::vector<double> lfunc(static_cast<size_t>(norder + 1) * (mord + 1) * ny * nx, 0.0);
        for (int m = 0; m <= mord; ++m) {
            for (int j = 0; j < ny; ++j) {
                for (int i = 0; i < nx; ++i) {
                    lfunc[idx4(0, m, j, i, norder, mord, ny, nx)] = 1.0;
                    if (norder >= 1) {
                        lfunc[idx4(1, m, j, i, norder, mord, ny, nx)] = (
                            1.0 - r2_over_sigma2[static_cast<size_t>(j) * nx + static_cast<size_t>(i)]
                            + static_cast<double>(m)
                        );
                    }
                }
            }
        }

        for (int n = 2; n <= norder; ++n) {
            for (int m = 0; m <= mord; ++m) {
                const double a_coeff = 2.0;
                const double b_coeff = 1.0 + (static_cast<double>(m) - 1.0) / static_cast<double>(n);
                for (int j = 0; j < ny; ++j) {
                    for (int i = 0; i < nx; ++i) {
                        const size_t base_idx = static_cast<size_t>(j) * nx + static_cast<size_t>(i);
                        const double r_term = (
                            static_cast<double>(m) - 1.0
                            - r2_over_sigma2[base_idx]
                        ) / static_cast<double>(n);
                        const double prev1 = lfunc[idx4(n - 1, m, j, i, norder, mord, ny, nx)];
                        const double prev2 = lfunc[idx4(n - 2, m, j, i, norder, mord, ny, nx)];
                        lfunc[idx4(n, m, j, i, norder, mord, ny, nx)] = (
                            (a_coeff + r_term) * prev1
                            - b_coeff * prev2
                        );
                    }
                }
            }
        }

        std::vector<std::complex<double>> chi(total_size);
        for (int nn = 0; nn <= norder; ++nn) {
            for (int mm = nn; mm >= 0; mm -= 2) {
                const int abs_m = std::abs(mm);
                const int c1 = (nn - abs_m) / 2;
                const int d1 = (nn + abs_m) / 2;
                const double factorial_c1 = std::tgamma(static_cast<double>(c1) + 1.0);
                const double factorial_d1 = std::tgamma(static_cast<double>(d1) + 1.0);
                double coeff_ratio = 0.0;
                if (factorial_d1 != 0.0) {
                    coeff_ratio = factorial_c1 / factorial_d1;
                }
                const double sign = (d1 % 2 == 0) ? 1.0 : -1.0;
                const double sqrt_coeff = std::sqrt(coeff_ratio);
                const std::complex<double> i_pow = pow_i(nn);
                for (int j = 0; j < ny; ++j) {
                    for (int i = 0; i < nx; ++i) {
                        const size_t base_idx = static_cast<size_t>(j) * nx + static_cast<size_t>(i);
                        const double r_pow = std::pow(
                            r2_over_sigma2[base_idx],
                            static_cast<double>(abs_m) / 2.0
                        );
                        const double laguerre = lfunc[idx4(
                            c1, abs_m, j, i, norder, mord, ny, nx
                        )];
                        std::complex<double> eul_pow = 1.0;
                        if (abs_m > 0) {
                            eul_pow = std::pow(eulfunc[base_idx], abs_m);
                        }
                        const std::complex<double> value = (
                            sign
                            * sqrt_coeff
                            * laguerre
                            * r_pow
                            * gauss[base_idx]
                            * i_pow
                        );
                        chi[idx4(nn, mm, j, i, norder, mord, ny, nx)] = value * eul_pow;
                    }
                }
            }
        }

        // Divide Gaussian kernel by npix**2 for inverse FFT normalization
        const double norm = static_cast<double>(npix) * static_cast<double>(npix);
        for (auto& val : chi) {
            val /= norm;
        }

        return {ny, nx, norder, std::move(chi)};
    }
} // namespace detail

inline py::object gauss_kernel_rfft(
    int ny,
    int nx,
    double sigma,
    double kmax,
    bool return_grid = false
) {
    const int nx_r = nx / 2 + 1;
    py::array_t<double> kernel({ny, nx_r});
    py::array_t<double> ygrid({ny, nx_r});
    py::array_t<double> xgrid({ny, nx_r});

    auto k = kernel.mutable_unchecked<2>();
    auto yg = ygrid.mutable_unchecked<2>();
    auto xg = xgrid.mutable_unchecked<2>();

    std::vector<double> xfreq = detail::make_rfftfreq(nx);
    std::vector<double> yfreq = detail::make_fftfreq(ny);

    const double kmax2 = kmax * kmax;
    const double sigma2 = sigma * sigma;

    for (int j = 0; j < ny; ++j) {
        const double ky = yfreq[j];
        for (int i = 0; i < nx_r; ++i) {
            const double kx = xfreq[i];
            const double r2 = kx * kx + ky * ky;
            const double mask = (r2 <= kmax2) ? 1.0 : 0.0;
            k(j, i) = std::exp(-r2 / (2.0 * sigma2)) * mask;
            yg(j, i) = ky;
            xg(j, i) = kx;
        }
    }

    if (!return_grid) {
        return kernel;
    }

    return py::make_tuple(kernel, py::make_tuple(ygrid, xgrid));
}

inline double m00_to_flux(
    double m00,
    double sigma_shapelets,
    double pixel_scale
) {
    const double sigma_use = sigma_shapelets / std::sqrt(2.0);
    const double sigma_pix = sigma_use / pixel_scale;
    const double ff = 4.0 * M_PI * sigma_pix * sigma_pix;
    return m00 * pixel_scale * pixel_scale * ff;
}

inline py::array_t<double>
m00_to_flux(
    const py::array_t<double>& m00,
    double sigma_shapelets,
    double pixel_scale
) {
    const double factor = m00_to_flux(
        1.0,
        sigma_shapelets,
        pixel_scale
    );

    int nn = m00.shape(0);
    auto in_r = m00.unchecked<1>();
    py::array_t<double> flux(nn);
    auto out_r = flux.mutable_unchecked<1>();
    for (py::ssize_t i = 0; i < nn; ++i) {
        out_r(i) = in_r(i) * factor;
    }
    return flux;
}

inline py::array_t<std::complex<double>>
shapelets2d_func(int npix, int norder, double sigma, double kmax) {
    detail::ShapeletResult result = detail::compute_shapelet_modes(
        npix, norder, sigma, kmax
    );
    const int ny = result.ny;
    const int nx = result.nx;
    const int nmode = (norder + 1) * (norder + 1);
    py::array_t<std::complex<double>> chi({nmode, ny, nx});
    auto out = chi.mutable_unchecked<3>();
    for (int nn = 0; nn <= norder; ++nn) {
        for (int mm = 0; mm <= norder; ++mm) {
            const int mode_index = nn * (norder + 1) + mm;
            for (int j = 0; j < ny; ++j) {
                for (int i = 0; i < nx; ++i) {
                    const size_t idx = detail::idx4(
                        nn, mm, j, i, norder, norder, ny, nx
                    );
                    out(mode_index, j, i) = result.chi[idx];
                }
            }
        }
    }
    return chi;
}

inline py::array_t<double>
shapelets2d(int npix, double sigma, double kmax) {
    int norder = 6; // default norder set to 6
    detail::ShapeletResult result = detail::compute_shapelet_modes(
        npix, norder, sigma, kmax
    );
    const int ny = result.ny;
    const int nx = result.nx;
    // the output modes are:
    //  "m00", "m20", "m22c", "m22s", "m40", "m42c", "m42s", "m44c", "m44s",
    //  "m60", "m64c", "m64s"
    static constexpr std::array<int, 12> indices = {
        0, 14, 16, 16, 28, 30, 30, 32, 32, 42, 46, 46
    };
    static constexpr std::array<bool, 12> use_imag = {
        false, false, false, true, false, false, true,
        false, true, false, false, true
    };

    py::array_t<double> chi({12, ny, nx});
    auto out = chi.mutable_unchecked<3>();
    for (size_t mode = 0; mode < indices.size(); ++mode) {
        const int idx_mode = indices[mode];
        const int nn = idx_mode / (norder + 1);
        const int mm = idx_mode % (norder + 1);
        for (int j = 0; j < ny; ++j) {
            for (int i = 0; i < nx; ++i) {
                const size_t idx = detail::idx4(
                    nn, mm, j, i, norder, norder, ny, nx
                );
                const std::complex<double>& value = result.chi[idx];
                out(mode, j, i) = use_imag[mode] ? value.imag() : value.real();
            }
        }
    }
    return chi;
}

inline py::array_t<std::complex<double>>
detlets2d(int npix, double sigma, double kmax) {
    const int ny = npix;
    const int nx = npix / 2 + 1;
    const int det_nrot = 4;
    const int nmode = 3 * det_nrot;

    py::array_t<std::complex<double>> psi({nmode, ny, nx});
    auto out = psi.mutable_unchecked<3>();

    std::vector<double> xfreq = detail::make_rfftfreq(npix);
    std::vector<double> yfreq = detail::make_fftfreq(npix);

    std::vector<double> gauss(static_cast<size_t>(ny) * nx, 0.0);
    std::vector<double> k1grid(static_cast<size_t>(ny) * nx, 0.0);
    std::vector<double> k2grid(static_cast<size_t>(ny) * nx, 0.0);

    const double kmax2 = kmax * kmax;
    const double sigma2 = sigma * sigma;

    for (int j = 0; j < ny; ++j) {
        const double ky = yfreq[j];
        for (int i = 0; i < nx; ++i) {
            const double kx = xfreq[i];
            const double r2 = kx * kx + ky * ky;
            const size_t idx = static_cast<size_t>(j) * nx + static_cast<size_t>(i);
            if (r2 <= kmax2) {
                gauss[idx] = std::exp(-r2 / (2.0 * sigma2));
            } else {
                gauss[idx] = 0.0;
            }
            k1grid[idx] = kx;
            k2grid[idx] = ky;
        }
    }

    const double norm = static_cast<double>(npix) * static_cast<double>(npix);
    std::vector<double> q1(static_cast<size_t>(ny) * nx, 0.0);
    std::vector<double> q2(static_cast<size_t>(ny) * nx, 0.0);
    std::vector<std::complex<double>> d1(static_cast<size_t>(ny) * nx);
    std::vector<std::complex<double>> d2(static_cast<size_t>(ny) * nx);

    for (int j = 0; j < ny; ++j) {
        for (int i = 0; i < nx; ++i) {
            const size_t idx = static_cast<size_t>(j) * nx + static_cast<size_t>(i);
            const double gaussian = gauss[idx] / norm;
            const double k1 = k1grid[idx];
            const double k2 = k2grid[idx];
            q1[idx] = ((k1 * k1) - (k2 * k2)) / sigma2 * gaussian;
            q2[idx] = (2.0 * k1 * k2) / sigma2 * gaussian;
            d1[idx] = std::complex<double>(0.0, -k1) * gaussian;
            d2[idx] = std::complex<double>(0.0, -k2) * gaussian;
            gauss[idx] = gaussian;
        }
    }

    for (int irot = 0; irot < det_nrot; ++irot) {
        const double angle = 2.0 * M_PI * static_cast<double>(irot) / det_nrot;
        const double cx = std::cos(angle);
        const double cy = std::sin(angle);
        for (int j = 0; j < ny; ++j) {
            for (int i = 0; i < nx; ++i) {
                const size_t idx = static_cast<size_t>(j) * nx + static_cast<size_t>(i);
                const double k1 = k1grid[idx];
                const double k2 = k2grid[idx];
                const double phase = k1 * cx + k2 * cy;
                const std::complex<double> foub = std::polar(1.0, phase);

                const std::complex<double> psi0 = gauss[idx] * (1.0 - foub);
                const std::complex<double> term1 = q1[idx] + cx * d1[idx] - cy * d2[idx];
                const std::complex<double> term2 = q2[idx] + cy * d1[idx] + cx * d2[idx];

                out(0 * det_nrot + irot, j, i) = psi0;
                out(1 * det_nrot + irot, j, i) = q1[idx] - term1 * foub;
                out(2 * det_nrot + irot, j, i) = q2[idx] - term2 * foub;
            }
        }
    }

    return psi;
}

inline double
get_kmax(
    const py::array_t<double>& psf_pow,
    double sigma,
    double kmax_thres = 1e-20
) {
    auto psf = psf_pow.unchecked<2>();
    const int npix = static_cast<int>(psf.shape(0));
    const int nx = static_cast<int>(psf.shape(1));

    int nx_full = (nx - 1) * 2;
    if (nx_full <= 0) {
        nx_full = nx;
    }
    std::vector<double> xfreq = detail::make_rfftfreq(nx_full);
    std::vector<double> yfreq = detail::make_fftfreq(npix);

    const double sigma2 = sigma * sigma;
    const double kmax = M_PI;
    const double kmax2 = kmax * kmax;

    double min_r = std::numeric_limits<double>::max();
    bool found = false;

    for (int j = 0; j < npix; ++j) {
        const double ky = yfreq[j];
        for (int i = 0; i < nx; ++i) {
            const double kx = xfreq[i];
            const double r2 = kx * kx + ky * ky;
            if (r2 > kmax2) {
                continue;
            }
            const double gaussian = std::exp(-r2 / (2.0 * sigma2));
            double psf_val = psf(j, i);
            double ratio;
            if (psf_val == 0.0) {
                ratio = std::numeric_limits<double>::infinity();
            } else {
                ratio = gaussian / psf_val;
            }
            if (ratio < kmax_thres) {
                const double r = std::sqrt(r2);
                if (r < min_r) {
                    min_r = r;
                    found = true;
                }
            }
        }
    }

    if (!found) {
        return static_cast<double>(npix / 2 - 1);
    }

    const double dk = 2.0 * M_PI / static_cast<double>(npix);
    double kmax_pix = std::round(min_r / dk);
    const double min_pix = static_cast<double>(npix / 5);
    const double max_pix = static_cast<double>(npix / 2 - 1);
    if (kmax_pix < min_pix) {
        kmax_pix = min_pix;
    }
    if (kmax_pix > max_pix) {
        kmax_pix = max_pix;
    }
    return kmax_pix;
}

void pyExportFpfsBase(py::module_& fpfs);
} // namespace fpfs
}

#endif // ANACAL_FPFS_BASE_H
