/********************************************************************
 **                Image Component Library (ICL)                    **
 **                                                                 **
 ** Copyright (C) 2006-2012 CITEC, University of Bielefeld          **
 **                         Neuroinformatics Group                  **
 ** Website: www.iclcv.org and                                      **
 **          http://opensource.cit-ec.de/projects/icl               **
 **                                                                 **
 ** File   : ICLIO/src/ImageUndistortion.cpp                        **
 ** Module : ICLIO                                                  **
 ** Authors: Christian Groszewski                                   **
 **                                                                 **
 **                                                                 **
 ** Commercial License                                              **
 ** ICL can be used commercially, please refer to our website       **
 ** www.iclcv.org for more details.                                 **
 **                                                                 **
 ** GNU General Public License Usage                                **
 ** Alternatively, this file may be used under the terms of the     **
 ** GNU General Public License version 3.0 as published by the      **
 ** Free Software Foundation and appearing in the file LICENSE.GPL  **
 ** included in the packaging of this file.  Please review the      **
 ** following information to ensure the GNU General Public License  **
 ** version 3.0 requirements will be met:                           **
 ** http://www.gnu.org/copyleft/gpl.html.                           **
 **                                                                 **
 ** The development of this software was supported by the           **
 ** Excellence Cluster EXC 277 Cognitive Interaction Technology.    **
 ** The Excellence Cluster EXC 277 is a grant of the Deutsche       **
 ** Forschungsgemeinschaft (DFG) in the context of the German       **
 ** Excellence Initiative.                                          **
 **                                                                 **
 *********************************************************************/
#include <ICLIO/ImageUndistortion.h>
#include <ICLUtils/ConfigFile.h>
#include <ICLUtils/XML.h>
#include <fstream>

namespace icl{

  struct ImageUndistortion::Impl{
    virtual ~Impl(){}
    Img32f warpMap;
    std::string model;
    std::vector<double> params;
    Size imageSize;
    virtual Point32f undistort(const Point32f &point) const{
      const double &x0 = params[0];
      const double &y0 = params[1];
      const double &f = params[2]/100000000.0;
      const double &s = params[3];
      
      float x = s*(point.x-x0);
      float y = s*(point.y-y0);
      float p = 1 - f * (x*x + y*y);
      float xd = (p*x + x0);
      float yd = (p*y + y0);
      return Point32f(xd,yd);
    }
  };
  
  struct MatlabModel5ParamsImpl : public ImageUndistortion::Impl{
    FixedMatrix<icl64f,3,3> Kinv;
    void updateKinv(){
      Kinv[0] = params[0]; 
      Kinv[1] = params[0]*params[4]; 
      Kinv[2] = params[2];
      Kinv[3] = 0.0; 
      Kinv[4] = params[1]; 
      Kinv[5] = params[3];
      Kinv[6] = 0.0; 
      Kinv[7] = 0.0; 
      Kinv[8] = 1.0;
      Kinv = Kinv.inv();
    }
    virtual Point32f undistort(const Point32f &point) const{
      FixedMatrix<icl64f,1,3> p; p[0] = point.x; p[1] = point.y; p[2] = 1.0;
      FixedMatrix<icl64f,1,3> rays = Kinv*p;
      
      FixedMatrix<icl64f,1,2> x;
      x[0] = rays[0]/rays[2];
      x[1] = rays[1]/rays[2];
      
      FixedMatrix<icl64f,1,5> k;
      k[0] = params[5]; k[1] = params[6]; k[2] = params[7]; k[3] = params[8]; k[4] = params[9];
      // Add distortion:
      FixedMatrix<icl64f,1,2> pd;
      
      double r2 = x[0]*x[0]+x[1]*x[1];
      double r4 = r2*r2;
      double r6 = r4*r2;
      
      // Radial distortion:
      double cdist = 1 + k[0] * r2 + k[1] * r4 + k[4] * r6;
      
      pd[0] = x[0] * cdist;
      pd[1] = x[1] * cdist;
      
      // tangential distortion:
      double a1 = 2*x[0]*x[1];
      double a2 = r2 + 2*x[0]*x[0];
      double a3 = r2 + 2*x[1]*x[1];
      
      FixedMatrix<icl64f,1,2> delta_x;
      delta_x[0]= k[2]*a1 + k[3]*a2 ;
      delta_x[1]= k[2] * a3 + k[3]*a1;
      
      pd = pd + delta_x;
      // Reconvert in pixels:
      double px2 = params[0]*(pd[0]+params[4]*pd[1])+params[2];
      double py2 = params[1]*pd[1]+params[3];
      
      return Point32f(px2,py2);
    }
  };
  
  ImageUndistortion::ImageUndistortion():impl(0){};
  
  ImageUndistortion::ImageUndistortion(const std::string &model, const std::vector<double> &params, const Size &imageSize){
    if(model == "SimpleARTBased"){
      impl = new Impl;
      impl->params = params;
    }else if(model == "MatlabModel5Params"){
      MatlabModel5ParamsImpl *impl = new MatlabModel5ParamsImpl;
      this->impl = impl;
      impl->params = params;
      impl->updateKinv();
    }else{
      throw ICLException("ImageUndistortion(...): invalid model name");
    }
    impl->model = model;
    impl->imageSize = imageSize;
    impl->warpMap.detach();
    
  }

