/*
 * =====================================================================================
 *
 *       Filename:  state.cpp
 *
 *    Description:  A Class for managing body and features.
 *
 *        Version:  1.0
 *        Created:  03/17/2017 07:55:56 PM
 *       Revision:  none
 *       Compiler:  gcc
 *
 *         Author:  Martin Miller (MHM), miller7@illinois.edu
 *   Organization:  Aerospace Robotics and Controls Lab (ARC)
 *
 * =====================================================================================
 */
#include "state.h"

State::State ( )
{
    body = new Body;
    P = Matrix<double,Dynamic,Dynamic>::Zero(STATESIZE,STATESIZE);
    P.block<3,3>(STATESIZE-3,STATESIZE-3) = covbias*Matrix3d::Identity();
    return ;
}		/* -----  end of method State::State  ----- */

void
State::vel ( const Matrix<double,3,1> &v )
{
    body->vel(v);
    return ;
}		/* -----  end of method State::vel  ----- */

void
State::pos ( const UTM &utm )
{
    utm_c = utm.zone_c;
    utm_i = utm.zone_i;
    body->pos(utm);
    return ;
}		/* -----  end of method State::pos  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: H
 * Description:  Returns the matrix H for the measurements z
 *--------------------------------------------------------------------------------------
 */
MatrixXd
#if STATESIZE==13
State::H ( const vector<measurement_t> &z )
#else
State::H ( const Quaterniond &q, const vector<measurement_t> &z )
#endif
{
#if STATESIZE==13
    Quaterniond q = body->qhat();
#endif
    Vector3d pos;
    pos = body->ned();
    MatrixXd h;
    // Determine the size of H
    int cols = STATESIZE + 3*features.size();
    int rows = Hrows(z);
    h = MatrixXd::Zero(rows,cols);
    for (int i=0,row=0; i<z.size(); ++i) {
        int col;
        Feature *fi;
        switch ( z[i].z_type ) {
            case BOTH:	
                col = rowById(z[i].id);
                if (col==-1) {
                    fprintf(stderr, "Reflection feature %d not found, quitting.\n", z[i].id);
                    exit(1);
                }

                fi = featureById(z[i].id);

                h.block<6,STATESIZE>(row,0) = fi->Hx(pos,q);
                h.block<6,3>(row,col) = fi->Hy(pos,q);
                row += 6;
                break;

            case MONO:	
                col = rowById(z[i].id);
                if (col==-1) {
                    fprintf(stderr, "Feature %d not found, quitting.\n", z[i].id);
                    exit(1);
                }
                fi = featureById( z[i].id );
                h.block<4,STATESIZE>(row,0) = fi->Hx(pos,q).block<4,STATESIZE>(0,0);
                h.block<4,3>(row,col) = fi->Hy(pos,q).block<4,3>(0,0);
                row += 4;
                break;

            case REFLECTION:
                col = rowById(z[i].id);
                if (col==-1) {
                    fprintf(stderr, "Feature %d not found, quitting.\n", z[i].id);
                    exit(1);
                }
                fi = featureById( z[i].id );
                h.block<4,STATESIZE>(row,0) = fi->Hx(pos,q).block<4,STATESIZE>(2,0);
                h.block<4,3>(row,col) = fi->Hy(pos,q).block<4,3>(2,0);
                row += 4;
                break;

            case HEIGHT:	
                h.block<1,STATESIZE>(row,0) = body->H();
                row += 1;
                break;

            default:	
                break;
        }				/* -----  end switch  ----- */
    }
    return h;
}		/* -----  end of method State::H  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: Hrows
 * Description:  Returns the number of rows in the measurement vector based on
 * the number and types of measurements.
 *--------------------------------------------------------------------------------------
 */
int
State::Hrows ( const vector<measurement_t> &z) 
{
    int rows = 0;
    for (auto i=z.begin(); i!=z.end(); ++i) {
        switch ( i->z_type ) {
            case BOTH:
                rows += 6;
                break;

            case REFLECTION:	
                rows += 4;
                break;

            case MONO:	
                rows += 4;
                break;

            case HEIGHT:	
                rows += 1;
                break;

            default:	
                break;
        }				/* -----  end switch  ----- */
    }
    return rows ;
}		/* -----  end of method State::Hrows  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: R
 * Description:  Returns the measurement noise matrix.
 *--------------------------------------------------------------------------------------
 */
