/*
* @Author: felicitywang
* @Date:   2016-08-25 11:57:30
* @Email:  cnfxwang@gmail.com
* @Last Modified time: 2016-08-25 16:14:07
*/



/**
 * final version of demonstration demo
 * with well-structured GUI and iteraction
 * image processing only mode
 * driving mode
 * TODO algorithms: LK and Farneback
 */

// TODO GPU
// TODO parameter setting...
// with opencv 3
// TODO 代码复用 重构 太不简洁！！！


#include <iostream>
#include <opencv2/opencv.hpp>
#include "eyebot.h"
#include <sstream>

using namespace cv;
using namespace std;


// TODO test best wait key delay time
// delay time for waitKey(ms)
#define wait_key_delay 15


// where to display on LCD screen
#define lcd_row 0


// thresholds for psds
#define left_thre 2100
#define right_thre 2100
#define middle_thre 2800

// save pictures to file if running on ssh
#define ON_SSH

// iteration time for the for loop
// not used in the new GUI
// #define iter_time 300

// const for eyebot vomega functions
// TODO find best
#define curve_dist 35
#define straight_dist 40
#define lin_speed 50
#define turn_speed 25  // used in VWTurn()
#define turn_ang 45   // unit: degree
#define curve_speed 25  // used in VWCurve()

// ignore some too-some optical flow for obstacle
// TODO how to define obstacle threshold
#define obstacle_thre 0.6

// whether should turn
#define turn_thre 0.7

// use avg, sum and obstacle altogether
#define w_obst  0.5
#define w_sum   0.25
#define w_avg   0.25

// for eyebot turn function
#define LEFT  true
#define RIGHT false

// for driving/showing mode

#define mode_driving 0
#define mode_showing 1


/**
 * distance between two points
 * opencv norm not as fast
 */
double dist(Point2f a, Point2f b) {
    return sqrt((a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y));
}

/**
 * whether to turn according to delta and threshold
 * turn when |delta_turn| > OBSTABLE_THRE
 * @param delta_turn
 * @return true if should turn, false otherwise
 */
bool should_turn(double delta_turn) {
    return delta_turn < -turn_thre || delta_turn > turn_thre;
}


/**
 * stop the eyebot when obstacles too near or at the end of the program
 */
void VWStop() {
    VWStraight(0, 0);
    VWDriveWait();
    VWSetSpeed(0, 0);
    VWSetPosition(0, 0, 0);
    VWDriveWait();
    VWSetSpeed(lin_speed, turn_speed);
}


/**
 * stop and go backwards if motor(s) stalled
 */
void check_stalled() {
    if (VWStalled() > 0) {
        VWStop();
        VWStraight(-1 * straight_dist, lin_speed);
        VWDriveWait();
    }
}

/**
 * go backwards
 * turn around (360 degrees)
 */
void VWReverse() {
    VWStop();
    VWStraight(-2 * straight_dist, lin_speed);
    VWDriveWait();
    VWTurn(3140, 8 * turn_speed);
    VWDriveWait();
}

/**
 * exit the program with key trigger KEY4
 */
void VWExit() {
    VWStop();
    VWSetSpeed(0, 0);
    VWSetPosition(0, 0, 0);
    VWDriveWait();
    LCDClear();
    LCDSetFont(COURIER, FONT_NORMAL);
    LCDSetFontSize(10);
    LCDSetPrintf(lcd_row + 4, 0, "Exiting the program....");
    usleep(2000);
    exit(-1);
}

/**
 * whether to use psd functions
 * can be turned on / off on homepage
 */
bool g_psd_on = false;

/**
 * detect with infrared light sensor for front, left, right
 * use psd when necessary
 * @return ture if actions changed by psd, false otherwise
 */
