// Autor: Lenka Sokova
// xlogin: xsokov01

#include <iostream>
#include <vector>
#include <bits/stdc++.h>
#include <iostream>

#include <filesystem> // file management
#include <map> // store and count time
#include <cmath> // brightness
#include <string>
#include <regex>

// opencv
#include "opencv2/imgcodecs.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/imgproc.hpp"
#include "opencv2/core.hpp"

// codes
#include "hough_transform.h"

using namespace std;
using namespace cv;
namespace fs = std::filesystem;

#define MAX_THRESH 255

struct ProgramArgs{
    fs::path file;
    fs::path directory;
    fs::path writeDirectory;
    int scale;
};


void resize_img_scale(Mat srcImg, Mat *OutputImg, uint scale){

    int new_width = srcImg.cols * scale / 100;
    int new_height = srcImg.rows * scale / 100;

    resize(srcImg, *OutputImg, Size(new_width, new_height), INTER_LINEAR);
}

// apply green mask
int apply_mask(Mat img, Mat& green, Scalar lower_color, Scalar upper_color){

    Mat hsv_img; Mat green_mask;

    // convert to hsv image
    cvtColor(img, hsv_img, COLOR_BGR2HSV);

    // get mask from hsv image
    inRange(hsv_img, lower_color, upper_color, green_mask);

    Mat green_imask = (green_mask > 0);

    double greenStatic = static_cast<double>(countNonZero(green_imask)) / (green_imask.cols * green_imask.rows) * 100;

    // check the ration of green pixels
    if (greenStatic < 3.0) {
        cout << "No green laser in image" << endl;
        return 1;
    }

    green = Mat::zeros(img.cols, img.rows, CV_8UC3);
    img.copyTo(green, green_imask);

    return 0;
}

bool isNumber(const std::string& str) {
    for (char c : str) {
        if (!std::isdigit(c)) {
            return false;
        }
    }
    return true;
}

bool isInRange(const std::string& str, int min, int max) {
    int num = std::stoi(str);
    return num >= min && num <= max;
}

// parse input arguments
int parseArguments(int argc, char **argv, ProgramArgs *programArgs) {

    // process file argument
    for (int i = 1; i < argc; i++) {
        if (string(argv[i]) == "-f") {
            if (i + 1 < argc) {
                programArgs->file = argv[i + 1];
            } else {
                std::cerr << "Error: no file name specified after -f option" << std::endl;
                return 1;
            }
        } else if (string(argv[i]) == "-s") { // process scale argument

            if (i + 1 < argc) {
                if (!isNumber(argv[i + 1])) {
                    std::cerr << "Error: Input is not a number" << std::endl;
                    return 1;
                } else if (!isInRange(argv[i + 1], 1, 100)) {
                    std::cerr << "Error: Input number is not within the range of 1 to 100" << std::endl;
                    return 1;
                } else {
                    programArgs->scale = stoi(argv[i + 1]);
                }
            } else {
                std::cerr << "Error: no scale number specified after -s option" << std::endl;
                return 1;
            }
        } else if (string(argv[i]) == "-d") { // process directory argument
            if (i + 1 < argc) {
                programArgs->directory = argv[i + 1];
            } else {
                std::cerr << "Error: no file name specified after -d option" << std::endl;
                return 1;
            }
        }
        // process write argument
        if (string(argv[i]) == "-w") {

            if (i + 1 < argc && argv[i + 1][0] != '-') {
                programArgs->writeDirectory = argv[i + 1];
            }

            else {
                programArgs->writeDirectory = "result";
            }
        }
    }

    return 0;
}


void count_thresh_value(Mat img, uint *thresh_value){
    Point pt; uchar *mdata;

    int width = img.cols;
    int height = img.rows;

    vector<int> counts(MAX_THRESH + 1, 0);
    int sum_count = 0; int value;

    // run through image and count the number of pixels in each thresh value
    for (pt.y = 0; pt.y < height; pt.y++){
        mdata = img.ptr(pt.y);

        for (pt.x = 0; pt.x < width; pt.x++){
            value = mdata[pt.x];

            counts[value] += 1;
            sum_count += value;
        }
    }

    // find the thresh value
    int thresh_id =  sum_count * 0.02;
    int F = 0;
    for (int i = MAX_THRESH; i >= 0; i--){
        F += i * counts[i];

        if (F >= thresh_id){
            *thresh_value = i;
            return;
        }
    }

}