MatrixXd
State::R ( const vector<measurement_t> &z )
{
    int rows = Hrows(z);
    MatrixXd r;
    r = MatrixXd::Zero(rows,rows);

    for (int i=0,row=0; i<z.size(); ++i) {
        Feature *fi;
        switch ( z[i].z_type ) {
            case BOTH:	
                fi = featureById(z[i].id);
                r.block<6,6>(row,row) = fi->R();
                row += 6;
                break;

            case MONO:	
                fi = featureById(z[i].id);
                r.block<4,4>(row,row) = fi->R().block<4,4>(0,0);
                row += 4;
                break;

            case REFLECTION:
                fi = featureById(z[i].id);
                r.block<4,4>(row,row) = fi->R().block<4,4>(2,2);
                row += 4;
                break;

            case HEIGHT:	
                r.block<1,1>(row,row) = body->R();
                row += 1;
                break;

            default:	
                break;
        }				/* -----  end switch  ----- */
    }

    return r ;
}		/* -----  end of method State::R  ----- */

MatrixXd
#if STATESIZE==13
State::partialUpdate( const MatrixXd &h, const MatrixXd &S, 
        const vector<measurement_t> &z)
#else
State::partialUpdate( const MatrixXd &h, const MatrixXd &S, 
        const vector<measurement_t> &z, const Quaterniond &q )
#endif
{
    MatrixXd K;
    // P^T is implied here since P is symmetric
    K = S.fullPivHouseholderQr().solve(h*P).transpose();

    // Compute the innovation or error
    Matrix<double,Dynamic,1> y;
#if STATESIZE==13
    y = innovation(z);
#else
    y = innovation(z,q);
#endif
    
    // Get the update
    Matrix<double,Dynamic,1> dx;
    dx = K*y;
    if (dx.segment<3>(3).norm()>1) {
        cerr << "dx: " << dx.head<STATESIZE>().transpose() << endl;
        return K;
    }
    assert (dx.rows()==STATESIZE+3*features.size());
    body->dx(dx.segment<STATESIZE>(0));
    {
        int row=STATESIZE;
        for (auto i=features.begin(); i!=features.end(); ++i, row+=3) {
            (*i)->dx(dx.segment<3>(row));
        }
    }
    return K;
}

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: kalmanUpdate
 * Description:  Performs the kalman update. 
 *--------------------------------------------------------------------------------------
 */
void
#if STATESIZE==13
State::kalmanUpdate( const MatrixXd &h, const MatrixXd &S, 
        const vector<measurement_t> &z)
#else
State::kalmanUpdate( const MatrixXd &h, const MatrixXd &S, 
        const vector<measurement_t> &z, const Quaterniond &q)
#endif
{
    MatrixXd K;
#if STATESIZE==13
    K = partialUpdate(h,S,z);
#else
    K = partialUpdate(h,S,z,q);
#endif


    // Compute P_k|k
    P = (MatrixXd::Identity(P.rows(),P.rows())-K*h)*P;
    P = 0.5*(P+P.transpose());
    
    return ;
}		/* -----  end of method State::kalmanUpdate  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: innovation
 * Description:  Returns the innovation vector z-hhat
 *--------------------------------------------------------------------------------------
 */
Matrix<double,Dynamic,1>
#if STATESIZE==13
State::innovation( const vector<measurement_t> &z)
#else
State::innovation( const vector<measurement_t> &z, const Quaterniond &q)
#endif
{
#if STATESIZE==13
    Quaterniond q = body->qhat();
#endif
    Matrix<double,Dynamic,1> y;
    int rows = Hrows(z);
    y = Matrix<double,Dynamic,1>::Zero(rows,1);

    for (int i=0,row=0; i<z.size(); ++i) {
        Feature *fi;
        measurement_type mt=z[i].z_type;
        if (mt==BOTH) {
            fi = featureById(z[i].id);
            Matrix<double,6,1> hi, zi;
            hi = fi->h(body->ned(), q);
            zi = fi->measurement_vector(z[i].source, z[i].reflection);
            y.segment<6>(row) = zi-hi;
            row += 6;
        } else if (mt==MONO) {
            fi = featureById(z[i].id);
            Matrix<double,4,1> hi, zi;
            hi = fi->h(body->ned(), q).segment<4>(0);
            zi = fi->measurement_vector(z[i].source);
            y.segment<4>(row) = zi-hi;
            row += 4;
        } else if (mt==REFLECTION) {
            fi = featureById(z[i].id);
            Matrix<double,4,1> hi, zi;
            hi = fi->h(body->ned(), q).segment<4>(2);
            zi = fi->measurement_vector(z[i].source, z[i].reflection).segment<4>(2);
            y.segment<4>(row) = zi-hi;
            row += 4;
        } else if (mt==HEIGHT) {
            Matrix<double,1,1> zi, hi;
            zi[0] = z[i].height;
            Vector3d xbw;
            xbw = body->ned();
            hi[0] = xbw[2];

            y.segment<1>(row) = zi-hi;
            row += 1;
        } else {
            fprintf(stderr, "Unknown measurement type\n");
            exit(1);
        }
    }

    return y;
}

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: handle_measurements
 * Description:  Update EKF, and adds new features to state.
 *--------------------------------------------------------------------------------------
 */