bool use_psd() {

    if (!g_psd_on)
        return false;

    int psd_left = PSDGetRaw(1);
    int psd_middle = PSDGetRaw(2);
    int psd_right = PSDGetRaw(3);

    // TODO: Huge Problem: always stalled
//    check_stalled();

    if (psd_middle > middle_thre) {
        LCDSetFont(COURIER, FONT_NORMAL);
        LCDSetPrintf(lcd_row + 11, 0, "front too close    ");
        LCDSetPrintf(lcd_row + 12, 0, "reverse            ");
        VWReverse();
        usleep(500);
        return true;
    } else if (psd_left > left_thre) {
        LCDSetFont(COURIER, FONT_NORMAL);
        LCDSetPrintf(lcd_row + 11, 0, "left  too close");
        LCDSetPrintf(lcd_row + 12, 0, "turning right...");
        VWTurn((int) (-turn_ang * 3.14 / 1.8), turn_speed);
        VWDriveWait();
        usleep(200);
        return true;
    } else if (psd_right > right_thre) {
        LCDSetFont(COURIER, FONT_NORMAL);
        LCDSetPrintf(lcd_row + 11, 0, "right too close");
        LCDSetPrintf(lcd_row + 12, 0, "turning  left...");
        VWTurn((int) (turn_ang * 3.14 / 1.8), turn_speed);
        VWDriveWait();
        usleep(200);
        return true;
    } else {
        LCDSetFont(COURIER, FONT_NORMAL);
        LCDSetPrintf(lcd_row + 11, 0, "PSD tested fine");
        LCDSetPrintf(lcd_row + 12, 0, "                ");
        return false;
    }
}


void turn(int mode, float angle) {

    if (angle > 0) {
        LCDSetFont(COURIER, FONT_NORMAL);
        LCDSetPrintf(lcd_row + 5, 0, "left     (*)");
        LCDSetPrintf(lcd_row + 6, 0, "straight ( )");
        LCDSetPrintf(lcd_row + 7, 0, "right    ( )");
    } else {
        LCDSetFont(COURIER, FONT_NORMAL);
        LCDSetPrintf(lcd_row + 5, 0, "left     ( )");
        LCDSetPrintf(lcd_row + 6, 0, "straight ( )");
        LCDSetPrintf(lcd_row + 7, 0, "right    (*)");
    }
    LCDSetFont(COURIER, FONT_NORMAL);
    LCDSetPrintf(lcd_row + 9, 0, "angle = %d   ", (int) (angle * 1.8 / 3.14));


    if (!use_psd() && mode == mode_driving) {
        VWCurve(curve_dist, (int) (angle * 3.14 / 1.8), curve_speed);
        usleep(500);
    }
}

/**
 * print going straight information on screen
 * @param mode
 */
void go_straight(int mode) {
    LCDSetFont(COURIER, FONT_NORMAL);
    LCDSetPrintf(lcd_row + 5, 0, "left     ( )");
    LCDSetPrintf(lcd_row + 6, 0, "straight (*)");
    LCDSetPrintf(lcd_row + 7, 0, "right    ( )");
    LCDSetPrintf(lcd_row + 9, 0, "angle =  0");
    if (!use_psd() && mode == mode_driving) {
        VWStraight(straight_dist, lin_speed);
    }
}

/**
 * information printed on lcd screen for main page
 */
void lcd_print_main() {

    LCDClear();

    LCDSetFont(COURIER, FONT_NORMAL);
    LCDSetFontSize(10);

    LCDSetPrintf(lcd_row + 1, 0, "Press *SHOW*  to show images");
    LCDSetPrintf(lcd_row + 3, 0, "Press *DRIVE* to start driving");
    LCDSetPrintf(lcd_row + 5, 0, "Press *PSD*   to switch psd mode");
    LCDSetPrintf(lcd_row + 7, 0, "Press *EXIT*  to exit the program");
    if (g_psd_on) {
        LCDSetFont(COURIER, FONT_NORMAL);
        LCDSetFontSize(10);
        LCDSetPrintf(lcd_row + 9, 0, "PSD on     ");
    } else {
        LCDSetFont(COURIER, FONT_NORMAL);
        LCDSetFontSize(10);
        LCDSetPrintf(lcd_row + 9, 0, "PSD off     ");
    }

    LCDMenu((char *) "SHOW", (char *) "DRIVE", (char *) "PSD", (char *) "EXIT");
}

/**
 * information printed on lcd screen for show and drive page
 */
void lcd_print() {
    LCDClear();

    LCDSetFont(COURIER, FONT_NORMAL);
    LCDSetFontSize(10);

    LCDSetPrintf(lcd_row + 0, 0, "Press *LK*     to use Lucas-Kanade Pyramid algorithm ");
    LCDSetPrintf(lcd_row + 3, 0, "Press *FB*     to use Farneback algorithm ");
    LCDSetPrintf(lcd_row + 5, 0, "Press *RETURN* to return to the welcome page");
    LCDSetPrintf(lcd_row + 7, 0, "Press *EXIT*   to exit the program");
    LCDMenu((char *) "LK", (char *) "FB", (char *) "RETURN", (char *) "EXIT");
}

/**
 * show image of drive with LK method
 */
