#include "Processor_ObjectExtraction.h"
#include <sstream>

#define _USE_MATH_DEFINES
#include <math.h>
#include <QVector3D>
#include <cmath>

#include <QFile>

int Processor_ObjectExtraction::start() {
    inputMat.push_back(new cv::Mat());
    outputMat.push_back(new cv::Mat());
    //outputMat.push_back(new cv::Mat());

    return 0;
}

Processor_ObjectExtraction::Processor_ObjectExtraction() {
    Processor();
    inputType = InputTypeMat;

    mode = Processor_ObjectExtraction_Mode_Normal;

    pMOG2_ = cv::createBackgroundSubtractorMOG2();

    m_median_khalfsize = 14;

    m_workspace_x_offset = 0.058;
    m_workspace_y_offset = 0.0535;
    m_workspace_z_offset = -0.005;

    m_workspace_x_size = 0.420;
    m_workspace_y_size = 0.297;

    m_object_height = 0.011;

    m_minimum_object_area = 1000;

    m_background_filename = "background.png";
}

cv::Point2d Processor_ObjectExtraction::pixel_to_meters(cv::Point2d pixel, cv::Size img_size) {
    cv::Point2d retval;
    
    retval.x = pixel.x/(img_size.width-1) * m_workspace_x_size - m_workspace_x_offset;
    retval.y = (1.0 - pixel.y/(img_size.height-1)) * m_workspace_y_size - m_workspace_y_offset;
    
    return retval;
}

int Processor_ObjectExtraction::process(int,int) {
    switch(mode) {
    case Processor_ObjectExtraction_Mode_GetBackground:
        pMOG2_->apply(*(inputMat[0]), *(outputMat[0]), 1.0);
        cv::cvtColor(*(outputMat[0]), *(outputMat[0]), CV_GRAY2BGR);
        mode = Processor_ObjectExtraction_Mode_Normal;
            break;
        case Processor_ObjectExtraction_Mode_Normal:
        case Processor_ObjectExtraction_Mode_Object_Selected:
        {
            pMOG2_->apply(*(inputMat[0]), *(outputMat[0]), 0.0);
            
            *(outputMat[0]) = *(outputMat[0]) == 255;
            
            cv::medianBlur(*(outputMat[0]), *(outputMat[0]), m_median_khalfsize*2+1);
            
            if (mode == Processor_ObjectExtraction_Mode_Normal) {
                cv::connectedComponentsWithStats(*(outputMat[0]), labels_, stats_, centroids_, 8, CV_16U);
            }
            
            QVector<sObject> objects;
            
            cv::Mat labels_filtered(labels_.rows, labels_.cols, labels_.type());
            labels_filtered.setTo(0);
            
            for (int idx = 1; idx < stats_.rows; idx++) {
                sObject temp;
                temp.image_area = stats_.at<int32_t>(idx, cv::CC_STAT_AREA);
                if (temp.image_area > m_minimum_object_area && temp.image_area < m_maximum_object_area) {
                    temp.image_x = centroids_.at<double>(idx, 0);
                    temp.image_y = centroids_.at<double>(idx, 1);
                    
                    temp.mask = (labels_ == idx);
                    
                    vector<vector<cv::Point> > contours;
                    vector<cv::Vec4i> hierarchy;
                    findContours(temp.mask, contours, hierarchy, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE, cv::Point(0, 0));
                    
                    if (contours.size() != 1) {
                        std::cerr << "Processor_ObjectExtraction::process() wrong contours " << contours.size() << std::endl;
                    } else {
                        temp.bounding_box = minAreaRect(cv::Mat(contours[0]));
                        if (temp.bounding_box.size.width > temp.bounding_box.size.height) {
                            temp.bounding_box.angle += 90.0;
                            temp.bounding_box.size = cv::Size(temp.bounding_box.size.height, temp.bounding_box.size.width);
                        }
                    }
                    
                    cv::Point2d centroid_m = pixel_to_meters(cv::Point2d(temp.image_x, temp.image_y), labels_.size());
                    
                    temp.workspace_x = centroid_m.x;
                    temp.workspace_y = centroid_m.y;
                    temp.workspace_z = m_workspace_z_offset;
                    
                    temp.size_x = -1;
                    temp.size_y = -1;
                    temp.size_z = m_object_height;
                    
                    objects.push_back(temp);
                    labels_filtered.setTo(objects.size(), temp.mask);
                }
            }
            
            if (screenshot_) {
                imwrite(screenshotFilename_.toStdString()+"screenshot.png", *(inputMat[0]));

                QFile objects_file(screenshotFilename_+"objects.txt");
                if(objects_file.open(QIODevice::WriteOnly | QIODevice::Truncate | QIODevice::Text)) {
                    QTextStream objects_stream(&objects_file);
                    for (int i = 0; i < objects.size(); i++) {
                        objects_stream << i << ", " << objects[i].to_string().c_str() << endl;
                        std::string mask_filename = screenshotFilename_.toStdString()+"object_mask_" + std::to_string(i) + ".png";
                        imwrite(mask_filename, objects[i].mask);
                    }
                    
                }
                screenshot_ = false;
            }
            
            emit object_data(*(inputMat[0]), objects);
            
            
            //cv::erode(*(outputMat[0]), *(outputMat[0]), cv::Mat(9,9,CV_8UC1));
            //cv::dilate(*(outputMat[0]), *(outputMat[0]), cv::Mat(9,9,CV_8UC1));
            
            /* DISPLAY */
            
            //double number_of_components;
            //minMaxLoc(labels_filtered, NULL, &number_of_components);
            
            labels_filtered.convertTo(*(outputMat[0]), CV_8UC1, 250.0/objects.size());
            cv::applyColorMap(*(outputMat[0]), *(outputMat[0]), cv::COLORMAP_JET);
            //cv::cvtColor(*(outputMat[0]), *(outputMat[0]), CV_GRAY2BGR);
            
            for (int i = 0; i < objects.size(); i++) {
                std::string text_label = "id=" + std::to_string(i);// + "\na=" + std::to_string(objects[i].bounding_box.angle);
                
                //putText(*(outputMat[0]), text_label, cv::Point(objects[idx].x, objects[idx].y), cv::FONT_HERSHEY_PLAIN, 1.0, cvScalar(0,0,0), 3);
                putText(*(outputMat[0]), text_label, cv::Point(objects[i].image_x-20, objects[i].image_y+20), cv::FONT_HERSHEY_PLAIN, 2.0, cvScalar(255,255,255), 2);
                
                cv::Point2f rect_points[4];
                objects[i].bounding_box.points(rect_points);
                for(int j = 0; j < 4; j++) {
                    line(*(outputMat[0]), rect_points[j], rect_points[(j+1)%4], cvScalar(255,255,255), 2, 8 );
                }
                
                line(*(outputMat[0]), cv::Point(objects[i].image_x, objects[i].image_y), cv::Point(objects[i].image_x+20*cos(objects[i].rz()), objects[i].image_y+20*sin(objects[i].rz())), cvScalar(0,0,255), 2, 8 );
            }
            
            /* CLICK */
            
            if (mode == Processor_ObjectExtraction_Mode_Object_Selected) {
                if (x_ < 0) x_ = 0;
                if (x_ > inputMat[0]->cols) x_ = inputMat[0]->cols-1;
                if (y_ < 0) y_ = 0;
                if (y_ > inputMat[0]->rows) y_ = inputMat[0]->rows-1;
                
                
                int idx = labels_filtered.at<uint16_t>(y_, x_);
                
                cv::Point2d centroid;
                
                if (idx > 0) {
                    idx--;
                    std::cerr << "IDX: " << idx << " " << objects[idx].image_x << "," << objects[idx].image_y << std::endl;
                    centroid.x = objects[idx].image_x;
                    centroid.y = objects[idx].image_y;
                    cv::Point2d centroid_m = pixel_to_meters(centroid, labels_.size());
                    
                    std::cerr << centroid_m.x << " " << centroid_m.y << std::endl;
                    
                    emit pickObject(QVector3D((float)centroid_m.x, centroid_m.y, m_workspace_z_offset), m_object_height, objects[idx].rz(), 1);
                    
                    qDebug() << QVector3D((float)centroid_m.x, centroid_m.y, m_workspace_z_offset + m_object_height) << objects[idx].rz() << 1;
                    
                    cv::circle(*(outputMat[0]), cv::Point(centroid.x, centroid.y), 10, cv::Scalar(255,255,255));
                }

                
                mode = Processor_ObjectExtraction_Mode_Normal;
            }
            break;
        }
    }
    
    return 0;
}