void
#if STATESIZE==13
State::handle_measurements ( const vector<measurement_t> &z)
#else
State::handle_measurements ( const vector<measurement_t> &z, const Quaterniond &q)
#endif
{
    vector<measurement_t> featuresToAdd;
    vector<measurement_t> featuresToUpdate;
    double zmeas = body->ned()[2];
    for (auto i=z.begin(); i!=z.end(); ++i) {
        if (i->z_type==HEIGHT) {
            featuresToUpdate.push_back(*i);
            zmeas = i->height;
        } else if (exists(i->id)) {
            Feature *ft;
            ft = featureById(i->id);
#ifdef INLIERTEST
#if STATESIZE==13
            Quaterniond q = body->qhat();
#endif
            if (ft->isInlier(*i, Pxx(), Pxy(i->id), Pyy(i->id), body->ned(), q, INLIER_THRESHOLD)) 
#endif
            {
                featuresToUpdate.push_back(*i);
            } 
        } else {
            featuresToAdd.push_back(*i);
        }
    }
    // Remove features that don't have measurements in this timestep.
    auto i=features.begin();
    int feati=0;
    while (i!=features.end()) {
        bool measured=false;
        for (auto j=featuresToUpdate.begin(); j!=featuresToUpdate.end(); ++j) {
            if (j->z_type==HEIGHT) {
                continue;
            } else if (j->id==(*i)->id()) {
                measured=true;
                break;
            }
        }
        if (measured==true) {
            ++i;
            ++feati;
        } else {
            i=removeFeature(i,feati);
        }
    }
    if (featuresToUpdate.size()>1) {
#ifdef  DORANSAC
#if STATESIZE==13
        ransacUpdate(featuresToUpdate);
#else
        ransacUpdate(featuresToUpdate,q);
#endif
#else      /* -----  not DORANSAC  ----- */
        MatrixXd h;
#if STATESIZE==13
        h = H(featuresToUpdate);
#else
        h = H(q,featuresToUpdate);
#endif
        MatrixXd S;
        S = h*P*h.transpose() + R(featuresToUpdate);
#if STATESIZE==13
        kalmanUpdate(h,S,featuresToUpdate);
#else
        kalmanUpdate(h,S,featuresToUpdate,q);
#endif
#endif     /* -----  not DORANSAC  ----- */
    }
#if STATESIZE==13
    addFeatures( featuresToAdd, zmeas);
#else
    addFeatures( featuresToAdd, q, zmeas);
#endif

    return ;
}		/* -----  end of method State::feature_update  ----- */

MatrixXd
#if STATESIZE==13
State::blockSI ( const vector<measurement_t> &z)
#else
State::blockSI ( const vector<measurement_t> &z, const Quaterniond &q )
#endif
{
#if STATESIZE==13
    Quaterniond q = body->qhat();
#endif
    fprintf(stderr, "blockSI not implemented\n");
    exit(1);
    int rows = Hrows(z);
    MatrixXd SI;
    SI = MatrixXd::Zero(rows,rows);
    {
        int row = 0;
        for (auto i=z.begin(); i!=z.end(); ++i) {
            if (i->z_type==HEIGHT) {
                Matrix<double,1,1> s;
                s = body->S(Pxx());
                SI.block<1,1>(row,row) = s.inverse();
                row += 1;
            } else if (i->z_type==REFLECTION) {
                Feature *f;
                Matrix<double,6,6> s;
                f = featureById(i->id);
                s = f->S(Pxx(), Pxy(i->id), Pyy(i->id), body->ned() ,q);
                SI.block<6,6>(row,row) = s.inverse();
                row += 6;
            } else if (i->z_type==MONO) {
                /*
                Feature *f;
                Matrix<double,6,6> s;
                f = featureById(i->id);
                s = f->S(Pxx(), Pxy(i->id), Pyy(i->id), body->ned() ,q);
                s = s.block<4,4>(0,0);
                SI.block<4,4>(row,row) = s.inverse();
                row += 4;
                */
            } else {
                fprintf(stderr, "Unrecognized feature type, quitting\n");
                exit(1);
            }
        }
    }

    return SI ;
}		/* -----  end of method State::blockSI  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: Pxy
 * Description:  Return the covariance block for feature with the given id.
 *--------------------------------------------------------------------------------------
 */
