#include <iostream>
#include <vector>
#include "lodepng.h"

// Get the median value of a vector of values
unsigned char getMedian(std::vector<unsigned char>& values) {
    std::sort(values.begin(), values.end());
    return values[values.size() / 2];
}

// Apply median filter to an image
void applyMedianFilter(std::vector<unsigned char>& image, int width, int height, int channels, int filterSize) {
    std::vector<unsigned char> result(image.size());

    int radius = filterSize / 2;

    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            for (int c = 0; c < channels; c++) {
                std::vector<unsigned char> values;

                for (int i = -radius; i <= radius; i++) {
                    for (int j = -radius; j <= radius; j++) {
                        int px = std::min(std::max(x + j, 0), width - 1);
                        int py = std::min(std::max(y + i, 0), height - 1);

                        values.push_back(image[(py * width + px) * channels + c]);
                    }
                }

                result[(y * width + x) * channels + c] = getMedian(values);
            }
        }
    }

    image = result;
}

// Apply median filter to an image file
void medianFilter(const char* inputFilename, const char* outputFilename, int filterSize) {
    std::vector<unsigned char> image;
    unsigned width, height;

    // Load image from file
    unsigned error = lodepng::decode(image, width, height, inputFilename);

    if (error) {
        std::cerr << "Error decoding image: " << lodepng_error_text(error) << std::endl;
        return;
    }

    // Apply median filter
    applyMedianFilter(image, width, height, 4, filterSize);

    // Save image to file
    error = lodepng::encode(outputFilename, image, width, height);

    if (error) {
        std::cerr << "Error encoding image: " << lodepng_error_text(error) << std::endl;
        return;
    }
}