//
// Created by brett on 7/17/23.
//

#ifndef PARKSNREC_V2UTIL_H
#define PARKSNREC_V2UTIL_H


namespace parks {
    
    enum class ParameterType {
        SCALAR, COLOR, VARIABLE, IMAGE
    };
    
    struct Image {
        std::shared_ptr<double> image;
        unsigned int width, height;
        
        void write(unsigned int x, unsigned int y, const blt::vec3d& color) const {
            auto pos = x * CHANNELS + y * width * CHANNELS;
            image.get()[pos] = color.x();
            image.get()[pos + 1] = color.y();
            image.get()[pos + 2] = color.z();
        }
        
        [[nodiscard]] blt::vec3d read(unsigned int x, unsigned int y) const {
            auto pos = x * CHANNELS + y * width * CHANNELS;
            auto r = image.get()[pos];
            auto g = image.get()[pos + 1];
            auto b = image.get()[pos + 2];
            return blt::vec3d{r, g, b};
        }
    };
    
    struct Parameter {
        private:
            ParameterType type;
            std::variant<double, blt::vec3d, unsigned int, Image> value;
        public:
            explicit Parameter(double s): value(s) { type = ParameterType::SCALAR; }
            
            explicit Parameter(blt::vec3d c): value(c) { type = ParameterType::COLOR; }
            
            explicit Parameter(unsigned int v): value(v) { type = ParameterType::VARIABLE; }
            
            explicit Parameter(Image&& i): value(i) { type = ParameterType::IMAGE; }
            
            [[nodiscard]] ParameterType getType() const {
                return type;
            }
            
            template<typename T>
            [[nodiscard]] T& get() {
                return std::get<T>(value);
            }
            
            template<typename T>
            [[nodiscard]] const T& get() const {
                return std::get<T>(value);
            }
            
            template<typename Op>
            Parameter apply(Op f, double scalar) const {
                if (type == ParameterType::SCALAR)
                    return Parameter{f(get<double>(), scalar)};
                else if (type == ParameterType::COLOR) {
                    auto color = get<blt::vec3d>();
                    return Parameter{blt::vec3d{f(color.x(), scalar), f(color.y(), scalar),
                                                f(color.z(), scalar)}};
                } else if (type == ParameterType::VARIABLE) {
                    return Parameter{f((double) get<unsigned int>(), scalar)};
                } else if (type == ParameterType::IMAGE) {
                    auto image = get<Image>();
                    auto newImage = Image{std::shared_ptr<double>(
                            new double[image.width * image.height * CHANNELS]
                    ), image.width, image.height};
                    
                    for (unsigned int i = 0; i < image.width; i++) {
                        for (unsigned int j = 0; j < image.height; j++) {
                            auto oldColor = image.read(i, j);
                            newImage.write(
                                    i, j, blt::vec3d{f(oldColor.x(), scalar),
                                                     f(oldColor.y(), scalar),
                                                     f(oldColor.z(), scalar)}
                            );
                        }
                    }
                    
                    return Parameter{std::move(newImage)};
                }
            }
            
            template<typename Op>
            Parameter apply(Op f, const blt::vec3d& color) const {
                if (type == ParameterType::SCALAR)
                    return Parameter{
                            blt::vec3d{f(get<double>(), color.x()), f(get<double>(), color.y()),
                                       f(get<double>(), color.z())}};
                else if (type == ParameterType::COLOR) {
                    auto ourColor = get<blt::vec3d>();
                    return Parameter{
                            blt::vec3d{f(ourColor.x(), color.x()), f(ourColor.y(), color.y()),
                                       f(ourColor.z(), color.z())}};
                } else if (type == ParameterType::VARIABLE) {
                    return Parameter{blt::vec3d{f(get<unsigned int>(), color.x()),
                                                f(get<unsigned int>(), color.y()),
                                                f(get<unsigned int>(), color.z())}};
                } else if (type == ParameterType::IMAGE) {
                    auto image = get<Image>();
                    auto newImage = Image{std::shared_ptr<double>(
                            new double[image.width * image.height * CHANNELS]
                    ), image.width, image.height};
                    
                    for (unsigned int i = 0; i < image.width; i++) {
                        for (unsigned int j = 0; j < image.height; j++) {
                            auto oldColor = image.read(i, j);
                            newImage.write(
                                    i, j, blt::vec3d{f(oldColor.x(), color.x()),
                                                     f(oldColor.y(), color.y()),
                                                     f(oldColor.z(), color.z())}
                            );
                        }
                    }
                    
                    return Parameter{std::move(newImage)};
                }
            }
            
            template<typename Op>
            Parameter apply(Op f, unsigned int variable) const {
                if (type == ParameterType::SCALAR)
                    return Parameter{f(get<double>(), (double) variable)};
                else if (type == ParameterType::COLOR) {
                    auto color = get<blt::vec3d>();
                    return Parameter{blt::vec3d{f(color.x(), (double) variable),
                                                f(color.y(), (double) variable),
                                                f(color.z(), (double) variable)}};
                } else if (type == ParameterType::VARIABLE) {
                    return Parameter{f(get<unsigned int>(), variable)};
                } else if (type == ParameterType::IMAGE) {
                    auto image = get<Image>();
                    auto newImage = Image{std::shared_ptr<double>(
                            new double[image.width * image.height * CHANNELS]
                    ), image.width, image.height};
                    
                    for (unsigned int i = 0; i < image.width; i++) {
                        for (unsigned int j = 0; j < image.height; j++) {
                            auto oldColor = image.read(i, j);
                            newImage.write(
                                    i, j, blt::vec3d{f(oldColor.x(), (double) variable),
                                                     f(oldColor.y(), (double) variable),
                                                     f(oldColor.z(), (double) variable)}
                            );
                        }
                    }
                    
                    return Parameter{std::move(newImage)};
                }
            }
            
            template<typename Op>
            Parameter apply(Op f, const Image& image) const {
                if (type == ParameterType::IMAGE) {
                    auto ourImage = get<Image>();
                    if (ourImage.width != image.width || ourImage.height != image.height) {
                        BLT_ERROR("Unable to apply to images of differing sizes!");
                        throw std::runtime_error("Unable to apply to images of differing sizes!");
                    }
                    auto newImage = Image{std::shared_ptr<double>(
                            new double[ourImage.width * ourImage.height * CHANNELS]
                    ), ourImage.width, ourImage.height};
                    
                    for (unsigned int i = 0; i < ourImage.width; i++) {
                        for (unsigned int j = 0; j < ourImage.height; j++) {
                            auto oldColor = ourImage.read(i, j);
                            auto newColor = image.read(i, j);
                            newImage.write(
                                    i, j, blt::vec3d{f(oldColor.x(), (double) newColor.x()),
                                                     f(oldColor.y(), (double) newColor.y()),
                                                     f(oldColor.z(), (double) newColor.z())}
                            );
                        }
                    }
                    
                    return Parameter{std::move(newImage)};
                } else {
                    BLT_WARN("Please apply non-images to images instead of images to non-images");
                }
            }
            
            template<typename Op>
            inline Parameter apply(Op f, const Parameter& param) const {
                if (param.type == ParameterType::SCALAR)
                    return apply(f, param.get<double>());
                else if (type == ParameterType::COLOR)
                    return apply(f, param.get<blt::vec3d>());
                else if (type == ParameterType::VARIABLE)
                    return apply(f, param.get<unsigned int>());
                else
                    return apply(f, param.get<Image>());
            }
    };
    
}


#endif //PARKSNREC_UTIL_H