#include "hoCuNDArray_utils.h"
#include "radial_utilities.h"
#include "cuNDArray_fileio.h"
#include "cuNDArray.h"
#include "imageOperator.h"
#include "identityOperator.h"
#include "hoPartialDerivativeOperator.h"
#include "hoCuConebeamProjectionOperator.h"
#include "cuConvolutionOperator.h"
#include "hoCuNDArray_blas.h"
#include "hoCuNDArray_elemwise.h"
#include "hoCuNDArray_blas.h"
#include "cgSolver.h"
#include "CBCT_acquisition.h"
#include "complext.h"
#include "encodingOperatorContainer.h"
#include "vector_td_io.h"
#include "hoPartialDerivativeOperator.h"
#include "hoCuGPBBSolver.h"
#include "hoCuTvOperator.h"
#include "hoCuTvPicsOperator.h"
#include "hoCuNCGSolver.h"
#include "hoCuCgDescentSolver.h"
#include "hoNDArray_utils.h"
#include "hoCuPartialDerivativeOperator.h"
#include "cuTvOperator.h"
#include "cuTv1dOperator.h"
#include "cuTvPicsOperator.h"
#include "CBSubsetOperator.h"
#include "osSPSSolver.h"
#include "osMOMSolver.h"
#include <iostream>
#include <algorithm>
#include <sstream>
#include <math_constants.h>
#include <boost/program_options.hpp>
#include <boost/make_shared.hpp>
#include "cuSolverUtils.h"
using namespace std;
using namespace Gadgetron;

namespace po = boost::program_options;

boost::shared_ptr<hoCuNDArray<float> > calculate_prior(boost::shared_ptr<CBCT_binning>  binning,boost::shared_ptr<CBCT_acquisition> ps, hoCuNDArray<float>& projections, std::vector<size_t> is_dims, floatd3 imageDimensions){
  std::cout << "Calculating FDK prior" << std::endl;
	boost::shared_ptr<CBCT_binning> binning_pics= binning->get_3d_binning();
	    std::vector<size_t> is_dims3d = is_dims;
	    is_dims3d.pop_back();
	    boost::shared_ptr< hoCuConebeamProjectionOperator >
	      Ep( new hoCuConebeamProjectionOperator() );
	    Ep->setup(ps,binning_pics,imageDimensions);
	    Ep->set_codomain_dimensions(ps->get_projections()->get_dimensions().get());
	    Ep->set_domain_dimensions(&is_dims3d);
	      Ep->set_use_filtered_backprojection(true);
	    boost::shared_ptr<hoCuNDArray<float> > prior3d(new hoCuNDArray<float>(&is_dims3d));
	    Ep->mult_MH(&projections,prior3d.get());

	    hoCuNDArray<float> tmp_proj(*ps->get_projections());
	    Ep->mult_M(prior3d.get(),&tmp_proj);
	    float s = dot(ps->get_projections().get(),&tmp_proj)/dot(&tmp_proj,&tmp_proj);
	    *prior3d *= s;
	    boost::shared_ptr<hoCuNDArray<float> > prior(new hoCuNDArray<float>(*expand( prior3d.get(), is_dims.back() )));
	    std::cout << "Prior complete" << std::endl;
	    return prior;
}

