/*
 * =====================================================================================
 *
 *       Filename:  vision.cpp
 *
 *    Description:  Class for vision tasks such as opening images, plotting
 *    points in the image frame, performing measurements.
 *
 *        Version:  1.0
 *        Created:  04/07/2017 11:34:16 AM
 *       Revision:  none
 *       Compiler:  gcc
 *
 *         Author:  Martin Miller (MHM), miller7@illinois.edu
 *   Organization:  Aerospace Robotics and Controls Lab (ARC)
 *
 * =====================================================================================
 */
#include "vision.h"
#include "types.h"

/*
 *--------------------------------------------------------------------------------------
 *       Class:  Vision
 *      Method:  Vision
 * Description:  constructor
 *--------------------------------------------------------------------------------------
 */
Vision::Vision ()
{
    ;
}  /* -----  end of method Vision::Vision  (constructor)  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  Vision
 *      Method:  Vision :: open
 * Description:  Opens the given image, undistorts.
 *--------------------------------------------------------------------------------------
 */
void
Vision::open ( const char *fn, const Camera &cam )
{
    Mat bayered, unbayered, undistorted ;
    bayered = cv::imread(fn, -1);
    if (!bayered.data) {
        fprintf(stderr, "Could not read %s.\n", fn);
        exit(1);
    }
    cv::cvtColor(bayered, unbayered, CV_BayerBG2BGR ,3);
    cv::undistort( unbayered, undistorted, cam.K(1), cam.d(1));
    undistorted.convertTo(original, CV_32FC3);
    cv::cvtColor(original, gray, CV_BGR2GRAY);
    display = undistorted.clone();
    return ;
}		/* -----  end of method Vision::open  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  Vision
 *      Method:  Vision :: drawMeasurements
 * Description:  Draws the measured points onto the image.
 *--------------------------------------------------------------------------------------
 */
void
Vision::drawMeasurements ( const vector<measurement_t> &z, const Camera &cam,
        const cv::Scalar &colsrc, const cv::Scalar &colref, bool drawEllipse)
{
    cv::Scalar ellcol(0,0,200);
    for (auto i=z.begin(); i!=z.end(); ++i) {
        if (i->z_type==HEIGHT) continue;
        Vector3d xis, xir;
        xis = cam.body2img(i->source);
        xis /= xis[2];
        xir = cam.body2img(i->reflection);
        xir /= xir[2];
        cv::Point ps(xis[0],xis[1]);
        cv::Point pr(xir[0],xir[1]);

        cv::circle(display, ps, 3, colsrc,2);
        cv::circle(display, pr, 3, colref);

        if (drawEllipse) {
            double xs, ys, xr, yr;
            xs = 2*sqrt(i->Ssrc[0]);
            ys = 2*sqrt(i->Ssrc[1]);
            xr = 2*sqrt(i->Sref[0]);
            yr = 2*sqrt(i->Sref[1]);
            
            // Create rotated rect for ellipse
            cv::RotatedRect rrs(ps, cv::Size(2*INLIER_THRESHOLD*xs, 2*INLIER_THRESHOLD*ys), 0.);
            cv::RotatedRect rrr(pr, cv::Size(2*INLIER_THRESHOLD*xr, 2*INLIER_THRESHOLD*yr), 0.);
            try {
                cv::ellipse(display, rrs, ellcol );
                cv::ellipse(display, rrr, ellcol);
            } catch (...) {
                cerr << "caught and error drawing ellipse" << endl;
            }
        }
    }
    return ;
}		/* -----  end of method Vision::drawMeasurements  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  Vision
 *      Method:  Vision :: show
 * Description:  Show the display image.
 *--------------------------------------------------------------------------------------
 */
void
Vision::show (  )
{
    //cv::imshow("Image view", display);
    //cv::waitKey(1);
    return ;
}		/* -----  end of method Vision::show  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  Vision
 *      Method:  Vision :: acquireFeatures
 * Description:  Finds new features, assigns IDs, extracts descriptors. The
 * vector curz contains a list of existing features i.e. areas to mask off when
 * runing goodfeaturestotrack.
 *--------------------------------------------------------------------------------------
 */