int loadImage(string imageName, Mat &img){

    // Loads input image
    img = imread(imageName, IMREAD_COLOR);

    // Check if image is loaded fine
    if(img.empty()){
        cerr << " Error opening image" << endl;
        cerr << " Program Arguments: [image_name -- default " << imageName << "]" << endl;
        return -1;
    }

    return 0;
}



int processImage(Mat src, string directoryPath, int scale){
    Mat img;

    // Resize image
    resize_img_scale(src, &img, scale);

    // Apply green mask
    Scalar lowerGreenBright = Scalar(30, 0, 0);
    Scalar upperGreenBright = Scalar(170, 255, 200);

    Mat green;
    if (int err = apply_mask(img, green, lowerGreenBright, upperGreenBright)) return 0;

    // Convert to grayscale
    cvtColor(green, img, COLOR_BGR2GRAY);

    // Apply smooth filter
    Mat blurredImg;
    if (mean(img)[0] > 10) { // Brighter image

        bilateralFilter(img, blurredImg, 6, 75, 75);

    }
    else { // Darker image

        GaussianBlur(img, blurredImg, cv::Size(5, 5), 0);
    }


    // Count thresh value
    uint thresh_value;
    count_thresh_value(blurredImg, &thresh_value);

    // apply threshold
    Mat threshImg;
    threshold(blurredImg, threshImg, thresh_value, 255, THRESH_BINARY);


    auto meanValue = mean(threshImg)[0];
    bool changed;

    // Adaptive threshold
    for (int i = 0; i < 5; i++) {

        if(meanValue < 0.001 and thresh_value > 70){
            thresh_value = 20;
            changed = true;
        }

        else if(meanValue < 0.001 and thresh_value > 50){
            thresh_value = 20;
            changed = true;
        }

        else if (meanValue < 0.001){
            thresh_value -= 10;
            changed = true;
        }

        else if (meanValue < 0.02 and thresh_value > 40) {
            thresh_value = 20;
            changed = true;

        }

        else if (meanValue < 0.02) {
            thresh_value -= 2;
            changed = true;

        }

        else if (meanValue < 0) {
            thresh_value -= 1;
            changed = true;

        }

        else if (meanValue > 1) {
            thresh_value += 1;
            changed = false;

        }

        else {
            break;
        }

        if (changed) {
            threshold(img, threshImg, thresh_value, 255, THRESH_BINARY);
            meanValue = mean(threshImg)[0];
        }
    }

    img = threshImg;

    // Find line using Hough transform
    vector <Vec4i> linesP;

    if (meanValue < 0.1)
        houghProbabilisticTransform(img, linesP, 1, CV_PI / 180, 50, 50, 100, 20);

    else if (meanValue < 0.5)
        houghProbabilisticTransform(img, linesP, 1, CV_PI / 180, 70, 100, 100, 20);

    else if (meanValue < 1)
        houghProbabilisticTransform(img, linesP, 1, CV_PI / 180, 100, 100, 100, 20);

    else
        houghProbabilisticTransform(img, linesP, 1, CV_PI / 180, 140, 100, 100, 20);


    float scale_x = (double)src.cols/ (double)img.cols;
    float scale_y = (double)src.rows/ (double)img.rows;

    double x1p = 0, y1p = 0, x2p = 0, y2p = 0;
    double k, q, x;
    int countLines = linesP.size();

    // write position on the line
    vector <int> linesOriginal;

    if (linesP.size() > 0) {

        for (size_t i = 0; i < linesP.size(); i++) {
            Vec4i l = linesP[i];

            // find the point in original image
            x1p += (l[0] * scale_x);
            y1p += (l[1] * scale_y);
            x2p += (l[2] * scale_x);
            y2p += (l[3] * scale_y);
        }


        if (countLines > 0){
            x1p = x1p / countLines;
            y1p = y1p / countLines;
            x2p = x2p / countLines;
            y2p = y2p / countLines;
        }


        double dx = x2p - x1p;
        double dy = y2p - y1p;

        if( dx != 0 ){
            k = dy / dx;
            q = ((y1p - k * x1p) + (y2p - k * x2p)) / 2;

            // check angle
            if (-0.1 < k and k < 0.1){
                cout << "No line detected" << endl;
                return 0;
            }
            cout << "- Line parameters:\tk = " << k << " q = " << q << endl;

            int left_y = k * 0 + q;
            int right_y = k * (src.cols - 1) + q;
            int top_x = (-q) / k;
            int bottom_x = (src.rows - 1 - q) / k;


            if (0 < left_y and left_y < src.rows){
                linesOriginal.push_back(0);
                linesOriginal.push_back(left_y);
            }

            if (0 < right_y and right_y < src.rows){
                linesOriginal.push_back(src.cols - 1);
                linesOriginal.push_back(right_y);
            }

            if (0 < top_x and top_x < src.cols){
                linesOriginal.push_back(top_x);
                linesOriginal.push_back(0);
            }

            if (0 < bottom_x and bottom_x < src.cols){
                linesOriginal.push_back(bottom_x);
                linesOriginal.push_back(src.rows - 1);
            }
        }

        else{
            x = x2p;
            linesOriginal.push_back(x); // x1
            linesOriginal.push_back(0); // y1
            linesOriginal.push_back(x); // x2
            linesOriginal.push_back(src.rows); // y2

            cout << "- Line parameters:\tx = " << x << endl;
        }

        cout << "- Position in Image:\tposX1 = " << to_string(linesOriginal[0]) << " posY1 = " << to_string(linesOriginal[1]) << " posX2 = " << to_string(linesOriginal[2]) << " posY2 = " << to_string(linesOriginal[3]) << endl;
        cout << "- Vector of line:\t(" << to_string(linesOriginal[2] - linesOriginal[0]) << ", " << to_string(linesOriginal[3] - linesOriginal[1]) << ")" << endl;

        line(src, Point(linesOriginal[0], linesOriginal[1]), Point(linesOriginal[2], linesOriginal[3]), Scalar(0, 255, 255), 3, LINE_AA);
    }

    else {
        cout << "No line detected" << endl;
    }

    return 0;
}