void LK(int mode) {

    int width, height, half_width;
    BYTE *img;
    LCDImageStart(160, 0, 320, 240);

    // TODO 160*120 for driving
//    if (mode == mode_driving) {
//        width = QQVGA_X;
//        height = QQVGA_Y;
//        half_width = width / 2;
//        img = new BYTE[QQVGA_SIZE];
//        IPSetSize(QQVGA);
//        LCDImageSize(QQVGA);
//    } else { // mode_showing
    // 320*240
    width = QVGA_X;
    height = QVGA_Y;
    half_width = width / 2;
    img = new BYTE[QVGA_SIZE];
    IPSetSize(QVGA);
    LCDImageSize(QVGA);
//    }


    // ignore some too-long optical flow
    // TODO how to define this threshold
    const int of_mod_thre = width / 4;

    // parameters
    TermCriteria term_criteria(TermCriteria::COUNT | TermCriteria::EPS, 20, 0.03);

    Size sub_pix_win_size(width / 30, width / 30), win_size(width / 20, width / 20);
    const int max_level = 4;
    const int max_corners = 1000;  // basically 50 on raspberry pi with 320*240 resolution
    const double quality_level = 0.01;
    const int min_distance = 10;
    const int block_size = 3;

    // init camera and set to resolution
    VideoCapture cap(0);
    if (!cap.isOpened()) {
        cout << "Cannot open the web cam" << endl;
        exit(-1);
    }

    cap.set(CV_CAP_PROP_FRAME_WIDTH, width);
    cap.set(CV_CAP_PROP_FRAME_HEIGHT, height);

    Mat src, curr_gray, prev_gray;
    vector <Point2f> prev_points, next_points;

    // init prev_gray and prev_points before loop
    if (!cap.read(src)) {
        cout << "cannot read from camera" << endl;
        exit(1);
    }
    cvtColor(src, prev_gray, COLOR_BGR2GRAY);
    goodFeaturesToTrack(prev_gray, prev_points, max_corners, quality_level, min_distance, Mat(), block_size);
    cornerSubPix(prev_gray, prev_points, sub_pix_win_size, Size(-1, -1), term_criteria);

    // drive a little bit so that the first image in the loop is different
    if (mode == mode_driving) {
        VWSetSpeed(lin_speed, turn_speed);
        VWStraight(straight_dist, lin_speed);
    }

    int frame_ind = 0;

    LCDClear();
    LCDMenu((char *) "", (char *) "", (char *) "RETURN", (char *) "EXIT");


//    for (int count = 0; count < iter_time; count++) {
    while (true) {  // not yet to exit

        int key_code = KEYRead();
        // TODO keys
        if (key_code == KEY4)
            VWExit();
        if (key_code == KEY3) {
            VWStop();
            return;
        }
        // showing images

        // the first window isn't at the moveWindow position, somehow
        // not necessary if use LCDImage()
        // if (frame_ind == 2) {
        //     LCDClear();
        //     LCDMenu((char *) "", (char *) "", (char *) "RETURN", (char *) "EXIT");
        // }

        // get image and turn to grayscale
        if (!cap.read(src)) {
            cout << "cannot read from camera" << endl;
            exit(1);
        }
        cvtColor(src, curr_gray, COLOR_BGR2GRAY);

        // draw a white line to separate the screen left and right
        line(src, Point(half_width, 0), Point(half_width, height), Scalar(255, 255, 255), 1, 8);

        // compute optical flow motion field
        if (!prev_points.empty()) {
            vector <uchar> status;  // whether corner found
            vector<float> err;

            calcOpticalFlowPyrLK(prev_gray, curr_gray, prev_points, next_points, status, err, win_size, max_level,
                                 term_criteria);

            /**
             * @algorithm
             * first calculate number, sum(of moduli) of all the optical flows on each side
             * then calculate the same for potential obstacles
             * then adopt balance strategy with the result
             */
            int num_left = 0;
            int num_right = 0;
            double sum_left = 0.0;
            double sum_right = 0.0;

            vector <Point2f> optical_flows;  // prev_corner -> curr_corner if curr_corner exists
            vector<double> moduli;  // modulus of every optical flow
            vector<bool> pos;  // whether on the LEFT or RIGHT side of the screen
            vector<bool> is_obstacle;  // whether the optical flow belongs to an obstacle according to its modulus via the avg of the opposite side of screen

            for (size_t i = 0; i < next_points.size(); i++) {
                // some corners lost
                if (!status[i]) {
                    // num_lost++;
                    continue;
                }

                double modulus = dist(prev_points[i], next_points[i]);

                // ignore optical flow if modulus too big(?)
                // draw the yellow
                if (modulus > of_mod_thre) {
                    circle(src, next_points[i], 3, Scalar(0, 255, 255), -1, 8);
                    line(src, prev_points[i], next_points[i], Scalar(0, 255, 255), 1, 8);
                }

                    // otherwise categorize left / right and save
                    // draw them green(left) / blue(right)
                    // circle of the optical flow indicates the current corner
                else {
                    optical_flows.push_back(
                            Point2f(next_points[i].x - prev_points[i].x, next_points[i].y - prev_points[i].y));
                    moduli.push_back(modulus);
                    // left
                    if (next_points[i].x < half_width) {
                        pos.push_back(LEFT);
                        num_left++;
                        sum_left += modulus;
                        circle(src, next_points[i], 3, Scalar(0, 255, 0), -1, 8);
                        line(src, prev_points[i], next_points[i], Scalar(0, 255, 0), 1, 8);
                    } else {  // right
                        pos.push_back(RIGHT);
                        num_right++;
                        sum_right += modulus;
                        circle(src, next_points[i], 3, Scalar(255, 0, 0), -1, 8);
                        line(src, prev_points[i], next_points[i], Scalar(255, 0, 0), 1, 8);
                    }
                }
            }


            // recompute corners for every frame as they may be easily lost if motion too fast
            next_points.clear();
            goodFeaturesToTrack(curr_gray, prev_points, max_corners, quality_level, min_distance, Mat(),
                                block_size);
            cornerSubPix(curr_gray, prev_points, sub_pix_win_size, Size(-1, -1), term_criteria);
            prev_gray = curr_gray.clone();


            // now try to find obstacle from corners
            // draw them red


            // avoid divided by 0 error
            double avg_left = 0.0001;
            double avg_right = 0.0001;
            if (num_left > 0)
                avg_left = sum_left / num_left;
            if (num_right > 0)
                avg_right = sum_right / num_right;

            double sum_left_obstacles = 0;
            double sum_right_obstacles = 0;
            int num_left_obstacles = 0;
            int num_right_obstacles = 0;

            // whether each corner is obsatcle
            // TODO what for...
            for (size_t i = 0; i < optical_flows.size(); i++) {
                double u = pos[i] == LEFT ? 1.0 - avg_right / moduli[i] : 1.0 - avg_left / moduli[i];

                if (u > obstacle_thre) {
                    circle(src, next_points[i], 4, Scalar(0, 0, 255), -1, 8);
                    if (pos[i] == LEFT) {
                        sum_left_obstacles += u;
                        num_left_obstacles++;
                    } else {
                        sum_right_obstacles += u;
                        num_right_obstacles++;
                    }
                }
            }


            // Balance strategy
            // TODO weight
            // sum / avg / obstacle

            double delta_sum = (sum_left - sum_right) / (sum_left + sum_right);
            double delta_avg = (avg_left - avg_right) / (avg_left + avg_right);
            double delta_obstacle = 0;
            if (sum_left_obstacles + sum_right_obstacles > 0)
                delta_obstacle =
                        (sum_left_obstacles - sum_right_obstacles) / (sum_left_obstacles + sum_right_obstacles);

            // double left_turn = sum_left * w_sum + avg_left * w_avg + sum_right_obstacles * w_obst;
            // double right_turn = sum_right * w_sum + avg_right * w_avg + sum_right_obstacles * w_obst;
            // double delta_turn = (left_turn - right_turn) / (left_turn + right_turn);

            double delta_turn = delta_sum * w_sum + delta_avg * w_avg + delta_obstacle * w_obst;

            int next_ori_turn = (int) (half_width * (1 + delta_turn));

            LCDSetFont(COURIER, FONT_NORMAL);
            LCDSetPrintf(lcd_row + 0, 0, "frame %d", ++frame_ind);
            LCDSetPrintf(lcd_row + 2, 0, "delta_turn ");
            LCDSetPrintf(lcd_row + 3, 0, "=%10lf", delta_turn);


            if (!should_turn(delta_turn)) {  // go straight
                line(src, Point(next_ori_turn, 0), Point(next_ori_turn, height), Scalar(203, 192, 255), 2, 8);

                // show picture before moving
                // moveWindow("LK", 155, 0);
                // imshow("LK", src);

                /**
 * save image as tmp file
 * and call IPReadFile() to display
 * TODO make faster directly with src.data
 */

                // save all the frames one by one if ON_SSH
#ifdef ON_SSH
                stringstream stream;
                stream << "/home/pi/usr/software/optical_flow/pic/";
                int k = frame_ind;
                int number = 5;
                while (k > 0) {
                    number--;
                    k /= 10;
                }
                for (int k = 0; k < number; k++)
                    stream << "0";
                stream << frame_ind << ".jpg";
                string file_name;
                stream >> file_name;
                imwrite(file_name, src);
#endif


                imwrite("/home/pi/usr/software/optical_flow/tmp.ppm", src);
                IPReadFile("/home/pi/usr/software/optical_flow/tmp.ppm", img);
                // cout << "read successful" << endl;
                LCDImage(img);

                if (mode == mode_showing)
                    go_straight(mode_showing);
                else
                    go_straight(mode_driving);


            } else {  // turn
                line(src, Point(next_ori_turn, 0), Point(next_ori_turn, height), Scalar(0, 0, 255), 2, 8);

//                moveWindow("LK", 155, 0);
//                imshow("LK", src);

                imwrite("/home/pi/usr/software/optical_flow/tmp.ppm", src);
                IPReadFile("/home/pi/usr/software/optical_flow/tmp.ppm", img);

                LCDImage(img);

                float curve_angle = -90 * (delta_turn);
                if (mode == mode_showing)
                    turn(mode_showing, curve_angle);
                else
                    turn(mode_driving, curve_angle);
            }
        }

        if (waitKey(wait_key_delay) == 27)
            exit(-1);

    }

    VWStop();
}