void
Vision::acquireFeatures ( const Camera &cam, vector<measurement_t> &z, const vector<measurement_t> &curz )
{
    vector<cv::Point2f> corners;

    Mat mask = Mat::zeros(gray.size(),CV_8UC1);
    // mask off the margins of the image
    cv::Rect roi(50,50,gray.cols-100, gray.rows-100);
    mask(roi) = 255*Mat::ones(roi.size(),CV_8UC1);

    // mask off the existing points
    for (auto i=curz.begin(); i!=curz.end(); ++i) {
        Vector3d xi;
        xi = cam.body2img(i->source);
        cv::Point tl;
        tl.x = xi(0);
        tl.y = xi(1);
        cv::Size sz(50,50);
        cv::Rect box(tl,sz);
        cv::rectangle(mask,box,cv::Scalar(0),-1);
    }

    cv::goodFeaturesToTrack(gray, corners, 20, 0.1, 50, mask);
    for (auto i=corners.begin(); i!=corners.end(); ++i) {
        Vector3d xi;
        xi << i->x, i->y, 1.;
        bool tooClose = false;
        for (auto j=curz.begin(); j!=curz.end(); ++j) {
            Vector3d xci = cam.body2img(j->source);
            if ((xci-xi).norm()<25) {
                tooClose = true;
                break;
            }
        }
        if (tooClose) continue;
        measurement_t m;
        m.id = _id++;
        m.z_type = MONO;
        m.source = cam.img2body(xi);
        m.reflection << 0,0,1;
        getTemplate(m.patch,xi);
        if (m.patch.cols==0 || m.patch.rows==0) continue;
        z.push_back(m);
    }
    return ;
}		/* -----  end of method Vision::acquireFeatures  ----- */

void
Vision::getTemplate( Mat &p, const Vector3d &pt)
{
    cv::Rect imrect(0,0,gray.cols,gray.rows);
    assert (PATCHSIZE%2==1); // patch size must be odd
    cv::Point pti;
    pti.x = (int) pt(0);
    pti.y = (int)pt(1);
    cv::Point topLeft = pti - cv::Point(PATCHSIZE/2,PATCHSIZE/2);
    cv::Size sz(PATCHSIZE,PATCHSIZE);
    cv::Rect roi(topLeft, sz);
    roi &= imrect;
    p = gray(roi).clone();
}

void
Vision::measurements ( const Camera &cam, const vector<measurement_t> &zin,
        vector<measurement_t> &zout)
{
    for (auto i=zin.begin(); i!=zin.end(); ++i) {
        measurement_t z;
        z = *i;
        measurement_type mt = z.z_type;
        if (i->z_type==MONO || i->z_type==BOTH) {
            measurement_t zins, zouts;
            zins = *i;
            zins.z_type = MONO;
            measure(cam,zins, zouts);
            z.xcorrmax = zouts.xcorrmax;
            z.source = zouts.source;
            if (z.xcorrmax<xcorrsrc) { //downgrade
                if (mt==BOTH) {
                    mt=REFLECTION;
                } else if (mt==MONO) {
                    continue;
                }
            }
        }
        if (i->z_type==REFLECTION || i->z_type==BOTH) {
            measurement_t zinr, zoutr;
            zinr = *i;
            zinr.z_type = REFLECTION;
            measure(cam, zinr, zoutr);
            z.reflection = zoutr.reflection;
            if (zoutr.xcorrmax<xcorrref) { //downgrade
                if (mt==BOTH) {
                    mt = MONO;
                } else if (mt==REFLECTION) {
                    continue;
                }
            }
        }
        z.z_type = mt;
        zout.push_back(z);
    }
    return ;
}		/* -----  end of method Vision::measurements  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  Vision
 *      Method:  Vision :: measure
 * Description:  Performs a measurement on zin and stores it in zout.
 *--------------------------------------------------------------------------------------
 */
