#ifndef ANACAL_MASK_H
#define ANACAL_MASK_H

#include <cmath>
#include "table.h"

namespace anacal {
namespace mask {
    void
    inline add_bright_star_mask(
        py::array_t<int16_t>& mask_array,
        const py::array_t<BrightStar>& star_array
    ) {
        auto star_r = star_array.unchecked<1>();
        int nn = star_array.shape(0);
        int ndim = mask_array.ndim();
        if (ndim != 2) {
            throw std::runtime_error(
                "Mask Error: The input mask array has an invalid shape."
            );
        }
        auto m_r = mask_array.mutable_unchecked<2>();
        int ny = m_r.shape(0);
        int nx = m_r.shape(1);
        for (int k = 0; k < nn; ++k) {
            int x = static_cast<int>(star_r(k).x + 0.5);
            int y = static_cast<int>(star_r(k).y + 0.5);
            float radius = star_r(k).r;
            if (!std::isfinite(radius) || radius < 0.0f) {
                continue;
            }
            int r = static_cast<int>(std::round(radius));
            if (r == 0) {
                r = 1;
            }
            int r2 = r * r;
            for (int j = y-r; j <= y+r; ++j) {
                if ((j < 0) || (j >= ny)) {
                    continue;
                }
                int dy2 = (j - y) * (j - y);
                for (int i = x-r; i <= x+r; ++i) {
                    if ((i < 0) || (i >= nx)) {
                        continue;
                    }
                    int dx2 = (i - x) * (i - x);
                    int d2 = dx2 + dy2;
                    if (d2 < r2) {
                        m_r(j, i) = m_r(j, i) | 1;
                    }
                }
            }
        }
        return;
    };

    void
    inline extend_mask_image(
        py::array_t<int16_t>& mask_array
    ) {
        // Get information of input array
        int ndim = mask_array.ndim();
        if (ndim != 2) {
            throw std::runtime_error(
                "Mask Error: The input mask array has an invalid shape."
            );
        }
        auto m_r = mask_array.mutable_unchecked<2>();
        int ny = m_r.shape(0);
        int nx = m_r.shape(1);

        // Initialize the convolved mask
        py::array_t<int16_t> mask_conv({ny, nx});
        auto conv_r = mask_conv.mutable_unchecked<2>();
        for (int j = 0; j < ny; ++j) {
            for (int i = 0; i < nx; ++i) {
                conv_r(j, i) = 0.0;
            }
        }

        // Convoved image
        for (int y = 0; y < ny; ++y) {
            for (int x = 0; x < nx; ++x) {
                if (m_r(y, x) > 0){
                    for (int j = y-2; j <= y+2; ++j) {
                        if ((j < 0) || (j >= ny)) {
                            continue;
                        }
                        for (int i = x-2; i <= x+2; ++i) {
                            if ((i < 0) || (i >= nx)) {
                                continue;
                            }
                            conv_r(j, i) = 1;
                        }
                    }
                }
            }
        }
        for (int j = 0; j < ny; ++j) {
            for (int i = 0; i < nx; ++i) {
                m_r(j, i) = conv_r(j, i);
            }
        }
        return;
    };

    void
    inline mask_galaxy_image(
        py::array_t<double>& gal_array,
        py::array_t<int16_t>& mask_array,
        bool do_extend_mask,
        const std::optional<py::array_t<BrightStar>>& star_array
    ) {
        if (do_extend_mask) {
            extend_mask_image(mask_array);
        }
        if (star_array.has_value()) {
            add_bright_star_mask(
                mask_array,
                *star_array
            );
        }

        auto img_r = gal_array.mutable_unchecked<2>();
        int ny = gal_array.shape(0);
        int nx = gal_array.shape(1);
        auto mask_r = mask_array.unchecked<2>();

        for (int j = 0; j < ny; ++j) {
            for (int i = 0; i < nx; ++i) {
                if (mask_r(j, i) > 0) {
                    img_r(j, i) = 0.0;
                }
            }
        }
    }

