#include "rocks.hpp"
#include "camera_transforms.hpp"
#include "camera_constants.hpp"
#include <stdio.h>
#include <algorithm>


using namespace std;
using namespace cv;

void drawRocks(Mat image, vector<Rock> &rocks, const Scalar &center_colour, const Scalar &circle_colour) {
    for (Rock rock : rocks) {
        Point2d center = realToPerspective(rock.center);
        double radius = rock.radius;

        /* Circle center */
        circle(image, center, 3, center_colour, -1, 8, 0);

        /* Circle radius */
        circle(image, center, radius, circle_colour, 3, 8, 0);
    }
}


/* Find rocks in a thresholded image */
void findThresholdedRocks(Mat image, RockColour colour, vector<Rock> &rocks)
{
    /* Find contours in the image, we'll use these to find rocks */
    vector< vector<Point> > contours;
    vector<Vec4i> hierarchy;
    findContours(image, contours, hierarchy, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);
    if (hierarchy.size() == 0) {
        return;
    }

    for (int idx = 0; idx >= 0; idx = hierarchy[idx][0]) {
        /* Check that contour appears to be rounded, hinting it's circular */
        if (contours[idx].size() < 10) {
            continue;
        }

        /* Undistort our points */
        vector<Point2d> conts;
        Mat(contours[idx]).copyTo(conts);

        vector<Point2d> undistorted;
        undistortPoints(conts, undistorted, CAMERA_MATRIX, DIST_COEFFS, Mat(), CAMERA_MATRIX);

        /* Convert contour points into real space */
        vector<Point2d> real_points;
        for (Point2d screen_point : undistorted) {
            Point2d real = screenToReal(screen_point);
            real_points.push_back(real);
        }

        /* Check that an ellipse covering the points is roughly circular */
        vector<Point2f> real_floats;
        Mat(real_points).copyTo(real_floats);
        RotatedRect rect = fitEllipse(real_floats);
        if (rect.size.height < rect.size.width / 1.2 || rect.size.height > rect.size.width * 1.2) {
            continue;
        }

        Point2f center;
        float radius;
        minEnclosingCircle(real_floats, center, radius);

        if (radius > ROCK_RADIUS_MIN && radius < ROCK_RADIUS_MAX) {
            Rock rock;
            rock.center = rect.center;
            rock.radius = radius / 5;
            rock.colour = colour;

            rocks.push_back(rock);
        }
    }
}


void findRocks(Mat background, Mat image, vector<Rock> &red_rocks, vector<Rock> &yellow_rocks)
{
    Mat subtracted;
    Mat hsv, hsv_background;
    cvtColor(image, hsv, CV_BGR2HSV);
    cvtColor(background, hsv_background, CV_BGR2HSV);

    absdiff(hsv_background, hsv, subtracted);

    Mat gray;
    cvtColor(subtracted, gray, CV_BGR2GRAY);
    GaussianBlur(gray, gray, Size(9, 9), 10, 10);
    GaussianBlur(hsv, hsv, Size(9, 9), 10, 10);

    Mat red, yellow;
    inRange(hsv, RED_LOW_HSV, RED_HIGH_HSV, red);
    inRange(hsv, YELLOW_LOW_HSV, YELLOW_HIGH_HSV, yellow);

    Mat thresholded;
    threshold(gray, thresholded, 20, 255, THRESH_BINARY);

    findThresholdedRocks(yellow & thresholded, YELLOW_ROCK, yellow_rocks);
    findThresholdedRocks(red & thresholded, RED_ROCK, red_rocks);

    /* Debug drawing */
    /*
    Mat undistorted;
    undistort(image, undistorted, CAMERA_MATRIX, DIST_COEFFS);

    drawRocks(undistorted, red_rocks, Scalar(0,0,255), Scalar(0,0,255));
    drawRocks(undistorted, yellow_rocks, Scalar(0,200,200), Scalar(0,200,200));

    imwrite("out-circles.jpg", undistorted);

    */

    subtracted.release();
    hsv.release();
    hsv_background.release();
    gray.release();
    red.release();
    yellow.release();
    thresholded.release();
}


static void rocksFromCircles(Mat image, RockColour colour, vector<Rock> &rocks)
{
    vector<Vec3f> circles;
    HoughCircles(image, circles, CV_HOUGH_GRADIENT, 1, 15, 1, 32, 30, 40);

    for (Vec3f circle : circles) {
        Rock rock;
        rock.center = perspectiveToReal(Point2d(circle[0], circle[1]));
        rock.radius = circle[2];
        rock.colour = RED_ROCK;

        rocks.push_back(rock);
    }
}