int main(int argc, char** argv)
{
  string acquisition_filename;
  string outputFile;
  uintd3 imageSize;
  floatd3 voxelSize;
  int device;
  unsigned int downsamples;
  unsigned int iterations;
  unsigned int subsets;
  float rho;
  float tv_weight,pics_weight;

  po::options_description desc("Allowed options");

  desc.add_options()
    ("help", "produce help message")
    ("acquisition,a", po::value<string>(&acquisition_filename)->default_value("acquisition.hdf5"), "Acquisition data")
    ("samples,n",po::value<unsigned int>(),"Number of samples per ray")
    ("output,f", po::value<string>(&outputFile)->default_value("reconstruction.real"), "Output filename")
    ("size,s",po::value<uintd3>(&imageSize)->default_value(uintd3(512,512,1)),"Image size in pixels")
    ("binning,b",po::value<string>(),"Binning file for 4d reconstruction")
    ("SAG","Use exact SAG correction if present")
    ("voxelSize,v",po::value<floatd3>(&voxelSize)->default_value(floatd3(0.488f,0.488f,1.0f)),"Voxel size in mm")
    ("dimensions,d",po::value<floatd3>(),"Image dimensions in mm. Overwrites voxelSize.")
    ("iterations,i",po::value<unsigned int>(&iterations)->default_value(10),"Number of iterations")
    ("device",po::value<int>(&device)->default_value(0),"Number of the device to use (0 indexed)")
    ("downsample,D",po::value<unsigned int>(&downsamples)->default_value(0),"Downsample projections this factor")
    ("subsets,u",po::value<unsigned int>(&subsets)->default_value(10),"Number of subsets to use")
    ("TV",po::value<float>(&tv_weight)->default_value(0),"Total variation weight")
    ("PICS",po::value<float>(&pics_weight)->default_value(0),"PICS weight")
    ("use_prior","Use an FDK prior")
    ;

  po::variables_map vm;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);

  if (vm.count("help")) {
    cout << desc << "\n";
    return 1;
  }

  std::cout << "Command line options:" << std::endl;
  for (po::variables_map::iterator it = vm.begin(); it != vm.end(); ++it){
    boost::any a = it->second.value();
    std::cout << it->first << ": ";
    if (a.type() == typeid(std::string)) std::cout << it->second.as<std::string>();
    else if (a.type() == typeid(int)) std::cout << it->second.as<int>();
    else if (a.type() == typeid(unsigned int)) std::cout << it->second.as<unsigned int>();
    else if (a.type() == typeid(float)) std::cout << it->second.as<float>();
    else if (a.type() == typeid(vector_td<float,3>)) std::cout << it->second.as<vector_td<float,3> >();
    else if (a.type() == typeid(vector_td<int,3>)) std::cout << it->second.as<vector_td<int,3> >();
    else if (a.type() == typeid(vector_td<unsigned int,3>)) std::cout << it->second.as<vector_td<unsigned int,3> >();
    else std::cout << "Unknown type" << std::endl;
    std::cout << std::endl;
  }

  cudaSetDevice(device);
  cudaDeviceReset();

  //Really weird stuff. Needed to initialize the device?? Should find real bug.
  cudaDeviceManager::Instance()->lockHandle();
  cudaDeviceManager::Instance()->unlockHandle();

  boost::shared_ptr<CBCT_acquisition> ps(new CBCT_acquisition());
  ps->load(acquisition_filename);
  ps->get_geometry()->print(std::cout);
	ps->downsample(downsamples);

  float SDD = ps->get_geometry()->get_SDD();
  float SAD = ps->get_geometry()->get_SAD();

  boost::shared_ptr<CBCT_binning> binning(new CBCT_binning());
  if (vm.count("binning")){
    std::cout << "Loading binning data" << std::endl;
    binning->load(vm["binning"].as<string>());
  } else binning->set_as_default_3d_bin(ps->get_projections()->get_size(2));
  binning->print(std::cout);

  floatd3 imageDimensions;
  if (vm.count("dimensions")){
    imageDimensions = vm["dimensions"].as<floatd3>();
    voxelSize = imageDimensions/imageSize;
  }
  else imageDimensions = voxelSize*imageSize;

  float lengthOfRay_in_mm = norm(imageDimensions);
  unsigned int numSamplesPerPixel = 3;
  float minSpacing = min(voxelSize)/numSamplesPerPixel;

  unsigned int numSamplesPerRay;
  if (vm.count("samples")) numSamplesPerRay = vm["samples"].as<unsigned int>();
  else numSamplesPerRay = ceil( lengthOfRay_in_mm / minSpacing );

  float step_size_in_mm = lengthOfRay_in_mm / numSamplesPerRay;
  size_t numProjs = ps->get_projections()->get_size(2);
  size_t needed_bytes = 2 * prod(imageSize) * sizeof(float);
  std::vector<size_t> is_dims = to_std_vector((uint64d3)imageSize);

  std::cout << "IS dimensions " << is_dims[0] << " " << is_dims[1] << " " << is_dims[2] << std::endl;
  std::cout << "Image size " << imageDimensions << std::endl;

  is_dims.push_back(binning->get_number_of_bins());
  osMOMSolver<cuNDArray<float>> solver;

  if (pics_weight > 0){
  	std::cout << "Calculating PICS prior" << std::endl;
  	hoCuConebeamProjectionOperator op;
  	auto bin3D = binning->get_3d_binning();
  	op.setup(ps,bin3D,imageDimensions);
  	hoCuNDArray<float> proj(*ps->get_projections());
  	op.offset_correct(&proj);
  	op.set_use_filtered_backprojection(true);

  	std::vector<size_t> is_dims3D = to_std_vector((uint64d3)imageSize);

  	hoCuNDArray<float> image(is_dims3D);

  	op.mult_MH(&proj,&image,false);
  	auto prior = boost::make_shared<cuNDArray<float>>(image);
  	auto PICS = boost::make_shared<cuTvPicsOperator<float,3>>();
  	PICS->set_prior(prior);
  	PICS->set_weight(pics_weight);
  	solver.add_nonlinear_operator(PICS);

  	//write_nd_array(prior.get(),"fdk_prior.real");
  }

  // Define encoding matrix
  auto E = boost::make_shared<CBSubsetOperator<cuNDArray> >(subsets);

  //E->setup(ps,binning,imageDimensions);
  E->setup(ps,binning,imageDimensions);
  E->set_domain_dimensions(&is_dims);
  E->set_codomain_dimensions(ps->get_projections()->get_dimensions().get());

  /*auto weight_array = boost::make_shared<cuNDArray<float>>(is_dims);
  {
  	boost::shared_ptr<linearOperator<cuNDArray<float>>> E2(E);
  	cuNDArray<float> tmp_proj(*ps->get_projections()->get_dimensions());
  	fill(&tmp_proj,1.0f);
  	E2->mult_MH(&tmp_proj,weight_array.get(),false);
  	clamp_min(weight_array.get(),1.0f);
  	reciprocal_inplace(weight_array.get());
  	//*weight_array *= *weight_array;
  }
  write_nd_array(weight_array.get(),"weights.real");
*/
  //hoCuCgDescentSolver<float> solver;

  //osSPSSolver<hoNDArray<float>> solver;
  //hoCuNCGSolver<float> solver;
  solver.set_encoding_operator(E);
  //solver.set_domain_dimensions(&is_dims);
  solver.set_max_iterations(iterations);
  solver.set_output_mode(osSPSSolver<cuNDArray<float>>::OUTPUT_VERBOSE);
  //solver.set_non_negativity_constraint(true);
  //solver.set_rho(rho);

  if (tv_weight > 0){

  	auto total_variation = boost::make_shared<cuTvOperator<float,4>>();
  	total_variation->set_weight(tv_weight);
  	//total_variation->set_weight_array(weight_array);
  	solver.add_nonlinear_operator(total_variation);
  	solver.set_kappa(tv_weight);
  	/*
  	auto total_variation_t = boost::make_shared<cuTv1DOperator<float,4>>();
  	total_variation_t->set_weight(tv_weight);
  	solver.add_nonlinear_operator(total_variation_t);
  	*/
/*
  	auto total_variation2 = boost::make_shared<cuWTvOperator<float,4>>();
  	total_variation2->set_step(2);
  	total_variation2->set_weight(tv_weight);
  	total_variation2->set_weight_array(weight_array);
  	solver.add_nonlinear_operator(total_variation2);
  	auto total_variation3 = boost::make_shared<cuWTvOperator<float,4>>();
  	total_variation3->set_step(3);
  	total_variation3->set_weight(tv_weight);
  	total_variation3->set_weight_array(weight_array);
  	solver.add_nonlinear_operator(total_variation3);
*/
  }

  cuNDArray<float> projections(*ps->get_projections());
  std::cout << "Projection norm:" << nrm2(&projections) << std::endl;
  //E->set_use_offset_correction(false);
  E->offset_correct(&projections);
  std::cout << "Projection norm:" << nrm2(&projections) << std::endl;

  {
  	E->set_use_offset_correction(false);
  	linearOperator<cuNDArray<float>>* E_all = E.get();
  	auto precon_image = boost::make_shared<cuNDArray<float>>(is_dims);
  	fill(precon_image.get(),1.0f);
  	cuNDArray<float> tmp_proj(projections.get_dimensions());

  	E_all->mult_M(precon_image.get(),&tmp_proj);
  	E_all->mult_MH(&tmp_proj,precon_image.get());

   	clamp_min(precon_image.get(),1e-6f);
  	reciprocal_inplace(precon_image.get());
  	solver.set_preconditioning_image(precon_image);

  	std::cout << "Precon mean: " << mean(precon_image.get()) << std::endl;
  	E->set_use_offset_correction(true);

  }

