#include "image_utils.h"
#include "core/exception.h"

namespace image
{
    Image make_manyimg_composite(int count_width, int count_height, int img_cnt, std::function<Image(int cnt)> get_img_func)
    {
        Image first_img = get_img_func(0);
        int img_w = first_img.width();
        int img_h = first_img.height();

        Image image_all(first_img.depth(), img_w * count_width, img_h * count_height, first_img.channels());

        first_img.clear();

        for (int row = 0; row < count_height; row++)
        {
            for (int column = 0; column < count_width; column++)
            {
                if (row * count_width + column >= img_cnt)
                    return image_all;

                image_all.draw_image(0, get_img_func(row * count_width + column), img_w * column, img_h * row);
            }
        }

        return image_all;
    }

    Image blend_images(std::vector<Image> &images)
    {
        if(images.size() == 0)
            throw satdump_exception("No images passed!");

        size_t width = images[0].width();
        size_t height = images[0].height();
        int channels = images[0].channels();
        bool are_rgba = channels == 4;

        for (Image &this_img : images)
        {
            if(this_img.depth() != 16)
                throw satdump_exception("blend_images must be the same bit depth, and 16");
            width = std::min(width, this_img.width());
            height = std::min(height, this_img.height());
            are_rgba &= this_img.channels() == 4;
        }

        Image img_b(16, width, height, channels);

        for (int c = 0; c < channels; c++)
        {
            if (are_rgba)
            {
                for (size_t i = 0; i < height * width; i++)
                {
                    if (c == 3)
                    {
                        double final_alpha = 0.0;
                        for (Image &this_img : images)
                            final_alpha = std::max(final_alpha, this_img.getf(3, i));
                        img_b.setf(c, i, final_alpha);
                    }
                    else
                    {
                        float num_layers = 0.0f;
                        double final_val = 0.0;
                        for (Image &this_img : images)
                        {
                            final_val += this_img.getf(c, i) * this_img.getf(3, i);
                            num_layers += this_img.getf(3, i);
                        }

                        img_b.setf(c, i, final_val / num_layers);
                    }
                }
            }
            else
            {
                for (size_t i = 0; i < height * width; i++)
                {
                    size_t num_layers = images.size();
                    double final_val = 0.0f;
                    for (Image& this_img : images)
                    {
                        float layer_val = this_img.getf(c, i);
                        if (layer_val == 0)
                            num_layers--;
                        else
                            final_val += layer_val;
                    }

                    img_b.setf(c, i, final_val / num_layers);
                }
            }
        }
        return img_b;
    }

    image::Image merge_images_opacity(image::Image &img1, image::Image &img2, float op)
    {
        if (img1.depth() != img2.depth() || img1.depth() != 16 || img2.depth() != 16)
            throw satdump_exception("merge_images_opacity must be the same bit depth, and 16");

        const size_t width = std::min<int>(img1.width(), img2.width());
        const size_t height = std::min<int>(img1.height(), img2.height());
        const int64_t size = width * height;
        const int channels_1 = img1.channels();
        const int channels_2 = img2.channels();
        const int color_channels = std::min(3, channels_1);
        image::Image ret(16, width, height, channels_1);

#pragma omp parallel for
        for (int64_t i = 0; i < size; i++)
        {
            float alpha_1 = channels_1 == 4 ? (float)img1.get(3, i) / 65535.0f : 1.0f;
            float alpha_2 = (channels_2 == 4 ? (float)img2.get(3, i) / 65535.0f : 1.0f) * op;
            float ret_alpha = alpha_2 + alpha_1 * (1.0f - alpha_2);
            for (int j = 0; j < color_channels; j++)
                ret.set(j, i, ((alpha_2 * ((float)img2.get(j, i) / 65535.0f) + alpha_1 * ((float)img1.get(j, i) / 65535.0f) * (1.0f - alpha_2)) / ret_alpha) * 65535.0f);

            if (channels_1 == 4)
                ret.set(3, i, ret_alpha * 65535.0f);
            else
                for (int j = 0; j < color_channels; j++)
                    ret.set(j, i, (float)ret.get(j, i) * ret_alpha);
        }

        return ret;
    }
}