// process input directory or folder
int processDirectory(ProgramArgs programArgs) {
    int cnt = 1; bool write = false; string outputName; fs::path outputPath;

    // supported png or jpg images
    std::regex img_regex(".*\\.(?:png|jpg)$", std::regex_constants::icase);

    Mat src;

    // update write parameters
    if (programArgs.writeDirectory.string().length() > 0){
        write = true;

        if (!fs::exists(programArgs.writeDirectory)){
            if(fs::create_directory(programArgs.writeDirectory)){
                std::cout << "Failed to create directory" << endl;
            }
        }
    }

    // process input file
    if (fs::exists(programArgs.file)) {

        // check if file is in supported format
        if (regex_match(programArgs.file.filename().string(), img_regex)) {
            cout << programArgs.file.filename() << endl;

            if (int err = loadImage(programArgs.file.string(), src)) return err;
            if (int err = processImage(src, programArgs.writeDirectory.string(), programArgs.scale)) return err;

            if (write){
                outputName = programArgs.file.stem().string() + "-result" + programArgs.file.extension().string();
                outputPath = programArgs.writeDirectory / outputName;

                imwrite(outputPath.string(), src);
            }


        } else {
            std::cerr << "Error: file is not in supported format. Supported formats are png, jpg" << std::endl;
        }

    } else if (fs::exists(programArgs.directory)) {
        // run through directory

        for (const auto &entry: fs::directory_iterator(programArgs.directory)) {

            // Check if the entry is a file
            if (fs::is_regular_file(entry.path()) && regex_match(entry.path().filename().string(), img_regex)) {
                cout << cnt++ << ": " << entry.path().filename() << endl;

                if (int err = loadImage(entry.path().string(), src)) return err;
                if (int err = processImage(src, programArgs.writeDirectory.string(), programArgs.scale)) return err;

                if (write){
                    outputName = entry.path().stem().string() + "-result" + entry.path().extension().string();
                    outputPath = programArgs.writeDirectory / outputName;

                    imwrite(outputPath.string(), src);

                }


                cout << endl;
            }
        }
    }

        else{
            std::cerr << "Error: no file name or directory name is specify or do not exists" << std::endl;
            return 1;
        }

        return 0;
        
}



int main(int argc, char **argv){

    // initialize input arguments
    ProgramArgs programArgs;
    programArgs.scale = 60; programArgs.file = ""; programArgs.directory = ""; programArgs.writeDirectory = "";

    // parse input arguments
    if (int err = parseArguments(argc, argv, &programArgs)) return err;
    if (int err = processDirectory(programArgs)) return err;


    return 0;
}