int Processor_ObjectExtraction::processSpecial(string s) {
    std::cerr << "Processor_ObjectExtraction::processSpecial(" << s << ")" << std::endl;
    
    
    if(s == "get_background") {
        mode = Processor_ObjectExtraction_Mode_GetBackground;
    }

    if(s == "save_background") {
        cv::Mat background;
        pMOG2_->getBackgroundImage(background);
        cv::imwrite("background.png", background);
        //mode = Processor_ObjectExtraction_Mode_Waiting;
    }
    if(s == "load_background") {
        cv::Mat background = cv::imread("background.png");
        if(!background.data) {
            std::cerr << "Processor_ObjectExtraction: could not load background data" << std::endl ;
        } else {
            pMOG2_->apply(background, background, 1.0);
        }
    }
    
    if(s.compare(0, sizeof("object_picked")-1, "object_picked") == 0) {
        stringstream ss(s);
        string cmd;
        ss >> cmd >> x_ >> y_;
        std::cerr << "object picked: " << x_ << "," << y_ << std::endl;
        mode = Processor_ObjectExtraction_Mode_Object_Selected;
    }
    
    if(s == "screenshot") {
        screenshot_ = true;
    }
    

    return 0;
}
void Processor_ObjectExtraction::screenshot(const QString& filename)
{
    screenshotFilename_ = filename;
    screenshot_ = true;

}

int Processor_ObjectExtraction::stop() {
    delete inputMat[0];
    delete outputMat[0];
    inputMat.clear();
    outputMat.clear();
    return 0;
}

void Processor_ObjectExtraction::writeXml(QXmlStreamWriter& stream) {
    stream.writeStartElement(sXmlKey());
    metaDataXmlWriter(stream,*this);
    stream.writeEndElement();
    
    cv::Mat background;
    pMOG2_->getBackgroundImage(background);
    cv::imwrite("background.png", background);
}

void Processor_ObjectExtraction::readXml(QXmlStreamReader& stream) {
    while(!(stream.tokenType() == QXmlStreamReader::EndElement && stream.name() == sXmlKey())) {
        cXmlIO::readXml(stream);
        stream.readNext();
    }
    cv::Mat background = cv::imread("background.png");
    pMOG2_->apply(background, background, 1.0);
}