Matrix<double,STATESIZE,3>
State::Pxy ( int id )
{
    int col = rowById(id);
    return P.block<STATESIZE,3>(0,col) ;
}		/* -----  end of method State::Pxy  ----- */

Matrix<double,3,3>
State::Pyy ( int id )
{
    int row = rowById(id);
    return P.block<3,3>(row,row) ;
}		/* -----  end of method State::Pyy  ----- */

void
#if STATESIZE==13
State::addFeatures ( vector<measurement_t> &F, double z)
#else
State::addFeatures(vector<measurement_t> &F, const Quaterniond &q, double z)
#endif
{
#if STATESIZE==13
    Quaterniond q =body->qhat();
#endif
    // Add new features
    Vector3d pos = body->ned();
    for (auto i=F.begin(); i!=F.end(); ++i) {
        if (features.size()>MAXFEATURES) break;
        if (std::find(blacklist.begin(),blacklist.end(), i->id)!=blacklist.end()) continue;
        // Create feature
        Feature *f;
        if (i->z_type==BOTH) {
#ifdef  INITDEPTH
            f = new Feature(i->id, i->source, i->reflection, pos, q, z, i->patch);
#else      /* -----  not INITDEPTH  ----- */
            f = new Feature(i->id, i->source, pos, q, i->patch);
#endif     /* -----  not INITDEPTH  ----- */
        } else if (i->z_type==MONO) {
            f = new Feature(i->id, i->source, pos, q, i->patch);
        } else {
            fprintf(stderr, "Measurement type: %d not supported by addFeatures, quitting\n", i->z_type);
            exit(1);
        }
        if (!f->sane()) {
            delete f;
            continue;
        }
        features.push_back(f);

        // Expand P
        expandP(f->P0(i->z_type));
    }
    return ;
}		/* -----  end of method State::addFeatures  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: initializePi
 * Description:  Zeros out the off-diagonal blocks and sets the Pii block to Pi.
 *--------------------------------------------------------------------------------------
 */
    void
State::initializePi ( int i, const Matrix<double,3,3> &Pi )
{
    int pt=STATESIZE+3*(i);
    P.block(pt,0,3,P.cols()) = MatrixXd::Zero(3,P.cols());
    P.block(0,pt,P.rows(),3) = MatrixXd::Zero(P.rows(),3);
    P.block<3,3>(pt,pt) = Pi;

    return ;
}		/* -----  end of method State::initializePi  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: Q
 * Description:  Compose full Q from body and features.
 *--------------------------------------------------------------------------------------
 */
MatrixXd
State::Q ( double dt )
{
    MatrixXd q;
    int rows = STATESIZE + 3*features.size();
    q = MatrixXd::Zero(rows,rows);
    q.topLeftCorner<STATESIZE,STATESIZE>() = body->Q(dt);
    { // limit i's scope
        auto i = features.begin();
        for (int row=STATESIZE; row<rows; row+=3, ++i) {
            q.block<3,3>(row,row) = (*i)->Q(dt);
        }
    }
    return q;
}		/* -----  end of method State::q  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: F
 * Description:  Compose full F from body and features
 *--------------------------------------------------------------------------------------
 */
MatrixXd
#if STATESIZE==13
State::F ( const Vector3d &w, double dt )
#else
State::F ( const Quaterniond &q, const Vector3d &w, double dt )
#endif
{
    Vector3d v;
    v = body->vel();
    // Allocate matrix F
    MatrixXd f;
    int rows = STATESIZE + 3*features.size();
    f = MatrixXd::Zero(rows,rows);

    // Set body F
#if STATESIZE==13
    f.topLeftCorner<STATESIZE,STATESIZE>() = body->F(w,dt);
#else
    f.topLeftCorner<STATESIZE,STATESIZE>() = body->F(w,q,dt);
#endif
    // Set Fxi Fyi
    { // limit i's scope
        auto i = features.begin();
        for (int row=STATESIZE; row<rows; row+=3,++i) {
            f.block<3,STATESIZE>(row,0) = (*i)->Fx(dt);
            f.block<3,3>(row,row) = (*i)->Fy(v,w,dt);
        }
    }

    return f ;
}		/* -----  end of method State::F  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: Pkk1
 * Description:  Updates P_k|k-1
 *--------------------------------------------------------------------------------------
 */