/*
    boost::shared_ptr<hoCuNDArray<float> > prior;

  if (vm.count("use_prior")) {
  	prior = calculate_prior(binning,ps,projections,is_dims,imageDimensions);
  	solver.set_x0(prior);
  }
*/

  auto result = solver.solve(&projections);

  write_nd_array<float>( result.get(), outputFile.c_str());
}


// Copyright (C) 2016 The Regents of the University of California (Regents).
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//
//     * Neither the name of The Regents or University of California nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Please contact the author of this library if you have any questions.
// Author: Chris Sweeney (cmsweeney@cs.ucsb.edu)

#include "theia/sfm/undistort_image.h"

#include <Eigen/Core>
#include <algorithm>
#include <limits>

#include "theia/image/image.h"
#include "theia/sfm/camera/camera.h"
#include "theia/sfm/camera/camera_intrinsics_model.h"
#include "theia/sfm/camera/fisheye_camera_model.h"
#include "theia/sfm/camera/pinhole_camera_model.h"
#include "theia/sfm/camera/pinhole_radial_tangential_camera_model.h"
#include "theia/sfm/reconstruction.h"
#include "theia/sfm/track.h"
#include "theia/sfm/types.h"
#include "theia/sfm/view.h"

namespace theia {
namespace {

// Set all radial distortion parameters to zero. This requires adjusting
// different parameters for each camera intrinsics model type.
void SetLensDistortionToZero(Camera* camera) {
  double* intrinsics = camera->mutable_intrinsics();
  switch (camera->GetCameraIntrinsicsModelType()) {
    case CameraIntrinsicsModelType::PINHOLE:
      intrinsics[PinholeCameraModel::RADIAL_DISTORTION_1] = 0.0;
      intrinsics[PinholeCameraModel::RADIAL_DISTORTION_2] = 0.0;
      break;
    case CameraIntrinsicsModelType::PINHOLE_RADIAL_TANGENTIAL:
      intrinsics[PinholeRadialTangentialCameraModel::RADIAL_DISTORTION_1] = 0.0;
      intrinsics[PinholeRadialTangentialCameraModel::RADIAL_DISTORTION_2] = 0.0;
      intrinsics[PinholeRadialTangentialCameraModel::RADIAL_DISTORTION_3] = 0.0;
      intrinsics[PinholeRadialTangentialCameraModel::TANGENTIAL_DISTORTION_1] =
          0.0;
      intrinsics[PinholeRadialTangentialCameraModel::TANGENTIAL_DISTORTION_2] =
          0.0;
      break;
    case CameraIntrinsicsModelType::FISHEYE:
      intrinsics[FisheyeCameraModel::RADIAL_DISTORTION_1] = 0.0;
      intrinsics[FisheyeCameraModel::RADIAL_DISTORTION_2] = 0.0;
      intrinsics[FisheyeCameraModel::RADIAL_DISTORTION_3] = 0.0;
      intrinsics[FisheyeCameraModel::RADIAL_DISTORTION_4] = 0.0;
      break;
    case CameraIntrinsicsModelType::FOV:
      intrinsics[FisheyeCameraModel::RADIAL_DISTORTION_1] = 0.0;
      break;
    default:
      LOG(FATAL) << "Invalid camera intrinsics type.";
      break;
  }
}

// We seek a mapping from undistorted pixels to distorted pixels such that a
// valid mapping exists for all undistorted pixels. This is effectively
// undistorting the image then cropping it so that no unmapped (black) pixels
// are in the undistorted image. To do this, we need to find the min and max x/y
// values of unmapped pixels so that we can determine the size and location of
// the crop.
void FindUndistortedImageBoundary(const Camera& distorted_camera,
                                  const Camera& undistorted_camera,
                                  Eigen::Vector4d* bounds) {
  const CameraIntrinsicsModel& distorted_intrinsics =
      *distorted_camera.CameraIntrinsics();
  const CameraIntrinsicsModel& undistorted_intrinsics =
      *undistorted_camera.CameraIntrinsics();

  // Find the max and min locations of the undistorted pixels.
  double left_max_x = std::numeric_limits<double>::lowest();
  double right_min_x = std::numeric_limits<double>::max();
  for (size_t y = 0; y < distorted_camera.ImageHeight(); ++y) {
    // Left border.
    const Eigen::Vector3d distorted_point1 =
        distorted_intrinsics.ImageToCameraCoordinates(
            Eigen::Vector2d(0.5, y + 0.0));
    const Eigen::Vector2d undistorted_point1 =
        undistorted_intrinsics.CameraToImageCoordinates(distorted_point1);
    left_max_x = std::max(left_max_x, undistorted_point1(0));

    // Right border.
    const Eigen::Vector3d distorted_point2 =
        distorted_intrinsics.ImageToCameraCoordinates(
            Eigen::Vector2d(distorted_camera.ImageWidth() - 0.5, y + 0.5));
    const Eigen::Vector2d undistorted_point2 =
        undistorted_intrinsics.CameraToImageCoordinates(distorted_point2);
    right_min_x = std::min(right_min_x, undistorted_point2(0));
  }

  // Determine min, max coordinates along left / right image border.
  double top_max_y = std::numeric_limits<double>::lowest();
  double bottom_min_y = std::numeric_limits<double>::max();
  for (size_t x = 0; x < distorted_camera.ImageWidth(); ++x) {
    // Top border.
    const Eigen::Vector3d distorted_point1 =
        distorted_intrinsics.ImageToCameraCoordinates(
            Eigen::Vector2d(x + 0.5, 0.5));
    const Eigen::Vector2d undistorted_point1 =
        undistorted_intrinsics.CameraToImageCoordinates(distorted_point1);
    top_max_y = std::max(top_max_y, undistorted_point1(1));

    // Bottom border.
    const Eigen::Vector3d distorted_point2 =
        distorted_intrinsics.ImageToCameraCoordinates(
            Eigen::Vector2d(x + 0.5, distorted_camera.ImageHeight() - 0.5));
    const Eigen::Vector2d undistorted_point2 =
        undistorted_intrinsics.CameraToImageCoordinates(distorted_point2);
    bottom_min_y = std::min(bottom_min_y, undistorted_point2(1));
  }

  *bounds = Eigen::Vector4d(left_max_x, right_min_x, top_max_y, bottom_min_y);
}

// Create an undistorted image from the distorted image given the distorted and
// undistorted camera parameters. This function only maps the pixels to create
// the undistorted image and assumes the camera parameters and image sizes have
// already been solved for.
void RemoveImageLensDistortion(const Camera& distorted_camera,
                               const FloatImage& distorted_image,
                               const Camera& undistorted_camera,
                               FloatImage* undistorted_image) {
  const CameraIntrinsicsModel& distorted_intrinsics =
      *distorted_camera.CameraIntrinsics();
  const CameraIntrinsicsModel& undistorted_intrinsics =
      *undistorted_camera.CameraIntrinsics();

  // For each pixel in the undistorted image, find the coordinate in the
  // distorted image and set the pixel color accordingly.
  const int num_channels = distorted_image.Channels();
  OpenImageIO::ImageBuf& undistorted_img =
      undistorted_image->GetOpenImageIOImageBuf();
  OpenImageIO::ImageBuf::Iterator<float> undistorted_it(undistorted_img);
  for (; !undistorted_it.done(); ++undistorted_it) {
    Eigen::Vector2d image_point(undistorted_it.x() + 0.5,
                                undistorted_it.y() + 0.5);

    // Camera models assume that the upper left pixel center is (0.5, 0.5).
    const Eigen::Vector3d distorted_point =
        undistorted_intrinsics.ImageToCameraCoordinates(image_point);
    const Eigen::Vector2d distorted_pixel =
        distorted_intrinsics.CameraToImageCoordinates(distorted_point);
    const Eigen::Vector2d pixel(std::round<int>(distorted_pixel.x() - 0.5),
                                std::round<int>(distorted_pixel.y() - 0.5));
    // Set all color channels appropriately. Note that we do not need to check
    // if the distorted pixel is within the image bounds since the
    // FindUndistortedImageBoundary function has guaranteed that all pixels will
    // be within the image borders.
    for (int c = 0; c < num_channels; c++) {
      undistorted_it[c] = distorted_image.BilinearInterpolate(
          distorted_pixel.x(), distorted_pixel.y(), c);
    }
  }
}

}  // namespace

bool UndistortImage(const Camera& distorted_camera,
                    const FloatImage& distorted_image,
                    const Camera& undistorted_camera,
                    FloatImage* undistorted_image) {
  // Undistort the image.
  if (distorted_image.Channels() == 1) {
    undistorted_image->ConvertToGrayscaleImage();
  } else {
    undistorted_image->ConvertToRGBImage();
  }
  undistorted_image->Resize(undistorted_camera.ImageWidth(),
                            undistorted_camera.ImageHeight());

  // Remap the distorted pixels into the undistorted image.
  RemoveImageLensDistortion(distorted_camera,
                            distorted_image,
                            undistorted_camera,
                            undistorted_image);

  return true;
}

// Create the undistorted camera by removing radial distortion parameters.
bool UndistortCamera(const Camera& distorted_camera,
                     Camera* undistorted_camera) {
  *undistorted_camera = distorted_camera;
  SetLensDistortionToZero(undistorted_camera);

  Eigen::Vector4d undistorted_image_boundaries;
  FindUndistortedImageBoundary(distorted_camera,
                               *undistorted_camera,
                               &undistorted_image_boundaries);

  // Given the locations of the min/max undistorted pixels, compute the scale
  // factor to resize the undistorted image.
  const double cx = undistorted_camera->PrincipalPointX();
  const double cy = undistorted_camera->PrincipalPointY();
  const double left_max_x = undistorted_image_boundaries(0);
  const double right_min_x = undistorted_image_boundaries(1);
  const double top_max_y = undistorted_image_boundaries(2);
  const double bottom_min_y = undistorted_image_boundaries(3);

  // Scale undistorted camera dimensions.
  const double scale_x =
      1.0 /
      std::max(cx / (cx - left_max_x),
               (distorted_camera.ImageWidth() - 0.5 - cx) / (right_min_x - cx));
  const double scale_y =
      1.0 / std::max(cy / (cy - top_max_y),
                     (distorted_camera.ImageHeight() - 0.5 - cy) /
                         (bottom_min_y - cy));

  undistorted_camera->SetImageSize(
      static_cast<int>(
          std::max(1.0, scale_x * undistorted_camera->ImageWidth())),
      static_cast<int>(
          std::max(1.0, scale_y * undistorted_camera->ImageHeight())));

  // Scale the principal point according to the new dimensions of the image.
  undistorted_camera->SetPrincipalPoint(
      undistorted_camera->PrincipalPointX() *
          static_cast<double>(undistorted_camera->ImageWidth()) /
          distorted_camera.ImageWidth(),
      undistorted_camera->PrincipalPointY() *
          static_cast<double>(undistorted_camera->ImageHeight()) /
          distorted_camera.ImageHeight());

  return true;
}

bool UndistortReconstruction(Reconstruction* reconstruction) {
  const auto view_ids = reconstruction->ViewIds();
  for (const ViewId view_id : view_ids) {
    View* view = reconstruction->MutableView(view_id);
    if (view == nullptr || !view->IsEstimated()) {
      continue;
    }

    // Undistort the image to obtain the undistorted camera.
    const Camera distorted_camera = view->Camera();
    Camera* undistorted_camera = view->MutableCamera();
    if (!UndistortCamera(distorted_camera, undistorted_camera)) {
      return false;
    }

    // The camera intrinsics models describe how to distort and undistort the
    // pixels.
    const CameraIntrinsicsModel& distorted_intrinsics =
        *distorted_camera.CameraIntrinsics();
    const CameraIntrinsicsModel& undistorted_intrinsics =
        *undistorted_camera->CameraIntrinsics();

    // Undisort all features seen by the view.
    const auto track_ids = view->TrackIds();
    for (const TrackId track_id : track_ids) {
      Track* track = reconstruction->MutableTrack(track_id);
      if (track == nullptr || !track->IsEstimated()) {
        continue;
      }

      // Get the undistorted feature.
      const Feature* distorted_feature = view->GetFeature(track_id);
      const Eigen::Vector3d distorted_point =
          distorted_intrinsics.ImageToCameraCoordinates(*distorted_feature);
      const Eigen::Vector2d undistorted_feature =
          undistorted_intrinsics.CameraToImageCoordinates(distorted_point);

      // Add the new undistorted feature to the undistorted view. First, remove
      // the feature since it is not mutable.
      CHECK(view->RemoveFeature(track_id));

      // Only add the feature back if it is within the new image
      // boundaries. Otherwise, remove the feature from the view and the track.
      if (undistorted_feature.x() < 0 ||
          undistorted_feature.x() > undistorted_camera->ImageWidth() - 1 ||
          undistorted_feature.y() < 0 ||
          undistorted_feature.y() > undistorted_camera->ImageHeight() - 1) {
        track->RemoveView(view_id);
        // Remove the track if there are no more observations.
        if (track->NumViews() == 0) {
          reconstruction->RemoveTrack(track_id);
        }
      } else {
        view->AddFeature(track_id, undistorted_feature);
      }
    }
  }

  return true;
}

}  // namespace theia