/**
 * OpenCV undistortion demonstration.
 *
 * Based on CvBase, see which.
 *
 * @author vona
 **/

#include "cvbase.hpp"

#include <iostream>
#include <iomanip>

using namespace std;

/**
 *  Extends base class to undistort given camera calibration in YAML format, as
 *  written by the OpenCV sample calibration program.
 **/
class CvUndistort : public CvBase {

  protected:

  /** Camera intrinsic parameters. **/
  CvMat *intrinsics;

  /** Camera distortion parameters. **/
  CvMat *distortions;

  /** Undistort map. **/
  IplImage *mapx, *mapy;

  /** Whether to undistort. **/
  bool undistort;

  /** The undistorted image. **/
  IplImage *undistortImage;


  /** Camera index is not optional on the command line. **/
  virtual bool camIndexOptional() { return false; }

  /** Adds command line help. **/
  virtual void cmdHelpExt() {
    cout<<"  where C is the YAML calibration file path"<<endl;
  }

  /** Adds GUI help. **/
  virtual void guiHelpExt() {
    cout<<"u -- toggle undistort on/off"<<endl;
  }

  /** Adds command line parameter help **/
  virtual std::string cmdHelpExtParams() { return "C "; }

  /** Undistorts each frame, if enabled. **/
  virtual IplImage *process(IplImage *frame) { 

    if (undistort) {
      
      if (undistortImage == 0) {
        //must allocate storage and compute the undistort map
        undistortImage = cvCloneImage(frame);
        mapx = cvCreateImage(cvGetSize(frame), IPL_DEPTH_32F, 1);
        mapy = cvCreateImage(cvGetSize(frame), IPL_DEPTH_32F, 1);
        cvInitUndistortMap(intrinsics, distortions, mapx, mapy);
      }
      
      cvRemap(frame, undistortImage, mapx, mapy);
      
      frame = undistortImage;
    }

    return frame;
  }

  /** Handles the undistort toggle key. **/
  virtual bool handleKeyExt(int code) {
    switch (code) {
    case 'u': {
      undistort = !undistort;
      cout<<"undistort "<<(undistort ? "on" : "off")<<endl;
      break;
    }
    }
    return true;
  }

  public:

  /** Consructor initializes members. **/
  CvUndistort() : CvBase("cvundistort") {
    undistort = false;
    undistortImage = 0;
    mapx = mapy = 0;
    intrinsics = distortions = 0;
  }

  /** Destructor releases memory. **/
  virtual ~CvUndistort() { 
    if (undistortImage != 0) cvReleaseImage(&undistortImage);
    if (mapx != 0) cvReleaseImage(&mapx);
    if (mapy != 0) cvReleaseImage(&mapy);
    if (intrinsics != 0) cvReleaseMat(&intrinsics);
    if (distortions != 0) cvReleaseMat(&distortions);
  }

  /** Parses base class plus extra args. **/
  virtual int initExt(int argc, char *argv[], int ate) {

    //must have at least one extra arg (plus the very first entry of argv which
    //was the invoked program name)
    if (argc < ate+2) { cmdHelp(); exit(-1); }
    const char *calfile = argv[++ate];
    
    CvFileStorage* fs = cvOpenFileStorage(calfile, 0, CV_STORAGE_READ);

    if (fs == 0) {
      cerr<<"error opening calibration file \""<<calfile<<"\""<<endl;
      exit(-1);
    }

    cout<<"reading camera calibration from \""<<calfile<<"\""<<endl;

    intrinsics = (CvMat*) cvReadByName(fs, 0, "camera_matrix");
    if (intrinsics == 0)
      { cerr<<"error reading camera_matrix"<<endl; exit(-1); }

    distortions = (CvMat*) cvReadByName(fs, 0, "distortion_coefficients");
    if (distortions == 0)
      { cerr<<"error reading distortion_coefficients"<<endl; exit(-1); }
    cvReshape(distortions, distortions, 1, 1); //ensure row vector

    cvReleaseFileStorage(&fs);

    ios_base::fmtflags fmtWas = cout.flags();
    cout<<scientific<<setprecision(3);

    cout<<"fx="<<CV_MAT_ELEM(*intrinsics, double, 0, 0)
        <<", fy="<<CV_MAT_ELEM(*intrinsics, double, 1, 1)
        <<", cx="<<CV_MAT_ELEM(*intrinsics, double, 0, 2)
        <<", cy="<<CV_MAT_ELEM(*intrinsics, double, 1, 2)<<endl;

    cout<<"k1="<<CV_MAT_ELEM(*distortions, double, 0, 0)
        <<", k2="<<CV_MAT_ELEM(*distortions, double, 0, 1)
        <<", p1="<<CV_MAT_ELEM(*distortions, double, 0, 2)
        <<", p2="<<CV_MAT_ELEM(*distortions, double, 0, 3)
        <<", k3="<<((distortions->cols > 4) ?
                    CV_MAT_ELEM(*distortions, double, 0, 4) : 0.0)<<endl;

    cout.flags(fmtWas);

    return ate;
  }
};

/** Program entry point. **/
int main(int argc, char *argv[]) {
  CvUndistort u;
  u.init(argc, argv);
  u.mainLoop();
}