    py::array_t<float>
    inline convolve_mask_gauss(
        const py::array_t<int16_t>& mask_array,
        double sigma,
        double scale
    ) {
        int ngrid = int(sigma / scale) * 6 + 1;
        int ngrid2 = int((ngrid - 1) / 2);

        py::array_t<float> kernel({ngrid, ngrid});
        auto kernel_r = kernel.mutable_unchecked<2>();
        // Compute the Gaussian kernel
        float A = float(scale * scale / (2.0 * M_PI * sigma * sigma));
        float sigma2 = -1.0 / (2 * float(sigma * sigma));
        for (int y = 0; y < ngrid; ++y) {
            for (int x = 0; x < ngrid; ++x) {
                float dx = (x - ngrid2) * scale;
                float dy = (y - ngrid2) * scale;
                float r2 = dx * dx + dy * dy;
                kernel_r(y, x) = A * std::exp(r2 * sigma2);
            }
        }

        auto mask_r = mask_array.unchecked<2>();
        int ny = mask_r.shape(0);
        int nx = mask_r.shape(1);
        py::array_t<float> mask_conv({ny, nx});
        auto conv_r = mask_conv.mutable_unchecked<2>();
        for (int j = 0; j < ny; ++j) {
            for (int i = 0; i < nx; ++i) {
                conv_r(j, i) = 0.0;
            }
        }

        for (int y = 0; y < ny; ++y) {
            for (int x = 0; x < nx; ++x) {
                if (mask_r(y, x) > 0) {
                    for (int j = -ngrid2; j <= ngrid2; ++j) {
                        if ((y + j < 0) || (y + j >= ny)) {
                            continue;
                        }
                        for (int i = -ngrid2; i <= ngrid2; ++i) {
                            if ((x + i < 0) || (x + i >= nx)) {
                                continue;
                            }
                            conv_r(y + j, x + i) += (
                                mask_r(y, x) * kernel_r(j + ngrid2, i + ngrid2)
                            );
                        }
                    }
                }
            }
        }
        return mask_conv;
    }

    inline py::array_t<int> convolve_mask(
        py::array_t<int> mask_array,
        py::array_t<int> kernel
    ) {
        auto img = mask_array.unchecked<2>();
        auto ker = kernel.unchecked<2>();

        const int height = img.shape(0);
        const int width  = img.shape(1);
        const int kh = ker.shape(0);
        const int kw = ker.shape(1);
        if (kh % 2 != 1 || kw % 2 != 1)
            throw std::invalid_argument("Kernel size must be odd");
        const int kh2 = kh / 2;
        const int kw2 = kw / 2;

        py::array_t<int> mask_conv({height, width});
        auto out = mask_conv.mutable_unchecked<2>();

        // Initialize mask_conv to zero
        for (int i = 0; i < height; ++i)
            for (int j = 0; j < width; ++j)
                out(i, j) = 0.0;

        // Loop over only nonzero pixels
        for (int y = kh2; y < height - kh2; ++y) {
            for (int x = kw2; x < width - kw2; ++x) {
                int v = img(y, x);
                if (v == 0.0) continue;  // Skip zeros
                for (int dy = -kh2; dy <= kh2; ++dy) {
                    const int ny = y + dy;
                    const int ky = dy + kh2;
                    for (int dx = -kw2; dx <= kw2; ++dx) {
                        const int nx = x + dx;
                        const int kx = dx + kw2;
                        out(ny, nx) += v * ker(ky, kx);
                    }
                }
            }
        }
        return mask_conv;
    };

    void
    inline add_pixel_mask_column(
        std::vector<table::galNumber>& catalog,
        const py::array_t<int16_t>& mask_array,
        double sigma,
        double scale
    ) {
        py::array_t<float> mask_conv = convolve_mask_gauss(
            mask_array, sigma, scale
        );

        auto conv_r = mask_conv.unchecked<2>();
        int ny = conv_r.shape(0);
        int nx = conv_r.shape(1);

        for (table::galNumber & src : catalog) {
            int y = static_cast<int>(
                std::round(src.model.x2.v / scale)
            );
            int x = static_cast<int>(
                std::round(src.model.x1.v / scale)
            );
            if (y>=0 && y< ny && x>=0 && x<nx) {
                src.mask_value = static_cast<int>(conv_r(y, x) * 1000);
            }
        }
        return;
    };

    void pyExportMask(py::module& m);

} // end mask
} // end anacal

#endif // MASK_H