void
Vision::measure ( const Camera &cam, const measurement_t &zin, measurement_t &zout )
{
    assert((zin.z_type==MONO) || (zin.z_type==REFLECTION));
    // Create ROIs
    RotatedRect searchrr; // Rotated rect describing search ellipse
    Rect searchr; // Rect describing search ellipse
    searchRegion(cam, zin, searchrr);
    searchr = searchrr.boundingRect();

    // Add margin around search region
    Rect imrect(cv::Point(0,0), gray.size());
    cv::Point offset(PATCHSIZE/2,PATCHSIZE/2);
    cv::Size enlarge(2*(PATCHSIZE/2),2*(PATCHSIZE/2));
    Rect roi = (searchr-offset) + enlarge;
    if ((roi & imrect)==roi) {
        Mat sr = gray(roi);

        // Match template
        Mat result;
        result.create( sr.cols-zin.patch.cols+1, sr.rows-zin.patch.rows+1, CV_32FC1 );
        if (zin.z_type==MONO) {
            matchTemplate( sr, zin.patch, result, CV_TM_CCOEFF_NORMED);
        } else {
            matchTemplate( sr, zin.refpatch, result, CV_TM_CCOEFF_NORMED);
        }

        // Mask the ellipse around the result.
        Mat mask = Mat::zeros(result.size(), CV_8UC1);
        RotatedRect ellipserr;
        cv::Size sz = searchr.size();
        ellipserr.size = sz;
        ellipserr.center = cv::Point2f(0.5*sz.width,0.5*sz.height);
        ellipserr.angle = 0;
        cv::ellipse(mask, ellipserr, cv::Scalar(255), -1 );
        double minval, maxval;
        int minidx[2], maxidx[2];
        cv::minMaxIdx(result, &minval, &maxval, minidx, maxidx, mask);
        cv::Point maxpt(maxidx[1], maxidx[0]);

        cv::Point zpi = maxpt+searchr.tl();
        Vector3d zi;
        zi << zpi.x, zpi.y, 1;

        zout.id = zin.id;
        if (zin.z_type==MONO) {
            zout.source = cam.img2body(zi);
            zout.z_type = MONO;
        } else {
            zout.reflection = cam.img2body(zi);
            zout.z_type = REFLECTION;
        }
        zout.xcorrmax = maxval;

        /*
        Mat sr8, patch8, rpatch8;
        sr.convertTo(sr8, CV_8UC1);
        zin.patch.convertTo(patch8, CV_8UC1);
        zin.refpatch.convertTo(rpatch8, CV_8UC1);
        cv::circle(sr8, maxpt+offset, 4, cv::Scalar(255), 1);
        cv::Point spt(PATCHSIZE/2, PATCHSIZE/2);
        cv::circle(patch8, spt, 4, cv::Scalar(255), 1);
        cv::circle(display, maxpt+searchr.tl(), 4, cv::Scalar(20,240,0),2);
        cv::imshow("s region", sr8);
        cv::imshow("patch", patch8);
        cv::imshow("rpatch", rpatch8);
        cv::imshow("result", result);
        cv::imshow("mask", mask);
        cv::imshow("display2", display);
        cerr << maxval << endl;
        cv::waitKey(0);
        */
    } else {
        zout.xcorrmax = -1.;
    }
    return ;
}		/* -----  end of method Vision::measure  ----- */

void
Vision::searchRegion ( const Camera &cam, const measurement_t &z, cv::RotatedRect &rr )
{
    assert((z.z_type==MONO) || (z.z_type==REFLECTION));
    Vector3d xis;
    if (z.z_type==MONO) {
        xis = cam.body2img(z.source);
    } else {
        xis = cam.body2img(z.reflection);
    }
    //xis /= xis[2];
    cv::Point ps;
    ps.x = xis[0];
    ps.y = xis[1];

    double xs, ys;
    if (z.z_type==MONO) {
        xs = 2*sqrt(z.Ssrc[0]);
        ys = 2*sqrt(z.Ssrc[1]);
    } else {
        xs = 2*sqrt(z.Sref[0]);
        ys = 2*sqrt(z.Sref[1]);
    }
    
    // Create rotated rect for ellipse
    rr = cv::RotatedRect(ps, cv::Size(2*INLIER_THRESHOLD*xs, 2*INLIER_THRESHOLD*ys), 0.);

    return ;
}		/* -----  end of method Vision::searchRegion  ----- */

int Vision::_id = 1;