void
#if STATESIZE==13
State::Pkk1 ( const Vector3d &ang, const double dt)
#else
State::Pkk1 ( const Vector3d &ang, const Quaterniond &q, const double dt)
#endif
{
    MatrixXd f,fullQ;
    fullQ = Q(dt);
#if STATESIZE==13
    f = F ( ang, dt );
#else
    f = F ( q, ang, dt );
#endif
    P = f*P*f.transpose()+fullQ;
    // Enforce symmetry
    P = 0.5*(P+P.transpose());
    return ;
}		/* -----  end of method State::Pkk1  ----- */

void
#if STATESIZE==13
State::motionModel ( const Vector3d &acc, const Vector3d &ang, const double dt)
#else
State::motionModel ( const Vector3d &acc, const Vector3d &ang,
    const Quaterniond &q, const double dt)
#endif
{
#if STATESIZE==13
    Pkk1(ang,dt);
    body->motionModel(acc,ang,dt);
#else
    Pkk1(ang,q,dt);
    body->motionModel(acc,ang,q,dt);
#endif

    Vector3d vel;
    vel = body->vel();
    
#ifdef  FASTMOTIONMODEL
    Matrix<double,Dynamic,6> fullL;
    fullL = L();

    Matrix<double,6,1> V;
    V << vel[1], vel[2], vel[0], ang[1], ang[2], ang[0];

    Matrix<double,Dynamic,1> f;
    f = fullL*V*dt;

    // Update features' states
    {
        int rows = 3*features.size();
        auto i = features.begin();
        for (int row=0; row<rows; row+=3, ++i) {
            Vector3d dx;
            dx = f.segment<3>(row);
            (*i)->dx(dx);
        }
    }
    
#else      /* -----  not FASTMOTIONMODEL  ----- */
    for (auto i=features.begin(); i!=features.end(); ++i) {
        (*i)->motionModel(ang,vel,dt);
    }
#endif     /* -----  not FASTMOTIONMODEL  ----- */

    // Remove features that leave FOV
    {
        auto i = features.begin();
        int j=0;
        while (i!=features.end()) {
            if ((*i)->inFOV()) {
                ++i;
                ++j;
            } else {
                i=removeFeature(i,false);
            }
        }
    }


    return ;
}		/* -----  end of method State::motionModel  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: iById
 * Description:  Returns the the feature position in the features vector.
 * Returns -1 if not found.
 *--------------------------------------------------------------------------------------
 */
int
State::iById ( int id )
{
    int i=0;
    for (auto f=features.begin(); f!=features.end(); ++f, ++i) {
        if ((*f)->id()==id) return i;
    }
    return -1;
}		/* -----  end of method State::iById  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: removeFeature
 * Description:  Remove the Feature i and Pj in a clean fashion.
 *--------------------------------------------------------------------------------------
 */
std::list<Feature *>::iterator
State::removeFeature ( std::list<Feature *>::iterator &i, bool bl )
{
    // It is important to shrink P first, because it depends on knowing the
    // current size of the feature vector.
    int j = iById((*i)->id());
    assert (j!=-1);
    shrinkP(j);
    if (bl) { // add to blacklist
        blacklist.push_back((*i)->id());
    }
    delete *i;
    i = features.erase(i);
    return i;
}		/* -----  end of method State::removeFeature  ----- */

std::list<Feature *>::iterator
State::removeFeature ( int id, bool bl )
{
    int j = iById(id);
    assert (j!=-1);
    shrinkP(j);
    if (bl) {
        blacklist.push_back(id);
    }
    Feature *f = featureById(id);
    auto i = std::find(features.begin(), features.end(), f);
    delete *i;
    i = features.erase(i);
    return i;
}		/* -----  end of method State::removeFeature  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: L
 * Description:  Return the composition of all Li's of the features
 *--------------------------------------------------------------------------------------
 */
Matrix<double,Dynamic,6>
State::L (  )
{
    Matrix<double,Dynamic,6> l;
    int rows = 3*features.size();
    l = Matrix<double,Dynamic,6>::Zero(rows,6);
    {
        auto i=features.begin();
        for (int row=0; row<rows; row+=3, ++i) {
            l.block<3,6>(row,0) = (*i)->L();
        }
    }
    return l;
}		/* -----  end of method State::L  ----- */

void
State::position_covariance ( const Matrix<double,3,3> &p )
{
    P.block(0,0,3,3) = p;
    return ;
}		/* -----  end of method State::position_covariance  ----- */

