#include "color_filter.h"
#include "hsv_color.h"
#include "../image_utilities.h"
#include <opencv2/imgproc.hpp>
#include <iostream>
#include <cmath>
#include <stdexcept>

using namespace std;

pair<cv::Mat, cv::Mat> filter_color(const cv::Mat& img, const vector<vector<vector<int>>>& ranges, int fillval /* = 255 */) {
    cv::Mat out = cv::Mat::ones(img.size(), img.type()) * fillval;
    cv::Mat hsv;
    cv::cvtColor(img, hsv, cv::COLOR_BGR2HSV);
    
    cv::Mat mask = cv::Mat::zeros(img.size(), CV_8UC1);
    
    for (const auto& range_pair : ranges) {
        cv::Scalar lower(range_pair[0][0], range_pair[0][1], range_pair[0][2]);
        cv::Scalar upper(range_pair[1][0], range_pair[1][1], range_pair[1][2]);
        
        cv::Mat temp_mask;
        cv::inRange(hsv, lower, upper, temp_mask);
        
        mask += temp_mask;
    }
    
    cv::Mat pos = mask == 255;
    for (int y = 0; y < img.rows; y++) {
        for (int x = 0; x < img.cols; x++) {
            if (pos.at<uchar>(y, x)) {
                out.at<cv::Vec3b>(y, x) = img.at<cv::Vec3b>(y, x);
            }
        }
    }
    
    return {out, mask};
}

pair<cv::Mat, string> color_fliter_process(const cv::Mat& img, int threshold) {
    try {
        cv::Mat grayscale;
        cv::cvtColor(img, grayscale, cv::COLOR_BGR2GRAY);
        
        cv::Mat binary = block_threshold(grayscale);
        cv::Mat mask = binary > 0;
        
        vector<string> color_names = {"red"};
        vector<cv::Scalar> color_values = {{0, 0, 255}};
        int best_color_index = 0;
        double max_color_area = -1e10;
        cv::Mat final_mask = mask.clone();
        
        for (int i = 0; i < color_names.size(); i++) {
            string color_name = color_names[i];
            vector<vector<vector<int>>> ranges = get_color_range(color_name).check;
            auto [_, color_mask] = filter_color(img, ranges);
            cv::Mat color_threshold = color_mask == 255;
            cv::Mat combined_mask;
            cv::bitwise_and(color_threshold, mask, combined_mask);
            double color_area = cv::countNonZero(combined_mask);
            
            if (color_area > max_color_area) {
                best_color_index = i;
                max_color_area = color_area;
                
                vector<vector<vector<int>>> full_ranges = get_color_range(color_name).full;
                vector<vector<vector<int>>> strict_ranges = get_color_range(color_name).ranges2;
                vector<vector<vector<int>>> combined_ranges;
                
                for (int r = 0; r < full_ranges.size(); r++) {
                    vector<vector<int>> range_pair;
                    vector<int> lower;
                    vector<int> upper;
                    for (int k = 0; k < 3; k++) {
                        int lower_val = static_cast<int>(full_ranges[r][0][k] * (100 - threshold)/100.0 + strict_ranges[r][0][k] * threshold/100.0);
                        int upper_val = static_cast<int>(full_ranges[r][1][k] * (100 - threshold)/100.0 + strict_ranges[r][1][k] * threshold/100.0);
                        lower.push_back(lower_val);
                        upper.push_back(upper_val);
                    }
                    range_pair.push_back(lower);
                    range_pair.push_back(upper);
                    combined_ranges.push_back(range_pair);
                }
                
                auto [__, final_threshold] = filter_color(img, combined_ranges);
                final_mask = final_threshold == 255;
                cv::bitwise_and(final_mask, mask, final_mask);
            }
        }
        
        cv::Mat output = cv::Mat(img.size(), img.type(), cv::Scalar(255, 255, 255));
        cv::Scalar color_value = color_values[best_color_index];
        
        for (int y = 0; y < img.rows; y++) {
            for (int x = 0; x < img.cols; x++) {
                if (final_mask.at<uchar>(y, x)) {
                    output.at<cv::Vec3b>(y, x) = cv::Vec3b(color_value[0], color_value[1], color_value[2]);
                }
            }
        }
        
        return {output, color_names[best_color_index]};
    } catch (const exception& e) {
        cout << "Exception in color_fliter_process: " << e.what() << endl;
        return {img, "red"};
    } catch (...) {
        cout << "Unknown exception in color_fliter_process" << endl;
        return {img, "red"};
    }
}
