#include "stdafx.h"
#include <opencv2/opencv.hpp>
#include <iostream>
#include <cstdlib>
#include <fstream>

using namespace std;

// Global storage
//
// Float, 3-channel images
//
cv::Mat image;
cv::Mat IavgF, IdiffF, IprevF, IhiF, IlowF;
cv::Mat tmp, tmp2, mask;

// Float, 1-channel images
//
vector<cv::Mat> Igray(3);
vector<cv::Mat> Ilow(3);
vector<cv::Mat> Ihi(3);

// Byte, 1-channel image
//
cv::Mat Imaskt;

// Thresholds
//
float high_thresh = 28.0;  //scaling the thesholds in backgroundDiff()
float low_thresh = 20.0;

// Counts number of images learned for averaging later
//
float Icount;

// I is just a sample image for allocation purposes
// (passed in for sizing)
//
void AllocateImages(const cv::Mat& I) {
    cv::Size sz = I.size();
    IavgF = cv::Mat::zeros(sz, CV_32FC3);
    IdiffF = cv::Mat::zeros(sz, CV_32FC3);
    IprevF = cv::Mat::zeros(sz, CV_32FC3);
    IhiF = cv::Mat::zeros(sz, CV_32FC3);
    IlowF = cv::Mat::zeros(sz, CV_32FC3);
    Icount = 0.00001; // Protect against divide by zero
    tmp = cv::Mat::zeros(sz, CV_32FC3);
    tmp2 = cv::Mat::zeros(sz, CV_32FC3);
    Imaskt = cv::Mat(sz, CV_32FC1);
}

// Learn the background statistics for one more frame
// I is a color sample of the background, 3-channel, 8u
//
void accumulateBackground(cv::Mat& I) {
    static int first = 1; // nb. Not thread safe
    I.convertTo(tmp, CV_32F); // convert to float
    if (!first) {
        IavgF += tmp;
        cv::absdiff(tmp, IprevF, tmp2);
        IdiffF += tmp2;
        Icount += 1.0;
    }
    first = 0;
    IprevF = tmp;
}

void setHighThreshold(float scale) {
    IhiF = IavgF + (IdiffF * scale);
    cv::split(IhiF, Ihi);
}

void setLowThreshold(float scale) {
    IlowF = IavgF - (IdiffF * scale);
    cv::split(IlowF, Ilow);
}

void createModelsfromStats() {
    IavgF *= (1.0 / Icount);
    IdiffF *= (1.0 / Icount);

    // Make sure diff is always something
    //
    IdiffF += cv::Scalar(1.0, 1.0, 1.0);
    setHighThreshold(high_thresh);
    setLowThreshold(low_thresh);
}

// Create a binary: 0,255 mask where 255 (red) means foreground pixel
// I      Input image, 3-channel, 8u
// Imask  Mask image to be created, 1-channel 8u
//
void backgroundDiff(
    cv::Mat& I,
    cv::Mat& Imask) {

    I.convertTo(tmp, CV_32F); // To float
    cv::split(tmp, Igray);

    // Channel 1
    //
    cv::inRange(Igray[0], Ilow[0], Ihi[0], Imask);

    // Channel 2
    //
    cv::inRange(Igray[1], Ilow[1], Ihi[1], Imaskt);
    Imask = cv::min(Imask, Imaskt);

    // Channel 3
    //
    cv::inRange(Igray[2], Ilow[2], Ihi[2], Imaskt);
    Imask = cv::min(Imask, Imaskt);

    // Finally, invert the results
    //
    Imask = 255 - Imask;
}

///////////////////
void help(char** argv) {
    cout << "\n"
        << "Train a background model on  the first <#frames to train on> frames of an incoming video, then run the model\n"
        << argv[0] << " <#frames to train on> <avi_path/filename>\n"
        << "For example:\n"
        << argv[0] << " 50 ../tree.avi\n"
        << "'A' or 'a' to adjust thresholds, esc, 'q' or 'Q' to quit"
        << endl;
}

void showForgroundInRed(char** argv, const cv::Mat &img) {
    cv::Mat rawImage;
    cv::split(img, Igray);
    Igray[2] = cv::max(mask, Igray[2]);
    cv::merge(Igray, rawImage);
    cv::imshow(argv[0], rawImage);
    cv::imshow("Segmentation", mask);
}

void adjustThresholds(char** argv, cv::Mat &img) {
    int key = 1;
    while ((key = cv::waitKey()) != 27 && key != 'Q' && key != 'q')  // Esc or Q or q to exit
    {
        if (key == 'L') { low_thresh += 0.2; }
        if (key == 'l') { low_thresh -= 0.2; }
        if (key == 'H') { high_thresh += 0.2; }
        if (key == 'h') { high_thresh -= 0.2; }
        cout << "H or h, L or l, esq or q to quit;  high_thresh = " << high_thresh << ", " << "low_thresh = " << low_thresh << endl;
        setHighThreshold(high_thresh);
        setLowThreshold(low_thresh);
        backgroundDiff(img, mask);
        showForgroundInRed(argv, img);
    }
}

////////////////////////////////////////////////////////////////
int main(int argc, char** argv) {
    cv::namedWindow(argv[0], cv::WINDOW_AUTOSIZE);
    cv::VideoCapture cap(0);
    /*if ((argc < 3) || !cap.open(argv[2])) {
        cerr << "Couldn't run the program" << endl;
        help(argv);
        cap.open(0);
        return -1;
    }*/
    int number_to_train_on = atoi(argv[1]);

    // FIRST PROCESSING LOOP (TRAINING):
    //
    int frame_count = 0;
    int key;
    bool first_frame = true;
    cout << "Total frames to train on = " << number_to_train_on << endl; //db
    while (1) {
        cout << "frame#: " << frame_count << endl;
        cap >> image;
        if (!image.data) exit(1); // Something went wrong, abort
        if (frame_count == 0) { AllocateImages(image); }
        accumulateBackground(image);
        cv::imshow(argv[0], image);
        frame_count++;
        if ((key = cv::waitKey(7)) == 27 || key == 'q' || key == 'Q' || frame_count >= number_to_train_on) break; //Allow early exit on space, esc, q
    }

    // We have accumulated our training, now create the models
    //
    cout << "Creating the background model" << endl;
    createModelsfromStats();
    cout << "Done!  Hit any key to continue into single step. Hit 'a' or 'A' to adjust thresholds, esq, 'q' or 'Q' to quit\n" << endl;

    // SECOND PROCESSING LOOP (TESTING):
    //
    cv::namedWindow("Segmentation", cv::WINDOW_AUTOSIZE); //For the mask image
    while ((key = cv::waitKey()) != 27 || key == 'q' || key == 'Q') { // esc, 'q' or 'Q' to exit
//    while (frame_count >= number_to_train_on) { // esc, 'q' or 'Q' to exit
        cap >> image;
        if (!image.data) exit(0);
        cout << frame_count++ << endl;
        backgroundDiff(image, mask);
        cv::imshow("Segmentation", mask);

        // A simple visualization is to write to the red channel
        //
        showForgroundInRed(argv, image);
        if (key == 'a') {
            cout << "In adjust thresholds, 'H' or 'h' == high thresh up or down; 'L' or 'l' for low thresh up or down." << endl;
            cout << " esq, 'q' or 'Q' to quit " << endl;
            adjustThresholds(argv, image);
            cout << "Done with adjustThreshold, back to frame stepping, esq, q or Q to quit." << endl;
        }
    }
    exit(0);
}