void show_LK() {
    LK(mode_showing);
}

void drive_LK() {
    LK(mode_driving);
}


/**
 * show image of drive with FB method
 */
void FB(int mode) {

}

// TODO
void drive_FB() {
    FB(mode_driving);
}

// TODO
void show_FB() {
    FB(mode_showing);
}

void show() {
    lcd_print();
    int key_code = KEYGet();
    while (key_code != KEY4) {
        if (key_code == KEY1) {
            LCDSetFont(COURIER, FONT_NORMAL);
            LCDSetFontSize(10);
            LCDSetPrintf(lcd_row + 10, 0, "Now show Lucas-Kanade Pyramid method");
            usleep(2000);
            show_LK();
            lcd_print();
        } else if (key_code == KEY2) {
            LCDSetFont(COURIER, FONT_NORMAL);
            LCDSetFontSize(10);
            LCDSetPrintf(lcd_row + 10, 0, "Now show Farneback method");
            usleep(2000);
            show_FB();
            lcd_print();
        } else if (key_code == KEY3) {
            // not exactly necessary, but still
            VWStop();
            return;
        }
        key_code = KEYRead();
    }
    VWExit();
}

void drive() {
    lcd_print();

    int key_code = KEYGet();
    while (key_code != KEY4) {
        if (key_code == KEY1) {
            LCDSetFont(COURIER, FONT_NORMAL);
            LCDSetFontSize(10);
            LCDSetPrintf(lcd_row + 10, 0, "Lucas-Kanade Pyramid Algorithm");
            usleep(2000);
            drive_LK();
            lcd_print();
        } else if (key_code == KEY2) {
            LCDSetFont(COURIER, FONT_NORMAL);
            LCDSetFontSize(10);
            LCDSetPrintf(lcd_row + 10, 0, "Farneback Algorithm");
            usleep(2000);
            drive_FB();
            lcd_print();
        } else if (key_code == KEY3) {
            VWStop();
            return;
        }
        key_code = KEYRead();
    }
    VWExit();
}


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


    // write cout information to a file
// #ifndef ON_SSH
    freopen("/home/pi/usr/software/optical_flow/final_log", "w", stdout);
// #endif
    lcd_print_main();

    int key_code = KEYGet();

    while (key_code != KEY4) {  // not yet to exit

        if (key_code == KEY1) {  // show
            show();
            lcd_print_main();
        } else if (key_code == KEY2) {  // drive
            drive();
            lcd_print_main();
        } else if (key_code == KEY3) {  // switch PSD mode
            g_psd_on = !g_psd_on;
            if (g_psd_on) {
                LCDSetFont(COURIER, FONT_NORMAL);
                LCDSetFontSize(10);
                LCDSetPrintf(lcd_row + 9, 0, "PSD on     ");
            } else {
                LCDSetFont(COURIER, FONT_NORMAL);
                LCDSetFontSize(10);
                LCDSetPrintf(lcd_row + 9, 0, "PSD off    ");
            }
        }
        key_code = KEYRead();
    }

    VWExit();

}