  ImageUndistortion::ImageUndistortion(const ImageUndistortion &other):impl(0){
    *this=other;
  }
  ImageUndistortion &ImageUndistortion::operator=(const ImageUndistortion &other){
    ICL_DELETE(impl);
    if(other.impl){
      if(other.impl->model == "SimpleARTBased") this->impl = new Impl(*other.impl);
      else this->impl = new MatlabModel5ParamsImpl(*reinterpret_cast<MatlabModel5ParamsImpl*>(other.impl));
    }
    return *this;
  }


  ImageUndistortion::ImageUndistortion(const std::string &filename):impl(0){
    std::ifstream s(filename.c_str());
    s >> (*this);
  }
  
  const Point32f ImageUndistortion::operator()(const Point32f &p) const{
    ICLASSERT_THROW(!isNull(), ICLException("unable to use function call operator () on a null-ImageUndistortion instance"));
    return impl->undistort(p);
  }
  
  const Size &ImageUndistortion::getImageSize() const{
    ICLASSERT_THROW(!isNull(), ICLException("unable to query the image size from a null-ImageUndistortion instance"));
    return impl->imageSize;
  }
  const std::vector<double> &ImageUndistortion::getParams() const{
    ICLASSERT_THROW(!isNull(), ICLException("unable to query the undistortion parameters from a null-ImageUndistortion instance"));
    return impl->params;

  
  }
  const std::string &ImageUndistortion::getModel() const{
    ICLASSERT_THROW(!isNull(), ICLException("unable to query the model from a null-ImageUndistortion instance"));
    return impl->model;
  }

  std::istream &operator>>(std::istream &is, ImageUndistortion &dest){
    XMLDocument *doc = new XMLDocument;
    doc->loadNext(is);
    ConfigFile f(doc);
    f.setPrefix("config.");
    std::vector<double> params;

#define CHECK_THROW(KEY)                                                \
    if(!f.contains("" #KEY)){                                           \
      throw ICLException("unsable to parse xml-file to "                \
                         "ImageUndistortion: entry " #KEY               \
                         " is missing!");                               \
    }
    CHECK_THROW(model);
    std::string model = f["model"];
    
    if(model == "MatlabModel5Params"){
#define LFS(KEY)                                                        \
      CHECK_THROW(KEY)                                                  \
      params.push_back(f["" #KEY]);                             
      
      LFS(intrin.fx);
      LFS(intrin.fy);
      LFS(intrin.ix);
      LFS(intrin.iy);
      LFS(intrin.skew);
      LFS(udist.k1);
      LFS(udist.k2);
      LFS(udist.k3);
      LFS(udist.k4);
      LFS(udist.k5);
    }else if(model == "SimpleARTBased"){
      LFS(center.x);
      LFS(center.y);
      LFS(udist.f);
      LFS(udist.scale);
    }else{
      throw ICLException("unable to parse xml-file to ImageUndistortion: invalide model name");
    }
    
    CHECK_THROW(size.width);
    CHECK_THROW(size.height);
    
    dest = ImageUndistortion(model, params, Size(f["size.width"], f["size.height"]));

#undef LSF
#undef CHECK_THROW
    return is;
  }
  
  std::ostream &operator<<(std::ostream &s,const ImageUndistortion &udist){
    ICLASSERT_THROW(!udist.isNull(), ICLException("unable to serialize a null-ImageUndistortion instance"));
    ConfigFile f;
    f.setPrefix("config.");
    
    f["size.width"] = udist.getImageSize().width;
    f["size.height"] = udist.getImageSize().height;
    f["model"] = udist.getModel();
    
    const std::vector<double> &p = udist.getParams();
    if(udist.getModel() == "MatlabModel5Params"){
      f["intrin.fx"] = p[0];
      f["intrin.fy"] = p[1];
      f["intrin.ix"] = p[2];
      f["intrin.iy"] = p[3];
      f["intrin.skew"] = p[4];
      f["udist.k1"] = p[5];
      f["udist.k2"] = p[6];
      f["udist.k3"] = p[7];
      f["udist.k4"] = p[8];
      f["udist.k5"] = p[9];
    }else{
      f["center.x"] = p[0];
      f["center.y"] = p[1];
      f["udist.f"] = p[2];
      f["udist.scale"] = p[3];
    }
    return s << f;
  }

  const Img32f &ImageUndistortion::createWarpMap() const{
    if(!impl->warpMap.getChannels()){
      impl->warpMap.setChannels(2);
      impl->warpMap.setSize(getImageSize());
      Channel32f cs[2] = { impl->warpMap[0], impl->warpMap[1] };
      const Size size = getImageSize();
      for(int xi=0;xi<size.width;++xi){
        for(int yi=0;yi<size.height; ++yi){
          Point32f p = impl->undistort(Point32f(xi,yi));
          cs[0](xi,yi) = p.x;
          cs[1](xi,yi) = p.y; 
        }
      }
    }
    return impl->warpMap;
  }
}