void findRocksCircles(Mat background, Mat image, vector<Rock> &red_rocks, vector<Rock> &yellow_rocks, bool debug, String output_file)
{
    Mat undistorted;
    undistort(image, undistorted, CAMERA_MATRIX, DIST_COEFFS, OPT_MATRIX);
    imwrite("border.jpg", undistorted);

    Mat flat_perspective;
    warpPerspective(undistorted, flat_perspective, PERSPECTIVE_MATRIX, Size(PERSPECTIVE_SIZE.width/3 + 500, PERSPECTIVE_SIZE.height/3 + 200));
    undistorted.release();

    Mat flat_hsv;
    cvtColor(flat_perspective, flat_hsv, CV_BGR2HSV);

    Mat red, yellow;
    inRange(flat_hsv, CIRCLES_RED_LOW_HSV, CIRCLES_RED_HIGH_HSV, red);
    inRange(flat_hsv, CIRCLES_YELLOW_LOW_HSV, CIRCLES_YELLOW_HIGH_HSV, yellow);

    vector<Mat> channels;
    split(flat_hsv, channels);
    flat_hsv.release();

    Mat blurred, blurred_red, blurred_yellow;
    medianBlur(channels[1] & red, blurred_red, 11);
    medianBlur(channels[1] & yellow, blurred_yellow, 11);
    medianBlur(channels[1], blurred, 11);

    red.release();
    yellow.release();
    for (Mat chan : channels) {
        chan.release();
    }

    rocksFromCircles(blurred_red, RED_ROCK, red_rocks);
    rocksFromCircles(blurred_yellow, YELLOW_ROCK, yellow_rocks);

    vector<Rock> possible_rocks;
    rocksFromCircles(blurred, YELLOW_ROCK, possible_rocks);

    blurred_red.release();
    blurred_yellow.release();
    blurred.release();

    drawRocks(flat_perspective, possible_rocks, Scalar(255,0,0), Scalar(255,0,0));
    drawRocks(flat_perspective, red_rocks, Scalar(0,0,255), Scalar(0,0,255));
    drawRocks(flat_perspective, yellow_rocks, Scalar(0,220,220), Scalar(0,220,220));

    if (debug) {
        imwrite(output_file, flat_perspective);
    }

    flat_perspective.release();
}


bool pointEq(Point2d p1, Point2d p2) {
    return p1 == p2;
}

/* Must be unique */
vector<Point2d> getProjectedCircle(size_t col, size_t row)
{
    Point2d real_point = screenToReal(Point2d(col, row));

    /* Get the projected circle points */
    vector<Point2d> points;
    for (Point2i circle_point : BASIC_CIRCLE) {
        Point2d pt = Point2d((double) circle_point.x, (double) circle_point.y);
        pt = real_point + (ROCK_RADIUS / BASIC_CIRCLE_RADIUS) * pt;

        points.push_back(realToScreen(pt));
    }

    unique(points.begin(), points.end(), pointEq);
    return points;
}


void addAccumulation(Mat &accumulation, size_t col, size_t row)
{
    Rect accum_rect(Point(), accumulation.size());
    vector<Point2d> circle_points = getProjectedCircle(col, row);

    for (Point2i point : circle_points) {
        if (accum_rect.contains(point)) {
            accumulation.at<uchar>(point.y, point.x) += 1;
        }
    }
}


Mat getSaturation(Mat& in){
    Mat out = Mat::zeros(in.size(), CV_8UC1);

    for(int row = 0; row < in.rows; ++row) {
        uchar* p = in.ptr(row);

        for(int col = 0; col < in.cols; ++col) {
            uchar s = max(p[0],max(p[1],p[2]))-min(p[0],min(p[1],p[2]));
            out.at<uchar>(row,col)=s;
            p += 3;
        }
    }

    return out;
}


/* Find rocks with a custom circle detector */
void findRocksCirclesFast(Mat background, Mat image, vector<Rock> &red_rocks, vector<Rock> &yellow_rocks)
{
    /* Subtract the background */
    Mat subtracted;
    absdiff(background, image, subtracted);

    /* Get saturation from the subtracted image */
    Mat saturation;
    saturation = getSaturation(subtracted);
    subtracted.release();

    /* Blur the subtracted, saturated image */
    Mat blurred;
    blur(saturation, blurred, Size(7,7));
    saturation.release();

    /* Perform edge detection on background subtracted image */
    Mat edges;
    Canny(blurred, edges, 1, 10);
    blurred.release();

    /* Build up the accumulation matrix for our Hough transform */
    Mat accumulation = Mat::zeros(image.size(), CV_8UC1);
    for (size_t row = 0; row < edges.rows; ++row) {
        for (size_t col = 0; col < edges.cols; ++col) {
            if (255 == edges.at<uchar>(row, col)) {
                addAccumulation(accumulation, col, row);
            }
        }
    }

    edges.release();

    imwrite("accumulation.jpg", accumulation);
    /* Find local maxima in the accumulation matrix, this gives us
       screen space coordinates for the rocks.
     */
    // vector<Point2i> screen_centers = findLocalMaxima(accumulation);

    accumulation.release();
}


vector<Point2i> BASIC_CIRCLE;
void setupBasicCircle() {
    int back_width = BASIC_CIRCLE_RADIUS * 2;

    int radius = BASIC_CIRCLE_RADIUS;
    Point2i center = Point2i(radius, radius);
    Mat back = Mat::zeros(back_width, back_width, CV_8UC1);

    circle(back, center, radius, Scalar(255,255,255));

    findNonZero(back, BASIC_CIRCLE);

    for (size_t i = 0; i < BASIC_CIRCLE.size(); ++i) {
        BASIC_CIRCLE[i] -= Point2i((double) radius, (double) radius);
    }
}


double rockDist(Rock r1, Rock r2)
{
    double dx = r1.center.x - r2.center.x;
    double dy = r1.center.y - r2.center.y;

    return sqrt(dx * dx + dy * dy);
}