void
State::velocity_covariance ( const Matrix<double,3,3> &p )
{
    P.block(3,3,3,3) = p;
    return ;
}		/* -----  end of method State::velocity_covariance  ----- */

void
State::unicsv ( const double time )
{
    body->unicsv(time);
    printf(",%d\n", features.size());
#ifdef FEATUREMAP
    for (auto feature:features ) {
        feature->unicsv(time,body->ned(),body->qhat());
    }
#endif
    return ;
}		/* -----  end of method State::unicsv  ----- */

void
State::accelerometer_bias ( const Vector3d &b)
{
    body->accelerometer_bias(b);
    return ;
}		/* -----  end of method State::accelerometer_bias  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: exists
 * Description:  Tests if the id is a current feature.
 *--------------------------------------------------------------------------------------
 */
bool
State::exists ( int id )
{
    for (auto i=features.begin(); i!=features.end(); ++i) {
        if ((*i)->id()==id) return true;
    }
    return false;
}		/* -----  end of method State::exists  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: rowById
 * Description:  Returns the first row of the feature with the given ID. Returns
 * -1 if not found.
 *--------------------------------------------------------------------------------------
 */
int
State::rowById ( int id )
{
    int j=STATESIZE;
    for (auto i=features.begin(); i!=features.end(); ++i, j+=3) {
        if ((*i)->id()==id) return j;
    }
    return -1;
}		/* -----  end of method State::rowById  ----- */

Feature *
State::featureById ( int id )
{
    for (auto i=features.begin(); i!=features.end(); ++i) {
        if ((*i)->id()==id) return *i;
    }
    return NULL;
}		/* -----  end of method State::featureById  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: expandP
 * Description:  Expand the P matrix to fit one more feature.
 *--------------------------------------------------------------------------------------
 */
void
State::expandP ( const Matrix3d &Pi)
{
    P.conservativeResizeLike(MatrixXd::Zero(3+P.cols(),3+P.cols()));
    P.bottomRightCorner<3,3>() = Pi;
    return ;
}		/* -----  end of method State::expandP  ----- */

Matrix<double,STATESIZE,STATESIZE>
State::Pxx ( )
{
    return P.topLeftCorner<STATESIZE,STATESIZE>() ;
}		/* -----  end of method State::Pxx  ----- */

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: shrinkP
 * Description:  Delete the ith feature from P
 *--------------------------------------------------------------------------------------
 */
void
State::shrinkP ( int i )
{
    int N = STATESIZE + 3*features.size();
    int I = STATESIZE + 3*i;

    P.block(I,0,N-I-3,I) = P.bottomLeftCorner(N-I-3,I);
    P.block(0,I,I,N-I-3) = P.topRightCorner(I,N-I-3);
    P.block(I,I,N-I-3,N-I-3) = P.bottomRightCorner(N-I-3,N-I-3);
    P.conservativeResize(N-3,N-3);

    return ;
}		/* -----  end of method State::shrinkP  ----- */

    void
#if STATESIZE==13
State::ransacUpdate ( vector<measurement_t> &z )
#else
State::ransacUpdate ( vector<measurement_t> &z, const Quaterniond &q )
#endif
{
#if STATESIZE==13
    Quaterniond q = body->qhat();
#endif
    Matrix<double,Dynamic,1> mu_old; // x_k|k-1
    mu_old = asVector();
    // Randomize the measurements
    std::random_shuffle ( z.begin(), z.end() );

    int N0 = z.size();
    int N = z.size();
    vector<measurement_t> inliers;
    for (int i=0; i<N && i<N0; ++i) {
        measurement_t zi = z[i];
        vector<measurement_t> Z;
        Z.push_back(zi);
        MatrixXd h;
#if STATESIZE==13
        h = H(Z);
#else
        h = H(q,Z);
#endif
        MatrixXd S;
        S = h*P*h.transpose() + R(Z);
#if STATESIZE==13
        partialUpdate(h,S,Z);
#else
        partialUpdate(h,S,Z,q);
#endif

        for (auto j=z.begin(); j!=z.end(); ++j) {
            if (j->id==zi.id) continue;
            if (j->z_type==HEIGHT) {
                Z.push_back(*j);
                continue;
            }
            Feature *ft = featureById(j->id);
            if (ft->isRansacInlier(*j, body->ned(), q)) {
                Z.push_back(*j);
            }
        }
        if (Z.size()>inliers.size()) {
            inliers.clear();
            for (auto j=Z.begin(); j!=Z.end(); ++j) {
                inliers.push_back(*j);
            }
            double eps = 1. - (double) inliers.size() / (double) z.size();
            //N = ceil(log(0.01)/log(eps));
        }
        // Restore original state before trying again.
        asVector(mu_old);
    }

    // Perform a full update using the inliers.
    MatrixXd h;
#if STATESIZE==13
    h = H(inliers);
#else
    h = H(q,inliers);
#endif
    MatrixXd S;
    S = h*P*h.transpose() + R(inliers);
#if STATESIZE==13
    kalmanUpdate(h, S, inliers);
#else
    kalmanUpdate(h, S, inliers, q);
#endif

    // High innovation update
    vector<measurement_t> hi_inliers;
    for (auto i=z.begin(); i!=z.end(); ++i) {
        bool found = false;
        for (auto j=inliers.begin(); j!=inliers.end(); ++j) {
            if (i->id==j->id) {
                found = true;
                break;
            }
        }
        if (!found) {
            Feature *ft = featureById(i->id);
            if (ft->isInlier(*i, Pxx(), Pxy(i->id), Pyy(i->id), body->ned(), q, ransac_hi_threshold)) {
                hi_inliers.push_back(*i);
            } else {
                removeFeature(i->id, true);
            }
        }
    }

#if STATESIZE==13
    h = H(hi_inliers);
#else
    h = H(q,hi_inliers);
#endif
    S = h*P*h.transpose() + R(hi_inliers);
#if STATESIZE==13
    kalmanUpdate(h, S, hi_inliers);
#else
    kalmanUpdate(h, S, hi_inliers, q);
#endif

    /*
    cout << "Z: " << z.size()
         << " LI: " << inliers.size()
         << " HI: " << hi_inliers.size() << endl;
         */

    return ;
}		/* -----  end of method State::ransacUpdate  ----- */

Matrix<double,Dynamic,1>
State::asVector ( )
{
    Matrix<double,Dynamic,1> m;
    int rows = STATESIZE + 3*features.size();
    m = Matrix<double,Dynamic,1>::Zero(rows,1);
    m.head<STATESIZE>() = body->asVector();
    int row = STATESIZE;
    for (auto i=features.begin(); i!=features.end(); ++i) {
        m.segment<3>(row) = (*i)->asVector();
        row += 3;
    }

    return m ;
}		/* -----  end of method State::asVector  ----- */

void
State::asVector ( const Matrix<double,Dynamic,1> &m )
{
    int rows = STATESIZE + 3*features.size();
    assert (rows==m.rows());
    body->asVector(m.head<STATESIZE>());

    int row = STATESIZE;
    for (auto i=features.begin(); i!=features.end(); ++i) {
        (*i)->asVector(m.segment<3>(row));
        row += 3;
    }
    return ;
}		/* -----  end of method State::asVector  ----- */

#if STATESIZE==13
Quaterniond
State::qhat ( )
{
    return body->qhat() ;
}		/* -----  end of method State::qhat  ----- */

void
State::qhat ( const Quaterniond &q )
{
    body->qhat(q);
    return ;
}		/* -----  end of method State::qhat  ----- */


void
State::quaternion_covariance ( const Matrix<double,4,4> &covq )
{
    P.block<4,4>(6,6) = covq;
    return ;
}		/* -----  end of method State::quaternion_covariance  ----- */

#endif

/*
 *--------------------------------------------------------------------------------------
 *       Class:  State
 *      Method:  State :: featuresAsMeasurements
 * Description:  Writes the feature locations to a measurement_t vector in body
 * coordinates.
 *--------------------------------------------------------------------------------------
 */
void
#if STATESIZE==13
State::featuresAsMeasurements ( vector<measurement_t> &yk, const Camera &cam )
#else
State::featuresAsMeasurements ( vector<measurement_t> &yk, const Camera &cam,
        const Quaterniond &q)
#endif
{
#if STATESIZE==13
    Quaterniond q = body->qhat();
#endif
    for (auto i=features.begin(); i!=features.end(); ++i) {
        // Get the predicted measurement
        Matrix<double,6,1> h;
        h = (*i)->h(body->ned(), q);

        Vector3d pbs, pbr;
        pbs << h(0), h(1), 1;
        pbr << h(4), h(5), 1;

        measurement_t z;
        z.id = (*i)->id();
        z.source = (*i)->p2x(pbs);
        z.reflection = (*i)->p2x(pbr);
        z.z_type = (domono) ? MONO : BOTH;
        z.patch = (*i)->warpedPatch(cam, q);
        z.refpatch = (*i)->reflectedPatch(cam, q);

        // Get ellipse
        Matrix<double,6,6> s;
        s = (*i)->S(Pxx(), Pxy((*i)->id()), Pyy((*i)->id()), body->ned(), q);
	z.Ssrc[0] = s(0,0);
	z.Ssrc[1] = s(1,1);
	z.Sref[0] = s(4,4);
	z.Sref[1] = s(5,5);

        yk.push_back(z);
    }
    return ;
}		/* -----  end of method State::featuresAsMeasurements  ----- */

void
#if STATESIZE==13
State::doMeasurements ( vector<measurement_t> &z, Vision &viz, Camera &cam )
#else
State::doMeasurements ( vector<measurement_t> &z, Vision &viz,
        const Quaterniond &q, Camera &cam)
#endif
{
    measurement_t hgt = z[0];
    assert (hgt.z_type==HEIGHT);
    double hgtmeas = hgt.height;
    
    // Measure existing features
    vector<measurement_t> currentFeatures;
#if STATESIZE==13
    featuresAsMeasurements(currentFeatures,cam);
#else
    featuresAsMeasurements(currentFeatures,cam,q);
#endif

    vector<measurement_t> newFeatures;
    if (features.size()<MINFEATURES) {
        viz.acquireFeatures(cam, newFeatures,currentFeatures);
    }
    viz.measurements(cam, currentFeatures, z);
    // Check for flipped correspondence pairs and flop them
    for (auto i=z.begin(); i!=z.end(); ++i) {
        if (i->z_type==BOTH && i->source[2]>i->reflection[2]) {
            // Flop the measurements
            Vector3d tmp;
            tmp = i->source;
            i->source = i->reflection;
            i->reflection = tmp;

            // Get a new patch
            Vector3d pi;
            cv::Mat patch;
            pi = cam.body2img(i->source);
            viz.getTemplate(patch, pi);

            // Reinitialize feature
#if STATESIZE==13
            Quaterniond q = body->qhat();
#endif
            Feature *ft = featureById(i->id);
            ft->reinitialize(i->source,i->reflection,body->ned(),q,patch);

            // reinitialize the covariance
            int j = iById(i->id);
            initializePi(j,ft->P0(BOTH));
        }
    }

#ifdef DORANSAC
#if STATESIZE==13
    ransacUpdate(z);
#else
    ransacUpdate(z,q);
#endif
#else      /* -----  not DORANSAC  ----- */
        MatrixXd h;
#if STATESIZE==13
        h = H(z);
#else
        h = H(q,z);
#endif
        MatrixXd S;
        S = h*P*h.transpose() + R(z);
#if STATESIZE==13
        kalmanUpdate(h,S,z);
#else
        kalmanUpdate(h,S,z,q);
#endif
#endif
    // Reset lastseen counter
    for (auto i=z.begin(); i!=z.end(); ++i) {
        if (i->z_type==HEIGHT) continue;
        Feature *ft = featureById(i->id);
        if (ft==NULL) {
            cerr << "id: " << i->id << "is null." << endl;
        } else {
            ft->seen();
        }
    }

    // Remove features that haven't been measured recently
    auto i=features.begin();
    while (i!=features.end()) {
        if ((*i)->since(5)) {
            ++i;
        } else {
            i=removeFeature(i,false);
        }
    }

    // Initialize new features
    if (features.size()<MINFEATURES) {
#if STATESIZE==13
        addFeatures(newFeatures, hgtmeas);
#else
        addFeatures(newFeatures, q, hgtmeas);
#endif
    }
    vector<measurement_t> updatedFeatures;
#if STATESIZE==13
    featuresAsMeasurements(updatedFeatures, cam);
#else
    featuresAsMeasurements(updatedFeatures, cam, q);
#endif
    int numboth = 0;
    for (auto i=z.begin(); i<z.end(); ++i) {
        if (i->z_type==BOTH) ++numboth;
    }
    cerr << z.size() << " " << numboth << endl;
    if (numboth<10) {
        xcorrsrc *= 0.9;
        xcorrref *= 0.9;
    } else if (numboth>12) {
        xcorrsrc *= 1.1;
        xcorrref *= 1.1;
    }
    //cerr << z.size() << endl;
    viz.drawMeasurements(z, cam, cv::Scalar(0,0,200), cv::Scalar(0,0,200), false);
    viz.drawMeasurements(updatedFeatures, cam, cv::Scalar(0,255,0), cv::Scalar(0,200,0), true);
    viz.show();

    return ;
}		/* -----  end of method State::doMeasurements  ----- */