/*
  High Performance Astrophysical Reconstruction and Processing (HARP)
  (c) 2014-2015, The Regents of the University of California,
  through Lawrence Berkeley National Laboratory.  See top
  level LICENSE file for details.
*/

#include <iostream>
#include <cstdio>

extern "C" {
  #include <unistd.h>
}

#include <boost/program_options.hpp>

#include <harp_mpi.hpp>

using namespace std;
using namespace harp;

namespace popts = boost::program_options;

int main ( int argc, char *argv[] ) {

  El::Initialize ( argc, argv );

  boost::mpi::environment env;

  double tstart;
  double tstop;

  cout.precision ( 10 );
  cerr.precision ( 10 );

  size_t lambda_width = 50;
  double lambda_res = 0.6;
  double lambda_b_min = 3579.0;
  double lambda_b_max = 5939.0;
  double lambda_r_min = 5635.0;
  double lambda_r_max = 7731.0;
  double lambda_z_min = 7445.0;
  double lambda_z_max = 9824.0;
  double lambda_min;
  double lambda_max;
  double lambda_hard_min;
  double lambda_hard_max;

  size_t spec_width = 0;
  size_t spec_overlap = 0;
  size_t spec_min = 0;
  size_t spec_max = 499;

  size_t res_width = 10;

  bool lambda_mask = false;

  string outdir = ".";
  string fibermap = "";
  string in_image = "";
  string in_psf = "";
  string in_sim = "";

  string band = "";
  string camera = "";
  string expid = "";

  string prefix = "harp:  ";

  bool quiet = false;
  bool debug = false;

  double global_start = MPI_Wtime();

  int gangsize = 0;

  // global communicator and grid

  boost::mpi::communicator comm;

  int np = comm.size();
  int myp = comm.rank();

  El::Grid grid ( El::mpi::COMM_WORLD );

  // Declare options

  popts::options_description desc ( "Allowed Options" );

  desc.add_options()
  ( "help,h", "display usage information" )
  ( "quiet,q", "supress information printing" )
  ( "debug,d", "write out intermediate data products for debugging" )
  ( "gangsize", popts::value < int > ( &gangsize ), "number of processes per gang (choose whole number of nodes and perfect square number of processes if possible)" )
  ( "image", popts::value < string > ( &in_image ), "input image file" )
  ( "psf", popts::value < string > ( &in_psf ), "input PSF file (XML)" )
  ( "psf_interp", "enable PSF interpolation" )
  ( "fibermap", popts::value < string > ( &fibermap ), "input fibermap" )
  ( "sim", popts::value < string > ( &in_sim ), "(optional) original simulated spectra" )
  ( "out", popts::value < string > ( &outdir ), "output directory (default = \".\")" )
  ( "spec_width", popts::value < size_t > ( &spec_width ), "number of spectra to solve at once" )
  ( "spec_overlap", popts::value < size_t > ( &spec_overlap ), "minimum spectra to overlap" )
  ( "spec_min", popts::value < size_t > ( &spec_min ), "first spectrum to solve (0)" )
  ( "spec_max", popts::value < size_t > ( &spec_max ), "last spectrum to solve (499)" )
  ( "lambda_width", popts::value < size_t > ( &lambda_width ), "number of wavelength points to solve at once" )
  ( "lambda_res", popts::value < double > ( &lambda_res ), "extraction resolution, in Angstroms" )
  ( "lambda_min", popts::value < double > ( &lambda_min ), "minimum extracted wavelength in Angstroms" )
  ( "lambda_max", popts::value < double > ( &lambda_max ), "maximum extracted wavelength in Angstroms" )
  ( "res_width", popts::value < size_t > ( &res_width ), "half-band width of diagonal resolution to keep (10)" )
  ;

  popts::variables_map vm;

  popts::store(popts::command_line_parser( argc, argv ).options(desc).run(), vm);

  popts::notify(vm);

  if ( ( ! vm.count("image") ) || ( ! vm.count("psf") ) || vm.count( "help" ) ) {
    if ( myp == 0 ) {
      cerr << endl;
      cerr << desc << endl;
    }
    El::Finalize();
    return 0;
  }

  if ( vm.count( "quiet" ) ) {
    quiet = true;
  }

  if ( vm.count( "debug" ) ) {
    debug = true;
  }

  size_t res_band = 2 * res_width + 1;

  // attempt to get metadata from the file name

  size_t baseoff = 0;
  size_t basetmp = in_image.find( '/', baseoff );
  while ( basetmp != string::npos ) {
    baseoff = basetmp + 1;
    basetmp = in_image.find( '/', baseoff );
  }

  size_t off = in_image.find ( '-', baseoff );
  if ( off != string::npos ) {
    off += 1;
    size_t pos = in_image.find ( '-', off );
    if ( pos != string::npos ) {
      band = in_image.substr ( off, 1 );
      camera = in_image.substr ( off, 2 );
      off = pos + 1;
      size_t ext = in_image.find ( '.', off );
      expid = in_image.substr ( off, ext - off );
    }
  }

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Image from camera \"" << camera << "\", exposure \"" << expid << "\"" << endl;
  }

  if ( ! vm.count( "lambda_min" ) ) {
    if ( band == "b" ) {
      lambda_min = lambda_b_min;
    } else if ( band == "r" ) {
      lambda_min = lambda_r_min;
    } else if ( band == "z" ) {
      lambda_min = lambda_z_min;
    } else {
      HARP_MPI_ABORT( myp, "image name does not conform to standard, and lambda min not specified" );
    }
  }

  if ( ! vm.count( "lambda_max" ) ) {
    if ( band == "b" ) {
      lambda_max = lambda_b_max;
    } else if ( band == "r" ) {
      lambda_max = lambda_r_max;
    } else if ( band == "z" ) {
      lambda_max = lambda_z_max;
    } else {
      HARP_MPI_ABORT( myp, "image name does not conform to standard, and lambda max not specified" );
    }
  }

  // we specify an extra overlap range in lambda, to remove
  // edge effects.

  size_t lambda_pre = (res_band * lambda_res);
  size_t lambda_post = (res_band * lambda_res);
  //size_t lambda_post = ((res_band + 1) * lambda_res);
  lambda_hard_min = lambda_min - lambda_pre;
  lambda_hard_max = lambda_max + lambda_post;

  // intra-gang communicator and grid

  if ( gangsize == 0 ) {
    gangsize = np;
  }

  int ngang = (int)( np / gangsize );
  int gangtot = ngang * gangsize;
  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Using " << ngang << " gangs of " << gangsize << " processes each" << endl;
  }
  if ( gangtot < np ) {
    if ( myp == 0 ) {
      cout << prefix << "  WARNING: " << (np-gangtot) << " processes are idle" << endl;
    }
  }
  int gang = (int)( myp / gangsize );
  int grank = myp % gangsize;
  if ( gang >= ngang ) {
    gang = MPI_UNDEFINED;
    grank = MPI_UNDEFINED;
  }

  boost::mpi::communicator gcomm = comm.split ( gang, grank );

  El::Grid gang_grid ( (MPI_Comm)gcomm );

  // inter-gang communicator

  boost::mpi::communicator rcomm = comm.split ( grank, gang );

  // Create input globally-distributed PSF from the serial one

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Creating PSF..." << endl;
  }

  tstart = MPI_Wtime();

  boost::property_tree::ptree pathprops;
  boost::property_tree::ptree props;

  pathprops.clear();
  pathprops.put ( "path", in_psf );
  pathprops.put ( "wavemin", lambda_hard_min );
  pathprops.put ( "wavemax", lambda_hard_max );
  pathprops.put ( "wavebin", lambda_res );
  pathprops.put ( "fibermin", 0 );
  pathprops.put ( "fibermax", 499 );
  pathprops.put ( "type", "xml" );
  if ( vm.count( "psf_interp" ) ) {
    pathprops.put ( "interpolation", 1 );
  } else {
    pathprops.put ( "interpolation", 0 );
  }

  props.clear();
  props.put ( "type", "specex" );
  props.put_child ( "props", pathprops );

  mpi_psf_p design = mpi_load_psf ( comm, props );

  size_t psf_imgrows = design->img_rows();
  size_t psf_imgcols = design->img_cols();
  size_t psf_npix = psf_imgrows * psf_imgcols;

  size_t psf_nspec = design->n_spec();
  size_t psf_nlambda = design->n_lambda();
  size_t psf_nbins = psf_nspec * psf_nlambda;

  vector_double psf_lambda = design->lambda();

  size_t nlambda = psf_nlambda - 2*res_band;

  vector_double lambda(nlambda);
  for ( size_t i = 0; i < nlambda; ++i ) {
    lambda[i] = psf_lambda[i + res_band];
  }

  tstop = MPI_Wtime();

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "  time = " << tstop-tstart << " seconds" << endl;
    cout << prefix << "  " << psf_nspec << " spectra each with " << psf_nlambda << " wavelength points" << endl;
    cout << prefix << "  image dimensions = " << psf_imgcols << " x " << psf_imgrows << endl;
  }

  // Read input image

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Creating input image..." << endl;
  }

  tstart = MPI_Wtime();

  pathprops.clear();
  pathprops.put ( "path", in_image );

  props.clear();
  props.put ( "type", "desi" );
  props.put_child ( "props", pathprops );

  mpi_image_p img = mpi_load_image ( comm, props );

  size_t imgrows = img->n_rows();
  size_t imgcols = img->n_cols();
  size_t npix = imgrows * imgcols;

  tstop = MPI_Wtime();

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "  time = " << tstop-tstart << " seconds" << endl;
    cout << prefix << "  dimensions = " << imgcols << " x " << imgrows << endl;
  }

  if ( ( imgrows != psf_imgrows ) || ( imgcols != psf_imgcols ) ) {
    HARP_MPI_ABORT( myp, "PSF and image have inconsistent dimensions" );
  }

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Reading input image..." << endl;
  }

  tstart = MPI_Wtime();

  elem_matrix_local measured ( npix );
  local_matrix_zero ( measured );

  elem_matrix_local invnoise ( npix );
  local_matrix_zero ( invnoise );

  vector_mask msk;
  msk.clear();

  img->values ( measured );
  img->inv_variance ( invnoise );
  img->mask ( msk );

  tstop = MPI_Wtime();

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "  time = " << tstop-tstart << " seconds" << endl;
  }

  // target list

  mpi_targets_p fibers;
  vector < object_p > target_list;

  if ( fibermap != "" ) {

    pathprops.clear();
    pathprops.put ( "path", fibermap );

    props.clear();
    props.put ( "type", "desi" );
    props.put_child ( "props", pathprops );

    fibers = mpi_load_targets ( comm, props );

    target_list = fibers->objects();

  }

  // select the global region of spec and wavelength space that we are
  // solving for.

  size_t global_first_spec = spec_min;
  size_t global_nspec = spec_max - spec_min + 1;

  size_t global_nbins = global_nspec * psf_nlambda;

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Solving globally for " << global_nspec << " spectra and " << psf_nlambda << " wavelength points" << endl;
  }

  // output spectral products

  mpi_matrix data_truth ( 0, 1, grid );
  mpi_matrix data_Rtruth ( 0, 1, grid );

  mpi_matrix data_f ( global_nbins, 1, grid );
  mpi_matrix_zero ( data_f );

  mpi_matrix data_Rf ( global_nbins, 1, grid );
  mpi_matrix_zero ( data_Rf );

  mpi_matrix data_err ( global_nbins, 1, grid );
  mpi_matrix_zero ( data_err );

  mpi_matrix data_Rdiag ( global_nbins, res_band, grid );
  mpi_matrix_zero ( data_Rdiag );

  // read truth spectra, if provided

  bool dotruth = false;

  if ( vm.count ( "sim" ) > 0 ) {

    if ( ( myp == 0 ) && ( ! quiet ) ) {
      cout << prefix << "Reading input truth spectra..." << endl;
    }

    tstart = MPI_Wtime();

    pathprops.clear();
    pathprops.put ( "path", in_sim );

    props.clear();
    props.put ( "type", "desisim" );
    props.put_child ( "props", pathprops );

    mpi_spec_p truth_spec = mpi_load_spec ( comm, props );
    size_t truth_nspec = truth_spec->n_spec();
    size_t truth_nlambda = truth_spec->n_lambda();
    size_t truth_nbins = truth_nspec * truth_nlambda;

    if ( truth_nbins != global_nbins ) {
      ostringstream o;
      o << "truth spectrum has " << truth_nbins << " spectral bins, but the PSF has " << global_nbins << " bins";
      HARP_MPI_ABORT( myp, o.str().c_str() );
    }

    data_truth.Resize ( global_nbins, 1 );
    mpi_matrix_zero ( data_truth );

    data_Rtruth.Resize ( global_nbins, 1 );
    mpi_matrix_zero ( data_Rtruth );

    vector_double truth_lambda;

    truth_spec->values ( data_truth );
    truth_spec->lambda ( truth_lambda );

    for ( size_t i = 0; i < nlambda; ++i ) {
      if ( fabs ( truth_lambda[i] - lambda[i] ) / lambda[i] > 1.0e-5 ) {
        ostringstream o;
        o << "wavelength point " << i << " does not match between PSF (" << lambda[i] << ") and truth spec (" << truth_lambda[i] << ")";
        HARP_MPI_ABORT( myp, o.str().c_str() );
      }
    }

    dotruth = true;

    tstop = MPI_Wtime();

    if ( ( myp == 0 ) && ( ! quiet ) ) {
      cout << prefix << "  time = " << tstop-tstart << " seconds" << endl;
    }

  }

  // Divide spectral data into overlapping regions

  if ( spec_width < 1 ) {
    // do all spectra
    spec_width = global_nspec;
  }

  if ( lambda_width < 1 ) {
    lambda_width = psf_nlambda;
  }

  mpi_spec_slice_p slice ( new mpi_spec_slice ( rcomm, gcomm, global_first_spec, res_band, global_nspec, nlambda, spec_width, lambda_width, spec_overlap, res_band ) );

  size_t gang_nregion = slice->regions().size();
  size_t nregion = 0;

  boost::mpi::reduce ( rcomm, gang_nregion, nregion, std::plus < size_t > (), 0 );

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Extracting " << nregion << " spectral chunks, each with " << spec_width << " spectra ( overlap = " << spec_overlap << " ) and " << lambda_width << " lambda points ( overlap = " << res_band << " )" << endl;
  }

  if ( debug && ( ! quiet ) ) {

    // each gang prints info about assigned slices

    if ( gcomm.rank() == 0 ) {

      vector < spec_slice_region > gslices = slice->regions();

      for ( size_t r = 0; r < rcomm.size(); ++r ) {

        if ( rcomm.rank() == r ) {

          cout << prefix << "  Gang " << r << " processing " << gslices.size() << " chunks:" << endl;
          for ( vector < spec_slice_region > :: const_iterator gsit = gslices.begin(); gsit != gslices.end(); ++gsit ) {
            cout << prefix << "    spec (" << gsit->first_good_spec << " - " << (gsit->first_good_spec + gsit->n_good_spec - 1) << ") wavelength (" << gsit->first_good_lambda << " - " << (gsit->first_good_lambda + gsit->n_good_lambda - 1) << ")" << endl;
          }

          cout.flush();

        }

        rcomm.barrier();

        cout.flush();

        rcomm.barrier();

      }

    }

  }

  // timing results

  map < string, double > timing;

  // Process all spectral slices

  string extract_prefix = "";
  if ( ! quiet ) {
    extract_prefix = prefix;
  }

  mpi_extract_slices ( slice, design, measured, invnoise, data_truth, res_band, data_Rdiag, data_Rf, data_f, data_err, data_Rtruth, timing, lambda_mask, prefix );

  /*
  string outres = outdir + "/debug_res.out";
  ofstream fres;
  fres.open(outres.c_str(), ios::out);
  El::Print ( data_Rdiag, "res_diag", fres );
  fres.close();
  */

  // subtract sky if needed

  // Total timings- get the mean, min and max from all processes

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Aggregate Timings:" << endl;
  }

  map < string, string > counter_desc;
  counter_desc [ "design" ] = "  Build design matrix =";
  counter_desc [ "inverse" ] = "  Build inverse covariance =";
  counter_desc [ "eigen" ] = "  Eigendecompose inverse =";
  counter_desc [ "norm" ] = "  Compute column norm =";
  counter_desc [ "nsespec" ] = "  Compute noise weighted spec =";
  counter_desc [ "extract" ] = "  Extract spectra =";

  double time_mean;
  double time_min;
  double time_max;

  for ( map < string, double > :: const_iterator itcounter = timing.begin(); itcounter != timing.end(); ++itcounter ) {

    time_mean = 0.0;
    time_min = 0.0;
    time_max = 0.0;

    boost::mpi::reduce ( comm, itcounter->second, time_mean, std::plus < double > (), 0 );
    time_mean /= (double)comm.size();

    boost::mpi::reduce ( comm, itcounter->second, time_min, boost::mpi::minimum < double > (), 0 );

    boost::mpi::reduce ( comm, itcounter->second, time_max, boost::mpi::maximum < double > (), 0 );

    if ( ( myp == 0 ) && ( ! quiet ) ) {
      cout << prefix << counter_desc [ itcounter->first ] << " " << time_mean << " ( " << time_min << " low / " << time_max << " high ) seconds" << endl;
    }

  }

  // Write outputs and compute chi-square if input truth is given

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Writing solution and error..." << endl;
  }

  tstart = MPI_Wtime();

  // get local copies.

  elem_matrix_local loc_truth;
  elem_matrix_local loc_Rtruth;
  elem_matrix_local loc_Rf;
  elem_matrix_local loc_f;
  elem_matrix_local loc_err;

  El::AxpyInterface < double > globloc;

  if ( dotruth ) {
    globloc.Attach( El::GLOBAL_TO_LOCAL, data_Rtruth );
    if ( myp == 0 ) {
      loc_Rtruth.Resize ( data_Rtruth.Height(), 1 );
      local_matrix_zero ( loc_Rtruth );
      globloc.Axpy ( 1.0, loc_Rtruth, 0, 0 );
    }
    globloc.Detach();

    globloc.Attach( El::GLOBAL_TO_LOCAL, data_truth );
    if ( myp == 0 ) {
      loc_truth.Resize ( data_truth.Height(), 1 );
      local_matrix_zero ( loc_truth );
      globloc.Axpy ( 1.0, loc_truth, 0, 0 );
    }
    globloc.Detach();
  }

  globloc.Attach( El::GLOBAL_TO_LOCAL, data_Rf );
  if ( myp == 0 ) {
    loc_Rf.Resize ( data_Rf.Height(), 1 );
    local_matrix_zero ( loc_Rf );
    globloc.Axpy ( 1.0, loc_Rf, 0, 0 );
  }
  globloc.Detach();

  globloc.Attach( El::GLOBAL_TO_LOCAL, data_f );
  if ( myp == 0 ) {
    loc_f.Resize ( data_f.Height(), 1 );
    local_matrix_zero ( loc_f );
    globloc.Axpy ( 1.0, loc_f, 0, 0 );
  }
  globloc.Detach();

  globloc.Attach( El::GLOBAL_TO_LOCAL, data_err );
  if ( myp == 0 ) {
    loc_err.Resize ( data_err.Height(), 1 );
    local_matrix_zero ( loc_err );
    globloc.Axpy ( 1.0, loc_err, 0, 0 );
  }
  globloc.Detach();

  boost::property_tree::ptree solution_props;
  string outfile;

  // define the clipping in wavelength to remove the extra overlap
  // region that we added at the beginning.

  size_t clip_nlambda = nlambda;
  //size_t clip_nlambda = psf_nlambda;

  vector_double & clip_lambda = lambda;
  //vector_double & clip_lambda = psf_lambda;

  size_t clip_lambda_start = res_band;
  //size_t clip_lambda_start = 0;

  size_t clip_lambda_stop = clip_lambda_start + clip_nlambda - 1;
  size_t clip_nbins = global_nspec * clip_nlambda;

  if ( myp == 0 ) {

    boost::property_tree::ptree child;

    solution_props.clear();

    // initialize to image metadata
    solution_props = boost::dynamic_pointer_cast < image_desi, image > ( img->local() )->meta();

    child.clear();
    child.put ( "VAL", clip_lambda[0] );
    child.put ( "TYPE", "F" );
    child.put ( "COM", "Starting wavelength [Angstroms]" );
    solution_props.put_child ( "CRVAL1", child );

    child.clear();
    child.put ( "VAL", lambda_res );
    child.put ( "TYPE", "F" );
    child.put ( "COM", "Wavelength step [Angstroms]" );
    solution_props.put_child ( "CDELT1", child );

    child.clear();
    child.put ( "VAL", "vac" );
    child.put ( "TYPE", "C" );
    child.put ( "COM", "Vacuum wavelengths" );
    solution_props.put_child ( "AIRORVAC", child );

    child.clear();
    child.put ( "VAL", 0 );
    child.put ( "TYPE", "I" );
    child.put ( "COM", "linear wavelength steps, not log10" );
    solution_props.put_child ( "LOGLAM", child );

    if ( dotruth ) {
      child.clear();
      child.put ( "VAL", in_sim );
      child.put ( "TYPE", "C" );
      child.put ( "COM", "Input simulation file" );
      solution_props.put_child ( "SIMFILE", child );
    }

    child.clear();
    child.put ( "VAL", in_psf );
    child.put ( "TYPE", "C" );
    child.put ( "COM", "Input spectral PSF" );
    solution_props.put_child ( "IN_PSF", child );

    child.clear();
    child.put ( "VAL", in_image );
    child.put ( "TYPE", "C" );
    child.put ( "COM", "Input image" );
    solution_props.put_child ( "IN_IMG", child );

    // write out solution

    vector_double errbuf;
    elem_to_ublas ( loc_err, errbuf );
    // convert from rms to inverse variance
    for ( size_t i = 0; i < errbuf.size(); ++i ) {
      errbuf[i] = 1.0 / ( errbuf[i] * errbuf[i] );
    }

    vector_double ubuf;
    elem_to_ublas ( loc_Rf, ubuf );

    //cerr << "clipping " << global_nspec << " spectra to lambda [0," << (psf_nlambda-1) << "] --> [" << clip_lambda_start << "," << (clip_lambda_stop) << "] (" << clip_nbins << " bins)" << endl;

    vector_double clip_errbuf( clip_nbins );
    vector_double clip_ubuf( clip_nbins );

    for ( size_t i = 0; i < global_nspec; ++i ) {
      for ( size_t j = 0; j < clip_nlambda; ++j ) {
        clip_errbuf[i * clip_nlambda + j] = errbuf[i * psf_nlambda + clip_lambda_start + j];
        clip_ubuf[i * clip_nlambda + j] = ubuf[i * psf_nlambda + clip_lambda_start + j];
      }
    }

    outfile = outdir + "/frame-" + camera + "-" + expid + ".fits";
    spec_desi::write ( outfile, solution_props, clip_ubuf, clip_errbuf, clip_lambda );

    if ( dotruth ) {

      elem_to_ublas ( loc_Rtruth, ubuf );

      for ( size_t i = 0; i < global_nspec; ++i ) {
        for ( size_t j = 0; j < clip_nlambda; ++j ) {
          clip_ubuf[i * clip_nlambda + j] = ubuf[i * psf_nlambda + clip_lambda_start + j];
        }
      }

      outfile = outdir + "/harp_rsim-" + camera + "-" + expid + ".fits";
      spec_desi::write ( outfile, solution_props, clip_ubuf, clip_errbuf, clip_lambda );

      vector_double data_chisq ( clip_nbins );

      double chisq_reduced = 0.0;

      for ( size_t i = 0; i < global_nspec; ++i ) {
        for ( size_t j = 0; j < clip_nlambda; ++j ) {
          size_t w = i * psf_nlambda + clip_lambda_start + j;
          size_t v = i * clip_nlambda + j;
          if ( loc_err.Get(w,0) > std::numeric_limits < double > :: epsilon() ) {
            data_chisq[v] = ( loc_Rf.Get(w,0) - loc_Rtruth.Get(w,0) ) / loc_err.Get(w,0);
            data_chisq[v] *= data_chisq[v];
          } else {
            data_chisq[v] = 0.0;
          }
          chisq_reduced += data_chisq[v];
        }
      }

      chisq_reduced /= (double)( clip_nbins - 1 );

      if ( ! quiet ) {
        cout << prefix << "  Reduced Chi square = " << chisq_reduced << endl;
      }

      outfile = outdir + "/harp_chisq-" + camera + "-" + expid + ".fits";
      spec_desi::write ( outfile, solution_props, data_chisq, clip_errbuf, clip_lambda );

    }

  }

  // Free unused buffers before writing resolution matrix.

  loc_truth.Empty();
  loc_Rtruth.Empty();
  loc_Rf.Empty();
  loc_err.Empty();

  measured.Empty();
  data_Rf.Empty();
  data_Rtruth.Empty();
  data_err.Empty();

  // The resolution matrix is too large to reduce to a single process
  // all at once.  We have no choice but to reduce this in buffers and append
  // to the FITS file...

  fitsfile * fp;
  int ret;
  int status = 0;

  long naxes[3];
  naxes[0] = clip_nlambda;
  naxes[1] = res_band;
  naxes[2] = global_nspec;

  int fitstype = fits::ftype < double > :: datatype();
  long fpixel[2];

  elem_matrix_local loc_Rdiag;

  size_t write_spec_chunk = 10;
  size_t write_spec_offset = 0;

  if (write_spec_chunk > global_nspec) {
    write_spec_chunk = global_nspec;
  }

  // write buffer
  vector_double resbuffer;

  if ( myp == 0 ) {
    fits::open_readwrite ( fp, outfile );
    fits::img_seek ( fp, 3 );

    ret = fits_create_img ( fp, fits::ftype< double >::bitpix(), 3, naxes, &status );
    fits::check ( status );

    fits::key_write ( fp, "EXTNAME", string("RESOLUTION"), "" );
  }

  while ( write_spec_offset < global_nspec ) {
    if ( write_spec_offset + write_spec_chunk > global_nspec ) {
      write_spec_chunk = global_nspec - write_spec_offset;
    }

    size_t write_spec_nbuf = write_spec_chunk * res_band * clip_nlambda;

    if ( myp == 0 ) {
      resbuffer.resize(write_spec_nbuf);
      resbuffer.clear();
    }

    // we have to copy the data into the buffer one spec
    // at a time, due to memory layout.
    for ( size_t k = 0; k < write_spec_chunk; ++k ) {

      globloc.Attach( El::GLOBAL_TO_LOCAL, data_Rdiag );
      //cerr << "data_Rdiag has dimensions " << data_Rdiag.Height() << " x " << data_Rdiag.Width() << endl;
      if ( myp == 0 ) {
        loc_Rdiag.Resize ( clip_nlambda, res_band );
        local_matrix_zero ( loc_Rdiag );
        globloc.Axpy ( 1.0, loc_Rdiag, (write_spec_offset+k)*psf_nlambda + clip_lambda_start, 0 );
      }
      globloc.Detach();

      if ( myp == 0 ) {
        size_t boff = k * res_band * clip_nlambda;
        //cerr << "copying resolution spec " << (k + write_spec_offset) << "(" << boff << "-" << (boff + res_band * clip_nlambda) << ") to memory buffer" << endl;

        for ( size_t j = 0; j < res_band; ++j ) {

          for ( size_t i = 0; i < clip_nlambda; ++i ) {

            size_t outdiag = (res_band - 1) - j;
            int64_t outlambda = -1;
            int64_t inlambda = i;

            int64_t loff = (int64_t)j - (int64_t)res_width;

            outlambda = inlambda + loff;
            if ( outlambda >= (int64_t)clip_nlambda ) {
              outlambda = -1;
            }

            if ( outlambda >= 0 ) {
              double val = loc_Rdiag.Get(inlambda, j);
              if ( fabs(val) < 1.0e-100 ) {
                val = 0.0;
              }
              resbuffer[boff + outdiag * clip_nlambda + outlambda] = val;
            }
          }
        }
      }
    }

    if ( myp == 0 ) {
      fpixel[0] = 1;
      fpixel[1] = 1;
      fpixel[2] = write_spec_offset + 1;
      long npix = (long)( write_spec_nbuf );

      ret = fits_write_pix ( fp, fitstype, fpixel, npix, &(resbuffer[0]), &status );
      fits::check ( status );
    }

    write_spec_offset += write_spec_chunk;
  }

  if ( myp == 0 ) {
    fits::close( fp );
  }

  tstop = MPI_Wtime();

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "  time = " << tstop-tstart << " seconds" << endl;
  }

  // free more memory

  loc_Rdiag.Empty();
  data_Rdiag.Empty();

  if ( debug ) {

    boost::property_tree::ptree child;
    child.clear();
    child.put ( "VAL", lambda[0] );
    child.put ( "TYPE", "F" );
    child.put ( "COM", "Starting wavelength [Angstroms]" );
    solution_props.erase("CRVAL1");
    solution_props.put_child ( "CRVAL1", child );

    if ( ( myp == 0 ) && ( ! quiet ) ) {
      cout << prefix << "Writing projected, deconvolved spectra..." << endl;
    }

    tstart = MPI_Wtime();

    map < size_t, set < size_t > > speclambda;

    //size_t slbins = 0;
    for ( size_t s = 0; s < global_nspec; ++s ) {
      for ( size_t l = 0; l < clip_nlambda; ++l ) {
        //cerr << "using spec " << (s + spec_min) << ", lambda " << (l + clip_lambda_start) << endl;
        speclambda[ s + spec_min ].insert ( l + clip_lambda_start );
        //++slbins;
      }
    }

    mpi_matrix_sparse AT ( comm, clip_nbins, psf_npix );

    design->project_transpose ( speclambda, AT );

    elem_matrix_local f_projected;

    elem_matrix_local newloc_f;
    newloc_f.Resize ( clip_nbins, 1 );
    local_matrix_zero ( newloc_f );

    for ( size_t i = 0; i < global_nspec; ++i ) {
      for ( size_t j = 0; j < clip_nlambda; ++j ) {
        newloc_f.Set( i * clip_nlambda + j, 0, loc_f.Get( i * psf_nlambda + clip_lambda_start + j, 0 ) );
      }
    }

    mpi_matrix new_f ( clip_nbins, 1, grid );
    mpi_matrix_zero ( new_f );

    globloc.Attach( El::LOCAL_TO_GLOBAL, new_f );
    if ( myp == 0 ) {
      globloc.Axpy ( 1.0, newloc_f, 0, 0 );
    }
    globloc.Detach();

    mpi_sparse_mv_trans ( AT, new_f, f_projected );

    vector_double u_f_projected;
    vector_double u_invnoise;

    if ( myp == 0 ) {

      elem_to_ublas ( f_projected, u_f_projected );
      f_projected.Empty();

      elem_to_ublas ( invnoise, u_invnoise );
      invnoise.Empty();

      outfile = outdir + "/harp_rpix-" + camera + "-" + expid + ".fits";
      image_desi::write ( outfile, solution_props, imgrows, u_f_projected, u_invnoise, msk );

    }

    loc_f.Empty();
    data_f.Empty();

    tstop = MPI_Wtime();

    if ( ( myp == 0 ) && ( ! quiet ) ) {
      cout << prefix << "  time = " << tstop-tstart << " seconds" << endl;
    }

    if ( dotruth ) {

      if ( ( myp == 0 ) && ( ! quiet ) ) {
        cout << prefix << "Writing projected truth spectra..." << endl;
      }

      tstart = MPI_Wtime();

      elem_matrix_local truth_projected;

      mpi_sparse_mv_trans ( AT, data_truth, truth_projected );

      // save more memory before allocating chisquare image

      AT.clear();
      data_truth.Empty();

      vector_double u_truth_projected;

      if ( myp == 0 ) {

        elem_to_ublas ( truth_projected, u_truth_projected );
        truth_projected.Empty();

        // pixel space chi-square

        vector_double data_chisq ( npix );

        double chisq_reduced = 0.0;

        for ( size_t i = 0; i < npix; ++i ) {
          if ( u_invnoise[i] > std::numeric_limits < double > :: epsilon() ) {
            data_chisq[i] = ( u_f_projected[i] - u_truth_projected[i] ) * sqrt( u_invnoise[i] );
            data_chisq[i] *= data_chisq[i];
          } else {
            data_chisq[i] = 0.0;
          }
          chisq_reduced += data_chisq[i];
        }

        chisq_reduced /= (double)( npix - 1 );

        if ( ! quiet ) {
          cout << prefix << "  Pixel space reduced Chi square = " << chisq_reduced << endl;
        }

        outfile = outdir + "/harp_rsimpix-" + camera + "-" + expid + ".fits";
        image_desi::write ( outfile, solution_props, imgrows, u_truth_projected, u_invnoise, msk );

        outfile = outdir + "/harp_chisqpix-" + camera + "-" + expid + ".fits";
        image_desi::write ( outfile, solution_props, imgrows, data_chisq, u_invnoise, msk );

      }

      tstop = MPI_Wtime();

      if ( ( myp == 0 ) && ( ! quiet ) ) {
        cout << prefix << "  time = " << tstop-tstart << " seconds" << endl;
      }

    }

  }

  // final timing

  double global_stop = MPI_Wtime();

  if ( ( myp == 0 ) && ( ! quiet ) ) {
    cout << prefix << "Total run time = " << global_stop-global_start << " seconds" << endl;
  }

  El::Finalize();

  return 0;
}


/*!
@file
Defines `boost::hana::Ring`.

@copyright Louis Dionne 2013-2016
Distributed under the Boost Software License, Version 1.0.
(See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
 */

#ifndef BOOST_HANA_CONCEPT_RING_HPP
#define BOOST_HANA_CONCEPT_RING_HPP

#include <boost/hana/fwd/concept/ring.hpp>

#include <boost/hana/config.hpp>
#include <boost/hana/core/default.hpp>
#include <boost/hana/core/tag_of.hpp>
#include <boost/hana/mult.hpp>
#include <boost/hana/one.hpp>

BOOST_HANA_NAMESPACE_BEGIN
    template <typename R>
    struct Ring {
        using Tag = typename tag_of<R>::type;
        static constexpr bool value = !is_default<one_impl<Tag>>::value &&
                                      !is_default<mult_impl<Tag, Tag>>::value;
    };
BOOST_HANA_NAMESPACE_END

#endif // !BOOST_HANA_CONCEPT_RING_HPP


//__BEGIN_LICENSE__
//  Copyright (c) 2009-2012, United States Government as represented by the
//  Administrator of the National Aeronautics and Space Administration. All
//  rights reserved.
//
//  The NGT platform is licensed under the Apache License, Version 2.0 (the
//  "License"); you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//  http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// __END_LICENSE__

#ifdef _MSC_VER
#pragma warning(disable:4244)
#pragma warning(disable:4267)
#pragma warning(disable:4996)
#endif

#include <iostream>
#include <fstream>
#include <iomanip>
#include <string>
#include <vector>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <time.h>
using namespace std;

#include <boost/program_options.hpp>
namespace po = boost::program_options;

#include <boost/filesystem/convenience.hpp>
namespace fs = boost::filesystem;

#include <vw/Core.h>
#include <vw/Image.h>
#include <vw/FileIO.h>
#include <vw/Cartography.h>
#include <vw/Math.h>
#include <photk/Photometry.h>

// ISIS
#include "Camera.h"
#include "CameraFactory.h"
#include "TProjection.h"
#include "ProjectionFactory.h"
#include "ProcessRubberSheet.h"
#include "IException.h"
#include "Pvl.h"
#include "IString.h"
#include "PushFrameCameraDetectorMap.h"
#include "Target.h"
#include "CameraDistortionMap.h"
#include "CameraFocalPlaneMap.h"
#include "SurfacePoint.h"
#include "Latitude.h"
#include "Longitude.h"

using namespace vw;
using namespace vw::math;
using namespace vw::cartography;
using namespace photometry;

// Compute mean and standard deviation of an image
template <class ImageT>
void compute_image_stats(ImageT const& I, double & mean, double & stdev){
  mean  = 0;
  stdev = 0;
  double sum = 0.0, sum2 = 0.0, count = 0.0;
  for (int col = 0; col < I.cols(); col++){
    for (int row = 0; row < I.rows(); row++){
      if (!is_valid(I(col, row))) continue;
      count++;
      double val = I(col, row).child();
      sum += val;
      sum2 += val*val;
    }
  }

  if (count > 0){
    mean = sum/count;
    stdev = sqrt( sum2/count - mean*mean );
  }

}

// Write an image together with georeference and nodata value.
template<class ImgT>
void write_image_nodata(std::string imgfile,
                        ImgT const& img, GeoReference geo, float nodata_val){

  std::cout << "Writing: " << imgfile << std::endl;
  DiskImageResourceGDAL::Options gdal_options;
  gdal_options["COMPRESS"] = "LZW";
  DiskImageResourceGDAL rsrc(imgfile, img.format(), Vector2i(256, 256),
                             gdal_options);
  rsrc.set_nodata_write(nodata_val);
  write_georeference(rsrc, geo);
  write_image(rsrc, img,
              TerminalProgressCallback("{Core}", "Processing:"));
}

double calc_max_abs(ImageView<PixelGray<float> >& DEM,
                    float DEMnodata){
  double mx = 0.0;
  for (int col = 0; col < DEM.cols(); col++){
    for (int row = 0; row < DEM.rows(); row++){
      if (DEM(col, row) == DEMnodata) continue;
      mx = std::max(mx, (double)fabs(DEM(col, row)));
    }
  }

  return mx;
}

void calc_min_max(ImageView<PixelGray<float> >& DEM,
                  float DEMnodata,
                  double & mn, double & mx){
  mn = DBL_MAX;
  mx = -DBL_MAX;
  for (int col = 0; col < DEM.cols(); col++){
    for (int row = 0; row < DEM.rows(); row++){
      if (DEM(col, row) == DEMnodata) continue;
      mx = std::max(mx, (double)DEM(col, row));
      mn = std::min(mn, (double)DEM(col, row));
    }
  }

}

// Utilities
string print_vec(vector<double> const& vec){
  ostringstream oss;
  oss.precision(18);
  oss << "(";
  for (int s = 0; s < (int)vec.size()-1; s++){
    oss << vec[s] << ", ";
  }
  if (!vec.empty()) oss << vec[vec.size()-1];
  oss << ")";
  return oss.str();
}

vector<double> diff(vector<double> const& vec1, vector<double> const& vec2){
  vector<double> ans = vec1;
  for (int i = 0; i < (int)vec1.size(); i++){
    ans[i] -= vec2[i];
  }
  return ans;
}

vector<double> prod(double val, vector<double> const& vec){
  vector<double> ans = vec;
  for (int i = 0; i < (int)vec.size(); i++){
    ans[i] *= val;
  }
  return ans;
}

vector<double> div(vector<double> const& vec, double val){
  vector<double> ans = vec;
  for (int i = 0; i < (int)vec.size(); i++){
    ans[i] /= val;
  }
  return ans;
}

vector<double> normalize(vector<double> const& vec){
  double len = 0;
  for (int i = 0; i < (int)vec.size(); i++){
    len += vec[i]*vec[i];
  }
  len = sqrt(len);
  vector<double> ans = vec;
  for (int i = 0; i < (int)vec.size(); i++){
    ans[i] /= len;
  }
  return ans;
}

// Abstract base class
class IsisInterface {
public:
  IsisInterface( string const& file );
  virtual ~IsisInterface(); // Can't be declared here since we have
  // incomplete types from Isis.

  virtual string type() = 0;
  static IsisInterface* open( string const& filename );

  // Standard Methods
  //------------------------------------------------------

  // These are the standard camera request; IsisInterface allows for
  // them to be customized for the type of camera so that they are
  // fast and not too full of conditionals.

  // Note: The smallest pixel is (0, 0), not (1, 1)!
  virtual vector<double>
  point_to_pixel( vector<double> const& point ) const = 0;
  virtual vector<double>
  pixel_to_vector( vector<double> const& pix ) const = 0;
  virtual vector<double>
  camera_center() const = 0;

protected:
  // Standard Variables
  //------------------------------------------------------
  Isis::Pvl * m_label;
  Isis::Camera * m_camera;
  Isis::Cube * m_cube;

};

// ISIS frame camera class
class IsisInterfaceFrame : public IsisInterface {

public:
  IsisInterfaceFrame( string const& filename );
  virtual ~IsisInterfaceFrame(){}

  virtual string type()  { return "Frame"; }

  // Standard Methods
  //-------------------------------------------------

  // Note: The smallest pixel is (0, 0), not (1, 1)!
  virtual vector<double>
  point_to_pixel( vector<double> const& point ) const;
  virtual vector<double>
  pixel_to_vector( vector<double> const& pix ) const;
  virtual vector<double>
  camera_center() const;

protected:

  // Custom Variables
  Isis::CameraDistortionMap *m_distortmap;
  Isis::CameraFocalPlaneMap *m_focalmap;
  Isis::CameraDetectorMap   *m_detectmap;
  mutable Isis::AlphaCube   m_alphacube;

  vector<double> m_center;
};

// ISIS linescan camera class
class IsisInterfaceLineScan : public IsisInterface {

public:
  IsisInterfaceLineScan( std::string const& filename );

  virtual ~IsisInterfaceLineScan() {}

  virtual std::string type()  { return "LineScan"; }

  // Standard Methods
  //-------------------------------------------------

  // Note: The smallest pixel is (0, 0), not (1, 1)!
  virtual vector<double>
  point_to_pixel( vector<double> const& point ) const;
  virtual vector<double>
  pixel_to_vector( vector<double> const& pix ) const;
  virtual vector<double>
  camera_center() const;

protected:

  // Custom Variables
  Isis::CameraDistortionMap *m_distortmap;
  Isis::CameraFocalPlaneMap *m_focalmap;
  Isis::CameraDetectorMap   *m_detectmap;
  mutable Isis::AlphaCube   m_alphacube; // Doesn't use const
};

// IsisInterface base class implementation

IsisInterface::IsisInterface( string const& file ) {
  // Opening labels and camera
  Isis::FileName ifilename( QString::fromStdString(file) );
  m_label = new Isis::Pvl();
  m_label->read( ifilename.expanded() );

  // Opening Isis::Camera
  m_cube = new Isis::Cube(QString::fromStdString(file));
  m_camera = Isis::CameraFactory::Create( *m_cube );
}

IsisInterface::~IsisInterface(){
  delete m_label;
  delete m_camera;
  delete m_cube;
}

IsisInterface* IsisInterface::open( string const& filename ) {
  IsisInterface* result;

  // Opening Labels (This should be done somehow though labels)
  Isis::FileName ifilename( QString::fromStdString(filename) );
  Isis::Pvl label;
  label.read( ifilename.expanded() );

  Isis::Cube tempCube(QString::fromStdString(filename));
  Isis::Camera* camera = Isis::CameraFactory::Create( tempCube );

  switch ( camera->GetCameraType() ) {
  case 0:
    // Framing Camera
    if ( camera->HasProjection() ){
      throw string("Don't support Isis Camera Type with projection at this moment");
      //result = new IsisInterfaceMapFrame( filename );
    }
    else{
      result = new IsisInterfaceFrame( filename );
    }
    break;
  case 2:
    // Linescan Camera
    if ( camera->HasProjection() ){
      throw string("Don't support Projected Linescan Isis Camera Type at this moment");
    }
    else{
      result = new IsisInterfaceLineScan( filename );
    }
    break;
  default:
    throw string("Don't support unknown Isis Camera Type at this moment");
  }
  return result;
}

// IsisInterfaceFrame class implementation

IsisInterfaceFrame::IsisInterfaceFrame( string const& filename ) :
  IsisInterface(filename), m_alphacube( *m_cube ) {

  // Gutting Isis::Camera
  m_distortmap = m_camera->DistortionMap();
  m_focalmap   = m_camera->FocalPlaneMap();
  m_detectmap  = m_camera->DetectorMap();

  // Calculating Center (just once)
  m_center.resize(3);
  m_camera->instrumentPosition(&m_center[0]);
  for (int i = 0; i < (int)m_center.size(); i++){
    m_center[i] *= 1000; // convert to meters
  }
}

vector<double>
IsisInterfaceFrame::point_to_pixel( vector<double> const& point ) const{

  // Note: The smallest pixel is (0, 0), not (1, 1)!

  vector<double> look = normalize( diff(point, m_center) );

  vector<double> lookB_copy(3);
  std::copy( look.begin(), look.end(), lookB_copy.begin() );
  lookB_copy = m_camera->bodyRotation()->J2000Vector(lookB_copy);
  lookB_copy = m_camera->instrumentRotation()->ReferenceVector(lookB_copy);
  std::copy( lookB_copy.begin(), lookB_copy.end(), look.begin() );
  look = div(prod(m_camera->FocalLength(), look), fabs(look[2]) );

  // Back Projecting
  m_distortmap->SetUndistortedFocalPlane( look[0], look[1] );
  m_focalmap->SetFocalPlane( m_distortmap->FocalPlaneX(),
                             m_distortmap->FocalPlaneY() );
  m_detectmap->SetDetector( m_focalmap->DetectorSample(),
                            m_focalmap->DetectorLine() );

  vector<double> ans(2);
  ans[0] = m_alphacube.BetaSample( m_detectmap->ParentSample() ) - 1;
  ans[1] = m_alphacube.BetaLine( m_detectmap->ParentLine() ) - 1;
  return ans;
}

vector<double>
IsisInterfaceFrame::pixel_to_vector( vector<double> const& px ) const {

  // Note: The smallest pixel is (0, 0), not (1, 1)!

  m_detectmap->SetParent( m_alphacube.AlphaSample(px[0]+1),
                          m_alphacube.AlphaLine(px[1]+1));
  m_focalmap->SetDetector( m_detectmap->DetectorSample(),
                           m_detectmap->DetectorLine() );
  m_distortmap->SetFocalPlane( m_focalmap->FocalPlaneX(),
                               m_focalmap->FocalPlaneY() );

  vector<double> look(3);
  look[0] = m_distortmap->UndistortedFocalPlaneX();
  look[1] = m_distortmap->UndistortedFocalPlaneY();
  look[2] = m_distortmap->UndistortedFocalPlaneZ();

  look = normalize( look );
  look = m_camera->instrumentRotation()->J2000Vector(look);
  look = m_camera->bodyRotation()->ReferenceVector(look);

  return look;
}

vector<double> IsisInterfaceFrame::camera_center() const{
  return m_center;
}

// IsisInterfaceLineScan class implementation
IsisInterfaceLineScan::IsisInterfaceLineScan( std::string const& filename ):
  IsisInterface(filename), m_alphacube( *m_cube ) {

  // Gutting Isis::Camera
  m_distortmap = m_camera->DistortionMap();
  m_focalmap   = m_camera->FocalPlaneMap();
  m_detectmap  = m_camera->DetectorMap();
}

vector<double>
IsisInterfaceLineScan::point_to_pixel( vector<double> const& point ) const{

  // Note: The smallest pixel is (0, 0), not (1, 1)!
  using namespace Isis;
  SurfacePoint surfP(Displacement(point[0], Displacement::Meters),
                     Displacement(point[1], Displacement::Meters),
                     Displacement(point[2], Displacement::Meters));

  vector<double> pix(2);
  if(m_camera->SetUniversalGround(surfP.GetLatitude().degrees(),
                                  surfP.GetLongitude().degrees(),
                                  surfP.GetLocalRadius().meters()
                                  )) {
    pix[0] = m_camera->Sample()-1;
    pix[1] = m_camera->Line()-1;

    return pix;
  }

  throw string("Could not project point into the camera.");
  return pix;
}

vector<double>
IsisInterfaceLineScan::pixel_to_vector( vector<double> const& px ) const {

  // Note: The smallest pixel is (0, 0), not (1, 1)!

  m_detectmap->SetParent( m_alphacube.AlphaSample(px[0]+1),
                          m_alphacube.AlphaLine(px[1]+1));
  m_focalmap->SetDetector( m_detectmap->DetectorSample(),
                           m_detectmap->DetectorLine() );
  m_distortmap->SetFocalPlane( m_focalmap->FocalPlaneX(),
                               m_focalmap->FocalPlaneY() );

  vector<double> look(3);
  look[0] = m_distortmap->UndistortedFocalPlaneX();
  look[1] = m_distortmap->UndistortedFocalPlaneY();
  look[2] = m_distortmap->UndistortedFocalPlaneZ();

  look = normalize( look );
  look = m_camera->instrumentRotation()->J2000Vector(look);
  look = m_camera->bodyRotation()->ReferenceVector(look);

  return look;
}

vector<double> IsisInterfaceLineScan::camera_center() const{
  throw string("Computation of camera center is not implemented")
    + " for linescan cameras.";
}

// Convolve a DEM with exp(-sigma*x^2). The input DEM must
// have its invalid pixels masked.

template <class ImageT>
class BlurDEM:
  public ImageViewBase< BlurDEM<ImageT> > {
  ImageT m_img;
  int m_search_dist; // half of size of window to convolve with
  ImageView<double> m_gauss_kernel;
  typedef typename ImageT::pixel_type PixelT;

public:

  typedef PixelT pixel_type;
  typedef PixelT result_type;
  typedef ProceduralPixelAccessor<BlurDEM> pixel_accessor;

  BlurDEM( ImageViewBase<ImageT> const& img,
                     double blur_sigma) :
    m_img(img.impl()) {
    VW_ASSERT(blur_sigma > 0,
              ArgumentErr() << "Expecting positive sigma.");

    // Cut the gaussian exp(-sigma*x^2) where its value is 'scale'.
    double scale = 0.001;
    m_search_dist = (int)ceil(sqrt(-log(scale)/blur_sigma));
    std::cout << "Blur search distance is " << m_search_dist << std::endl;

    // The gaussian kernel
    int h = m_search_dist;
    m_gauss_kernel.set_size(2*h+1, 2*h+1);
    for (int c = 0; c < m_gauss_kernel.cols(); c++){
      for (int r = 0; r < m_gauss_kernel.rows(); r++){
        double r2 = double(c-h)*(c-h) + double(r-h)*(r-h);
        m_gauss_kernel(c, r) = exp(-blur_sigma*r2);
      }
    }

  }

  inline int32 cols() const { return m_img.cols(); }
  inline int32 rows() const { return m_img.rows(); }
  inline int32 planes() const { return 1; }

  inline pixel_accessor origin() const { return pixel_accessor(*this); }

  inline result_type operator()( size_t i, size_t j, size_t p=0 ) const {
    vw_throw( NoImplErr() << "BlurDEM: operator() not implemented.\n" );
  }

  typedef CropView< ImageView<PixelT> > prerasterize_type;
  inline prerasterize_type prerasterize( BBox2i const& bbox ) const {

    // Crop into an expanded box as to have enough pixels to do
    // the blurring at every pixel in the current box.
    int h = m_search_dist; // shorten
    BBox2i biased_box = bbox;
    biased_box.expand(h+1);
    biased_box.crop(bounding_box(m_img));
    ImageView<PixelT> img( crop( m_img, biased_box ) );
    ImageView<PixelT> filled_img = copy(img);

    int nc = img.cols(), nr = img.rows(); // shorten
    for (int row = 0; row < nr; row++){
      for (int col = 0; col < nc; col++){
        PixelT V; V.validate();
        double sum = 0.0;
        for (int c = std::max(col-h, 0); c <= std::min(col+h, nc-1); c++){
          for (int r = std::max(row-h, 0); r <= std::min(row+h, nr-1); r++){
            if (!is_valid(img(c, r))) continue;
            double wt = m_gauss_kernel(c-col+h, r-row+h);
              V   += wt*img(c, r);
              sum += wt;
          }
        }
        if (sum > 0) filled_img(col, row) = V/sum;

      }
    }

    return prerasterize_type(filled_img,
                             -biased_box.min().x(), -biased_box.min().y(),
                             cols(), rows());

  }
  template <class ImgT>
  inline void rasterize( ImgT const& img, BBox2i const& bbox ) const {
    vw::rasterize( prerasterize(bbox), img, bbox );
  }

};

template <class ImgT>
BlurDEM<ImgT>
blur_dem( ImageViewBase<ImgT> const& img, double blur_sigma) {
  typedef BlurDEM<ImgT> result_type;
  return result_type( img.impl(), blur_sigma);
}

template<class ImageT>
void calc_image_and_reflectance_at_pt(ImageView<PixelGray<float> > const& DEM,
                                      ImageT const& img,
                                      GeoReference const& DEMGeo,
                                      float DEMnodata,
                                      ModelParams const & img_params,
                                      GlobalParams const& globalParams,
                                      int col, int row,
                                      IsisInterface const* cam,
                                      PixelMask<float> & imgval,
                                      PixelMask<PixelGray<float> > & refl
                                      ){

  refl.invalidate();
  imgval.invalidate();

  double h = DEM(col, row);
  if (h == DEMnodata) return;

  ImageViewRef<typename ImageT::pixel_type> interp_img
    = interpolate(img,
                  BicubicInterpolation(), ZeroEdgeExtension());

  Vector2 lonlat = DEMGeo.pixel_to_lonlat(Vector2(col, row));
  Vector3 llh(lonlat[0], lonlat[1], h);
  Vector3 xyz = DEMGeo.datum().geodetic_to_cartesian(llh);

  // Switch from VW vector to vector<double>
  vector<double> xyz2(3);
  for (int k = 0; k < (int)xyz.size(); k++) xyz2[k] = xyz[k];

  vector<double> cpix = cam->point_to_pixel(xyz2);

  if (cpix[0] < 1 || cpix[0] > interp_img.cols() - 2 ||
      cpix[1] < 1 || cpix[1] > interp_img.rows() - 2){
    imgval.invalidate();
  }else{
    imgval = interp_img(cpix[0], cpix[1]);
  }

  bool savePhaseAngle = false;
  float outputReflectanceVal, phaseAngleVal;
  computeReflectanceAtPixel(col, row, DEM, DEMGeo, DEMnodata, img_params,
                            globalParams, savePhaseAngle,
                            outputReflectanceVal,
                            phaseAngleVal);
  refl = outputReflectanceVal;
  if (refl.child() <= 0){
    //imgval.invalidate();
    refl.invalidate();
  }

}

void dump_iter(int iter,
               double a, double b,
               GeoReference const& DEMGeo,
               ImageView<PixelMask<float> > const& I,
               ImageView<PixelMask<PixelGray<float> > > const& R
               ){

  // The image intensity at every DEM point (I), and the reflectance
  // at every DEM point (R) are floating-point quantities.  Towards
  // the end of optimization we will have I is about a*R + b.

  // Here, we will clamp I to [mean(I) - sigma*stdev(I), mean(I) + sigma*stddev(I)],
  // and cast to int. We will clamp a*R + b to the same interval for easy comparison.

  double imgmean, imgstdev;
  compute_image_stats(I, imgmean, imgstdev);

  double sigma = 5.0, mn = imgmean - sigma*imgstdev, mx = imgmean + sigma*imgstdev;
  ostringstream os; os << "_iter" << iter << "_int.tif";
  float nodata = 0;
  ImageView<uint8> V(I.cols(), I.rows());

  for (int col = 0; col < V.cols(); col++){
    for (int row = 0; row < V.rows(); row++){
      int v;
      if (!is_valid(I(col, row)))
        v = 0;
      else
        v = (int)round(255*(I(col, row).child() - mn)/(mx - mn));
      v = std::max(0, v);
      v = std::min(255, v);
      V(col, row) = uint8(v);
    }
  }
  write_image_nodata("demimg" + os.str(), V, DEMGeo, nodata);

  for (int col = 0; col < V.cols(); col++){
    for (int row = 0; row < V.rows(); row++){
      int v;
      if (!is_valid(R(col, row)))
        v = 0;
      else
        v = (int)round(255*((a*R(col, row).child()+b) - mn)/(mx - mn));
      v = std::max(0, v);
      v = std::min(255, v);
      V(col, row) = uint8(v);
    }
  }
  write_image_nodata("demrefl" + os.str(), V, DEMGeo, nodata);

}

double calc_cost_fun(ImageView<PixelGray<float> >& DEM,
                     GeoReference const& DEMGeo,
                     float DEMnodata,
                     ModelParams const& img_params,
                     GlobalParams const& globalParams,
                     IsisInterface const* cam,
                     ImageView<float> & img,
                     double a, double b,
                     ImageView<PixelMask<float> > & I,
                     ImageView<PixelMask<PixelGray<float> > > & R
                     ){

  // Compute the cost function.

  I.set_size(DEM.cols(), DEM.rows());
  R.set_size(DEM.cols(), DEM.rows());
  for (int col = 0; col < DEM.cols(); col++){
    for (int row = 0; row < DEM.rows(); row++){
      R(col, row).invalidate();
      I(col, row).invalidate();
    }
  }

  double cost = 0.0;

  // We will keep the border pixels fixed
  for (int col = 1; col < DEM.cols()-1; col++){
    for (int row = 1; row < DEM.rows()-1; row++){

      if (DEM(col, row) == DEMnodata) continue;

      bool bad_value = false;
      double cost2 = 0.0;
      for (int c = col-1; c <= col+1; c++){
        for (int r = row-1; r <= row+1; r++){

          PixelMask<float> imgval;
          PixelMask<PixelGray<float> > refl;
          calc_image_and_reflectance_at_pt(DEM, img,
                                           DEMGeo, DEMnodata,
                                           img_params, globalParams,
                                           c, r, cam,
                                           imgval, refl
                                           );

          if (!is_valid(imgval) || !is_valid(refl)){
            bad_value = true;
            continue;
          }
          if (c == col && r == row){
            I(c, r) = imgval;
            R(c, r) = refl; //a*refl + b;
          }

          double v = imgval.child() - a*refl.child() - b;
          cost2 += v*v;
        }
      }
      if (!bad_value){
        cost += cost2;
      }

    }
  }

  return cost;
}

void calc_update(ImageView<PixelGray<float> >& DEM,
                 GeoReference const& DEMGeo,
                 float DEMnodata,
                 ModelParams const& img_params,
                 GlobalParams const& globalParams,
                 IsisInterface const* cam,
                 ImageView<float> & img,
                 double a, double b, double deltah,
                 ImageView<PixelGray<float> > & update
                 ){

  // Calculate the update via gradient descent (the gradient is
  // computed numerically).

  update.set_size(DEM.cols(), DEM.rows());
  for (int col = 0; col < DEM.cols(); col++){
    for (int row = 0; row < DEM.rows(); row++){
      update(col, row) = 0;
    }
  }

  // We will keep the border pixels fixed
  for (int col = 1; col < DEM.cols()-1; col++){
    for (int row = 1; row < DEM.rows()-1; row++){

      if (DEM(col, row) == DEMnodata) continue;

      bool bad_value = false;

      double cost = 0.0;
      for (int c = col-1; c <= col+1; c++){
        for (int r = row-1; r <= row+1; r++){

          PixelMask<float> imgval;
          PixelMask<PixelGray<float> > refl;
          calc_image_and_reflectance_at_pt(DEM, img,
                                           DEMGeo, DEMnodata,
                                           img_params, globalParams,
                                           c, r, cam,
                                           imgval, refl
                                           );

          if (!is_valid(imgval) || !is_valid(refl)){
            bad_value = true;
            continue;
          }
          double v = imgval.child() - a*refl.child() - b;
          cost += v*v;
        }
      }

      // Now bias the DEM at current location
      DEM(col, row) += deltah;

      double cost2 = 0.0;
      for (int c = col-1; c <= col+1; c++){
        for (int r = row-1; r <= row+1; r++){

          PixelMask<float> imgval;
          PixelMask<PixelGray<float> > refl;
          calc_image_and_reflectance_at_pt(DEM, img,
                                           DEMGeo, DEMnodata,
                                           img_params, globalParams,
                                           c, r, cam,
                                           imgval, refl
                                           );

          if (!is_valid(imgval) || !is_valid(refl)){
            bad_value = true;
            continue;
          }
          double v = imgval.child() - a*refl.child() - b;
          cost2 += v*v;
        }
      }

      // Must bias the DEM back
      DEM(col, row) -= deltah;

      if (!bad_value){
        update(col, row) = (cost2-cost)/deltah;
      }

    }
  }

}

void calc_image_and_reflectance(ImageView<PixelGray<float> >& DEM,
                                GeoReference const& DEMGeo,
                                float DEMnodata,
                                ModelParams const& img_params,
                                GlobalParams const& globalParams,
                                int num_iter
                                ){

  std::string imgfile = img_params.inputFilename;

  boost::shared_ptr<IsisInterface> cam(IsisInterface::open(imgfile));

  boost::shared_ptr<DiskImageResource>
    img_rsrc( DiskImageResource::open(imgfile));
  ImageView<float> img = copy(DiskImageView<float>( img_rsrc ));

  // Calculate a and b in the camera transfer function model
  // I = a*R + b.
  double a = 1.0, b = 0.0;
  ImageView<PixelMask<float> > I;
  ImageView<PixelMask<PixelGray<float> > > R;
  calc_cost_fun(DEM, DEMGeo, DEMnodata,
                img_params,  globalParams, cam.get(), img, a, b,
                I, R
                );
  double imgmean, imgstdev, refmean, refstdev;
  compute_image_stats(I, imgmean, imgstdev);
  compute_image_stats(R, refmean, refstdev);
  a = imgstdev/refstdev;
  b = imgmean - a*refmean;

  // The cost fun is sum (I-a*R-b)^2. Each time the height of a DEM
  // point is changed, I and R change only at that point and its
  // immediate neighbors.

  // Will vary height by this much to compute the numerical gradient
  double max_dem = calc_max_abs(DEM, DEMnodata);
  double deltah = 0.00001*max_dem;
  std::cout << "deltah is " << deltah << std::endl;

  // Initialize the update
  ImageView<PixelGray<float> > update;
  calc_update(DEM, DEMGeo, DEMnodata,  img_params,
              globalParams, cam.get(), img, a, b, deltah, update
              );

  // Determine how to scale the updates, so they are not too big or
  // too small.
  double max_up  = calc_max_abs(update, 0.0);
  double mn, mx;
  calc_min_max(DEM, DEMnodata, mn, mx);
  double pct = 0.05;
  double factor = pct*(mx - mn)/max_up;
  if (factor == 0){
    // Handle the case when the input DEM is constant
    factor = 2.0/max_up; // 2 meter update
  }

  double cost = calc_cost_fun(DEM, DEMGeo, DEMnodata,
                              img_params,  globalParams, cam.get(), img,  a, b,
                              I, R
                              );

  vector<double> Cost;
  Cost.push_back(cost);
  ofstream cs("cost.txt");
  cs << cost << std::endl;

  dump_iter(0, a, b, DEMGeo, I, R);

  for (int i = 1; i <= num_iter; i++){

    std::cout << "iteration: " << i << std::endl;

    double sigma = 1.5;
    ImageView<PixelGray<float> > blurredDEM
      = apply_mask(blur_dem
                   (create_mask(DEM, DEMnodata), sigma),
                   DEMnodata);
    DEM = copy(blurredDEM);

    // Cost after blur
    cost = calc_cost_fun(DEM, DEMGeo, DEMnodata,
                         img_params,  globalParams, cam.get(), img, a, b,
                         I, R
                         );
    Cost.push_back(cost);
    cs << cost << std::endl;

    calc_update(DEM, DEMGeo, DEMnodata,  img_params,
                globalParams, cam.get(), img, a, b, deltah, update
                );

    double max_up = 0.0;
    for (int col = 1; col < DEM.cols()-1; col++){
      for (int row = 1; row < DEM.rows()-1; row++){
        if (DEM(col, row) == DEMnodata) continue;
        double u = factor*update(col, row);
        DEM(col, row) -= u;
        max_up = std::max(max_up, u);
      }
    }

    std::cout << "max update is " << max_up << std::endl;

    // Cost after update
    cost = calc_cost_fun(DEM, DEMGeo, DEMnodata,
                         img_params,  globalParams, cam.get(), img,  a, b,
                         I, R
                         );
    Cost.push_back(cost);
    cs << cost << std::endl;

    if (i%10 == 0) dump_iter(i, a, b, DEMGeo, I, R);

  }

  cs.close();

  dump_iter(num_iter, a, b, DEMGeo, I, R);
}

int main( int argc, char *argv[] ) {

  // Do shape-from shading. See the INSTALL file about how to compile
  // and run this program.

  // Important: The step size for gradient descent needs tweaking for
  // best performance.  Ideally it would be adaptive, growing or
  // shrinking depending on how the cost function is decreasing.

  // To do: Remove the unneeded initial blur

  if (argc < 5){
    std::cerr << "Usage: " << argv[0] << " image.cub DEM.tif metadata_dir num_iter" << std::endl;
    exit(1);
  }

  vector<string> images;
  images.push_back(argv[1]);

  std::string DEMFile = argv[2];
  std::string meta_dir = argv[3];
  int num_iter = atoi(argv[4]);

  ImageView<PixelGray<float> > DEMTile = copy(DiskImageView<PixelGray<float> >(DEMFile));
  GeoReference DEMGeo;
  read_georeference(DEMGeo, DEMFile);
  float DEMnodata = -32768;
  if (readNoDEMDataVal(DEMFile, DEMnodata)){
    std::cout << "Found DEM nodata value: " << DEMnodata << std::endl;
  }

  GlobalParams globalParams;
  globalParams.sunPosFile = meta_dir + "/sunpos.txt";
  globalParams.spacecraftPosFile = meta_dir + "/spacecraftpos.txt";
  globalParams.reflectanceType = LUNAR_LAMBERT;
  globalParams.phaseCoeffC1    = 1.383488;
  globalParams.phaseCoeffC2    = 0.501149;
  //PrintGlobalParams(globalParams);

  std::map<std::string, Vector3> sunPositions;
  std::map<std::string, Vector3> spacecraftPositions;
  ReadSunOrSpacecraftPosition(globalParams.sunPosFile, // Input
                              sunPositions             // Output
                              );
  ReadSunOrSpacecraftPosition(globalParams.spacecraftPosFile, // Input
                              spacecraftPositions             // Output
                              );

  std::vector<ModelParams> modelParamsArray;
  std::cout.precision(18);
  modelParamsArray.resize(images.size());
  for (int k = 0; k < (int)images.size(); k++){
    std::string prefix = getFirstElevenCharsFromFileName(images[k]);

    modelParamsArray[k].inputFilename = images[k];

    if ( sunPositions.find(prefix) == sunPositions.end()){
      std::cerr << "Could not find the sun position for the image file: "
                << images[k] << std::endl;
      exit(1);
    }
    // Go from kilometers to meters
    modelParamsArray[k].sunPosition = 1000*sunPositions[prefix];
    std::cout << "sun position: " <<  modelParamsArray[k].inputFilename
              << ' ' <<  modelParamsArray[k].sunPosition << std::endl;

    if (spacecraftPositions.find(prefix) == spacecraftPositions.end()){
      std::cerr << "Could not find the spacecraft position for the DRG file: "
                << images[k] << std::endl;
      exit(1);
    }
    // Go from kilometers to meters
    modelParamsArray[k].spacecraftPosition = 1000*spacecraftPositions[prefix];
    std::cout << "spacecraft position: " <<  modelParamsArray[k].inputFilename
              << ' ' <<  modelParamsArray[k].spacecraftPosition << std::endl;
  }

#if 0
  // If to start with a constant initial guess
  for (int col = 0; col < DEMTile.cols(); col++){
    for (int row = 0; row < DEMTile.rows(); row++){
      DEMTile(col, row) = mean;
    }
  }
#endif

  calc_image_and_reflectance(DEMTile, DEMGeo,
                             DEMnodata,
                             modelParamsArray[0],
                             globalParams,
                             num_iter
                             );
  return 0;
}


#include <cstddef>
#include <limits>
#include <complex>
#include <tuple>
#include <exception>
#include <fstream>
#include <sstream>
#include <map>
#include <boost/optional.hpp>
#include <gtest/gtest.h>
#include <mpi.h>
#include <petscksp.h>
#include <json.hpp>
#include "util/MPIPrettyUnitTestResultPrinter.h"
#include "util/util.h"
#include "util/lattice.h"
#include "grid_basis.h"
#include "fermi_surface.h"
#include "dyn_dm_graph.h"
#include "berry.h"
#include "driving.h"
#include "models/wannier90_Hamiltonian.h"
#include "observables/energy.h"
#include "observables/rho0.h"
#include "observables/current.h"
#include "disorder/disorder.h"
#include "disorder/collision.h"

using json = nlohmann::json;

int main(int argc, char* argv[]) {
  ::testing::InitGoogleTest(&argc, argv);
  PetscInitialize(&argc, &argv, nullptr, nullptr);

  int rank;
  MPI_Comm_rank(PETSC_COMM_WORLD, &rank);

  ::testing::TestEventListeners& listeners = ::testing::UnitTest::GetInstance()->listeners();

  testing::TestEventListener *default_rp = listeners.default_result_printer();
  listeners.Release(default_rp);
  listeners.Append(new anomtrans::MPIPrettyUnitTestResultPrinter(default_rp, rank));

  int test_result = RUN_ALL_TESTS();

  int ierr = PetscFinalize();CHKERRXX(ierr);

  return test_result;
}

/** @brief Check that extraction of Wannier90 hr.dat file is performed correctly.
 */
TEST( Wannier90_hr_load, Wannier90_hr_load ) {
  boost::optional<std::string> test_data_dir = anomtrans::getenv_optional("ANOMTRANS_TEST_DATA_DIR");
  if (not test_data_dir) {
    throw std::runtime_error("Could not get ANOMTRANS_TEST_DATA_DIR environment variable for regression test data");
  }

  std::stringstream hr_path;
  hr_path << *test_data_dir << "/WSe2/wannier/WSe2_hr.dat";

  const std::size_t k_dim = 2;
  anomtrans::kComps<k_dim> Nk = {8, 8};
  anomtrans::DimMatrix<k_dim> D = {{{1.659521, 1.659521},
                                   {-2.874374, 2.874374}}};

  auto tb = anomtrans::extract_Wannier90_Hamiltonian(hr_path.str(), Nk, D);

  // Check that Nbands and Nr are correct.
  std::size_t Nbands_expected = 22;
  std::size_t Nr_expected = 91;

  ASSERT_EQ( tb.kmb.Nbands, Nbands_expected );
  ASSERT_EQ( tb.Hrs.size(), Nr_expected );

  // Check some values to ensure they loaded correctly.
  std::map<anomtrans::LatVec<k_dim>, unsigned int> expected_degens {
    {{-6, -3}, 3}, {{-5, -4}, 2}
  };

  using HrElem = std::tuple<anomtrans::LatVec<k_dim>, std::size_t, std::size_t>;

  std::map<HrElem, std::complex<double>> expected_values {
    {HrElem{{-6, -3}, 0, 0}, std::complex<double>(0.001540, 0.0)},
    {HrElem{{-6, -3}, 5, 0}, std::complex<double>(0.0, -0.000101)},
    {HrElem{{-5, -4}, 0, 0}, std::complex<double>(0.000448, 0.0)}
  };

  auto macheps = std::numeric_limits<double>::epsilon();
  for (auto it = expected_values.begin(); it != expected_values.end(); ++it) {
    anomtrans::LatVec<k_dim> r;
    std::size_t ip, i;
    std::tie(r, ip, i) = it->first;
    unsigned int degen = expected_degens[r];
    std::complex<double> expected_val = it->second / static_cast<double>(degen);

    std::complex<double> elem = tb.Hrs.at(r)(ip, i);

    ASSERT_TRUE( anomtrans::scalars_approx_equal(elem, expected_val, 10.0*macheps, 10.0*macheps) );
  }

  // Check that H(R) = H(-R)^{\dagger}, as required for Hermiticity.
  for (auto it = tb.Hrs.begin(); it != tb.Hrs.end(); ++it) {
    anomtrans::LatVec<k_dim> r = it->first;
    anomtrans::LatVec<k_dim> minus_r = {-r.at(0), -r.at(1)};

    ASSERT_TRUE( it->second.isApprox(tb.Hrs.at(minus_r).adjoint(), 10.0*macheps) );
  }
}

/** @brief Dynamic electric response to second order with WSe2 model from Wannier90 hr.dat.
 */
TEST( Wannier90_WSe2_dynamic, DISABLED_Wannier90_WSe2_dynamic ) {
  boost::optional<std::string> test_data_dir = anomtrans::getenv_optional("ANOMTRANS_TEST_DATA_DIR");
  if (not test_data_dir) {
    throw std::runtime_error("Could not get ANOMTRANS_TEST_DATA_DIR environment variable for regression test data");
  }

  std::stringstream hr_path;
  hr_path << *test_data_dir << "/WSe2/wannier/WSe2_hr.dat";

  const std::size_t k_dim = 2;
  anomtrans::kComps<k_dim> Nk = {4, 4};

  anomtrans::DimMatrix<k_dim> D = {{{1.659521, 1.659521}, // Angstrom
                                   {-2.874374, 2.874374}}};

  auto H = anomtrans::extract_Wannier90_Hamiltonian(hr_path.str(), Nk, D);
  const auto& kmb = H.kmb;
  const auto Nbands = H.kmb.Nbands;

  // Choose berry_broadening ~ optical broadening.
  // What is appropriate value?
  double berry_broadening = 1e-3; // eV

  PetscReal max_energy_difference = anomtrans::find_max_energy_difference(kmb, H);
  double beta_max = anomtrans::get_beta_max(max_energy_difference);
  double beta = beta_max / 2.0;

  if (beta > beta_max) {
    PetscPrintf(PETSC_COMM_WORLD, "Warning: beta > beta_max: beta = %e ; beta_max = %e\n", beta, beta_max);
  }

  double sigma_min = anomtrans::DeltaGaussian::get_sigma_min(max_energy_difference);
  double sigma = 2.0 * sigma_min;

  if (sigma < sigma_min) {
    PetscPrintf(PETSC_COMM_WORLD, "Warning: sigma < sigma_min: sigma = %e ; sigma_min = %e\n", sigma, sigma_min);
  }

  anomtrans::DeltaGaussian delta(sigma);

  // U0 = how far can bands be driven from their average energy?
  double U0 = 1e-3; // eV

  auto Ekm = anomtrans::get_energies(kmb, H);

  auto v_op = anomtrans::calculate_velocity(kmb, H);
  auto spin_op = anomtrans::calculate_spin_operator(kmb, H);

  std::size_t Nk_tot = anomtrans::get_Nk_total(Nk);
  for (std::size_t d = 0; d < k_dim; d++) {
    Nk_tot *= kmb.Nk.at(d);
  }
  double U0_sq = U0*U0;
  double disorder_coeff = U0_sq / Nk_tot;

  auto disorder_term = [Nbands, H, disorder_coeff](PetscInt ikm1, PetscInt ikm2)->double {
    return disorder_coeff*anomtrans::on_site_diagonal_disorder_band_preserved(Nbands, H,
        ikm1, ikm2);
  };
  auto disorder_term_od = [Nbands, H, disorder_coeff](PetscInt ikm1, PetscInt ikm2,
      PetscInt ikm3)->std::complex<double> {
    return disorder_coeff*anomtrans::on_site_diagonal_disorder(Nbands, H, ikm1, ikm2, ikm3);
  };

  auto collision = anomtrans::make_collision(kmb, H, disorder_term, delta);

  // Create the linear solver context.
  KSP ksp;
  PetscErrorCode ierr = KSPCreate(PETSC_COMM_WORLD, &ksp);CHKERRXX(ierr);
  // This uses collision again as the preconditioning matrix.
  // TODO - is there a better choice?
  ierr = KSPSetOperators(ksp, collision.first.M, collision.first.M);CHKERRXX(ierr);
  // Could use KSPSetFromOptions here. In this case, prefer to keep options
  // hard-coded to have identical output from each test run.

  const unsigned int deriv_approx_order = 2;
  anomtrans::DerivStencil<1> stencil(anomtrans::DerivApproxType::central, deriv_approx_order);
  auto d_dk_Cart = anomtrans::make_d_dk_Cartesian(D, kmb, stencil);

  // E in x and y directions
  std::array<double, k_dim> Ex = {1.0, 0.0};
  std::array<double, k_dim> Ey = {0.0, 1.0};

  // Maximum number of elements expected for sum of Cartesian derivatives.
  PetscInt Ehat_grad_expected_elems_per_row = stencil.approx_order*k_dim*k_dim*k_dim;

  auto Ex_dot_grad_k = anomtrans::Mat_from_sum_const(anomtrans::make_complex_array(Ex),
      anomtrans::unowned(d_dk_Cart), Ehat_grad_expected_elems_per_row);
  auto Ey_dot_grad_k = anomtrans::Mat_from_sum_const(anomtrans::make_complex_array(Ey),
      anomtrans::unowned(d_dk_Cart), Ehat_grad_expected_elems_per_row);

  auto R = anomtrans::make_berry_connection(kmb, H, berry_broadening);
  auto Ex_dot_R = anomtrans::Mat_from_sum_const(anomtrans::make_complex_array(Ex),
      anomtrans::unowned(R), kmb.Nbands);
  auto Ey_dot_R = anomtrans::Mat_from_sum_const(anomtrans::make_complex_array(Ey),
      anomtrans::unowned(R), kmb.Nbands);

  // Chemical potential in the gap.
  // Valence band maximum at K: -0.504 eV, Gamma: -1.005 eV.
  // VBM S-O splitting at K: -0.807 eV - (-1.276 eV) = 0.469 eV.
  // Conduction band minimum at K: 0.762 eV, Q: 0.798 eV.
  // CBM S-O splitting at K: 0.484 eV - 0.444 eV = 0.04 eV.
  double mu = 0.0;

  // Gap = 1.266 eV.
  // Choose omega \approx gap so that we produce resonant excitations from VBM to CBM.
  // TODO - interested in what happens when omega is below the gap.
  double omega = 1.266 + 0.004;

  // Equilibrium density matrix <rho_{0,0}>. Make two since we will have two DM graphs,
  // one for E ~ \hat{x} cos(\omega t) and one for E ~ \hat{y} sin(\omega t).
  auto dm_rho0_Ex_cos = anomtrans::make_eq_node<anomtrans::DynDMGraphNode>(Ekm.v, beta, mu);
  auto dm_rho0_Ey_sin = anomtrans::make_eq_node<anomtrans::DynDMGraphNode>(Ekm.v, beta, mu);

  auto rho0_km = anomtrans::make_Vec_with_structure(Ekm.v);
  ierr = MatGetDiagonal(dm_rho0_Ex_cos->rho.M, rho0_km.v);CHKERRXX(ierr);

  // Get normalized version of rho0 to use for nullspace.
  // TODO can we safely pass a nullptr instead of rho0_orig_norm?
  auto rho0_normalized = anomtrans::make_Vec_with_structure(rho0_km.v);
  ierr = VecCopy(rho0_km.v, rho0_normalized.v);CHKERRXX(ierr);

  PetscReal rho0_orig_norm;
  ierr = VecNormalize(rho0_normalized.v, &rho0_orig_norm);CHKERRXX(ierr);

  // Set nullspace of K: K rho0_km = 0.
  // Note that this is true regardless of the value of mu
  // (energies only enter K through differences).
  // It is also true for any value of beta (the Fermi-Dirac distribution
  // function does not appear in K, only energy differences).
  // TODO does this mean that the nullspace has dimension larger than 1?
  MatNullSpace nullspace;
  ierr = MatNullSpaceCreate(PETSC_COMM_WORLD, PETSC_FALSE, 1, &(rho0_normalized.v), &nullspace);CHKERRXX(ierr);
  ierr = MatSetNullSpace(collision.first.M, nullspace);CHKERRXX(ierr);
  // NOTE rho0_normalized must not be modified after this call until we are done with nullspace.

  // Add <rho_{1,1}> and <rho_{-1,1}> for (Ex, cos) and (Ey, sin) cases.
  anomtrans::add_dynamic_electric_n_nonzero(dm_rho0_Ex_cos, boost::none, 1, omega, kmb, Ex_dot_grad_k.M,
      Ex_dot_R.M, H, berry_broadening, anomtrans::DynVariation::cos);
  anomtrans::add_dynamic_electric_n_nonzero(boost::none, dm_rho0_Ex_cos, -1, omega, kmb, Ex_dot_grad_k.M,
      Ex_dot_R.M, H, berry_broadening, anomtrans::DynVariation::cos);
  auto dm_rho_1_1_Ex_cos = dm_rho0_Ex_cos->children[anomtrans::DynDMDerivedBy::omega_inv_DE_up];
  auto dm_rho_m1_1_Ex_cos = dm_rho0_Ex_cos->children[anomtrans::DynDMDerivedBy::omega_inv_DE_down];

  anomtrans::add_dynamic_electric_n_nonzero(dm_rho0_Ey_sin, boost::none, 1, omega, kmb, Ey_dot_grad_k.M,
      Ey_dot_R.M, H, berry_broadening, anomtrans::DynVariation::sin);
  anomtrans::add_dynamic_electric_n_nonzero(boost::none, dm_rho0_Ey_sin, -1, omega, kmb, Ey_dot_grad_k.M,
      Ey_dot_R.M, H, berry_broadening, anomtrans::DynVariation::sin);
  auto dm_rho_1_1_Ey_sin = dm_rho0_Ey_sin->children[anomtrans::DynDMDerivedBy::omega_inv_DE_up];
  auto dm_rho_m1_1_Ey_sin = dm_rho0_Ey_sin->children[anomtrans::DynDMDerivedBy::omega_inv_DE_down];

  // Add <rho_{0, 2}> for (Ex, cos) and (Ey, sin) cases.
  anomtrans::add_dynamic_electric_n_zero(dm_rho_m1_1_Ex_cos, dm_rho_1_1_Ex_cos, omega, kmb,
      Ex_dot_grad_k.M, Ex_dot_R.M, ksp, H, disorder_term_od, delta,
      berry_broadening, anomtrans::DynVariation::cos);
  anomtrans::add_dynamic_electric_n_zero(dm_rho_m1_1_Ey_sin, dm_rho_1_1_Ey_sin, omega, kmb,
      Ey_dot_grad_k.M, Ey_dot_R.M, ksp, H, disorder_term_od, delta,
      berry_broadening, anomtrans::DynVariation::sin);

  // Valley (anomalous) hall effect at second order in electric field
  // from <rho_{0, 2}>. Intrinsic contribution. Current in y, electric field in x.
  auto dm_n_0_2_Ex_cos = dm_rho_m1_1_Ex_cos->children[anomtrans::DynDMDerivedBy::Kdd_inv_DE_up];
  auto dm_S_0_2_intrinsic_Ex_cos = dm_rho_m1_1_Ex_cos->children[anomtrans::DynDMDerivedBy::P_inv_DE_up];
  auto dm_S_0_2_extrinsic_Ex_cos = dm_rho_m1_1_Ex_cos->children[anomtrans::DynDMDerivedBy::P_inv_Kod_up];

  auto dm_n_0_2_Ey_sin = dm_rho_m1_1_Ey_sin->children[anomtrans::DynDMDerivedBy::Kdd_inv_DE_up];
  auto dm_S_0_2_intrinsic_Ey_sin = dm_rho_m1_1_Ey_sin->children[anomtrans::DynDMDerivedBy::P_inv_DE_up];
  auto dm_S_0_2_extrinsic_Ey_sin = dm_rho_m1_1_Ey_sin->children[anomtrans::DynDMDerivedBy::P_inv_Kod_up];

  Mat n_0_2_plus, S_0_2_intrinsic_plus, S_0_2_extrinsic_plus;
  ierr = MatDuplicate(dm_n_0_2_Ex_cos->rho.M, MAT_COPY_VALUES,
      &n_0_2_plus);CHKERRXX(ierr);
  ierr = MatAXPY(n_0_2_plus, 1.0, dm_n_0_2_Ey_sin->rho.M,
      DIFFERENT_NONZERO_PATTERN);CHKERRXX(ierr);
  ierr = MatDuplicate(dm_S_0_2_intrinsic_Ex_cos->rho.M, MAT_COPY_VALUES,
      &S_0_2_intrinsic_plus);CHKERRXX(ierr);
  ierr = MatAXPY(S_0_2_intrinsic_plus, 1.0, dm_S_0_2_intrinsic_Ey_sin->rho.M,
      DIFFERENT_NONZERO_PATTERN);CHKERRXX(ierr);
  ierr = MatDuplicate(dm_S_0_2_extrinsic_Ex_cos->rho.M, MAT_COPY_VALUES,
      &S_0_2_extrinsic_plus);CHKERRXX(ierr);
  ierr = MatAXPY(S_0_2_extrinsic_plus, 1.0, dm_S_0_2_extrinsic_Ey_sin->rho.M,
      DIFFERENT_NONZERO_PATTERN);CHKERRXX(ierr);

  bool ret_Mat = false;
  auto sigma_n_0_2_plus = anomtrans::calculate_current_ev(kmb, v_op, n_0_2_plus, ret_Mat);
  auto sigma_S_0_2_intrinsic_plus = anomtrans::calculate_current_ev(kmb, v_op,
      S_0_2_intrinsic_plus, ret_Mat);
  auto sigma_S_0_2_extrinsic_plus = anomtrans::calculate_current_ev(kmb, v_op,
      S_0_2_extrinsic_plus, ret_Mat);
  PetscReal sigma_n_0_2_xx_plus = sigma_n_0_2_plus.at(0).first.real();
  PetscReal sigma_n_0_2_xy_plus = sigma_n_0_2_plus.at(1).first.real();
  PetscReal sigma_S_0_2_intrinsic_xx_plus = sigma_S_0_2_intrinsic_plus.at(0).first.real();
  PetscReal sigma_S_0_2_intrinsic_xy_plus = sigma_S_0_2_intrinsic_plus.at(1).first.real();
  PetscReal sigma_S_0_2_extrinsic_xx_plus = sigma_S_0_2_extrinsic_plus.at(0).first.real();
  PetscReal sigma_S_0_2_extrinsic_xy_plus = sigma_S_0_2_extrinsic_plus.at(1).first.real();

  auto spin_n_0_2_plus = anomtrans::calculate_spin_ev(kmb, spin_op, n_0_2_plus, ret_Mat);
  auto spin_S_0_2_intrinsic_plus = anomtrans::calculate_spin_ev(kmb, spin_op,
      S_0_2_intrinsic_plus, ret_Mat);
  auto spin_S_0_2_extrinsic_plus = anomtrans::calculate_spin_ev(kmb, spin_op,
      S_0_2_extrinsic_plus, ret_Mat);
  PetscReal spin_n_0_2_z_plus = spin_n_0_2_plus.at(2).first.real();
  PetscReal spin_S_0_2_intrinsic_z_plus = spin_S_0_2_intrinsic_plus.at(2).first.real();
  PetscReal spin_S_0_2_extrinsic_z_plus = spin_S_0_2_extrinsic_plus.at(2).first.real();

  auto js_n_0_2_plus = anomtrans::calculate_spin_current_ev(kmb, spin_op, v_op, n_0_2_plus,
      ret_Mat);
  auto js_S_0_2_intrinsic_plus = anomtrans::calculate_spin_current_ev(kmb, spin_op, v_op,
      S_0_2_intrinsic_plus, ret_Mat);
  auto js_S_0_2_extrinsic_plus = anomtrans::calculate_spin_current_ev(kmb, spin_op, v_op,
      S_0_2_extrinsic_plus, ret_Mat);
  PetscReal js_n_0_2_sz_vy_plus = js_n_0_2_plus.at(2).at(1).first.real();
  PetscReal js_S_0_2_intrinsic_sz_vy_plus = js_S_0_2_intrinsic_plus.at(2).at(1).first.real();
  PetscReal js_S_0_2_extrinsic_sz_vy_plus = js_S_0_2_extrinsic_plus.at(2).at(1).first.real();

  // Done with PETSc data.
  ierr = MatNullSpaceDestroy(&nullspace);CHKERRXX(ierr);
  ierr = KSPDestroy(&ksp);CHKERRXX(ierr);

  json j_out = {
    {"sigma_n_0_2_xx_plus", sigma_n_0_2_xx_plus},
    {"sigma_n_0_2_xy_plus", sigma_n_0_2_xy_plus},
    {"sigma_S_0_2_intrinsic_xx_plus", sigma_S_0_2_intrinsic_xx_plus},
    {"sigma_S_0_2_intrinsic_xy_plus", sigma_S_0_2_intrinsic_xy_plus},
    {"sigma_S_0_2_extrinsic_xx_plus", sigma_S_0_2_extrinsic_xx_plus},
    {"sigma_S_0_2_extrinsic_xy_plus", sigma_S_0_2_extrinsic_xy_plus},
    {"spin_n_0_2_z_plus", spin_n_0_2_z_plus},
    {"spin_S_0_2_intrinsic_z_plus", spin_S_0_2_intrinsic_z_plus},
    {"spin_S_0_2_extrinsic_z_plus", spin_S_0_2_extrinsic_z_plus},
    {"js_n_0_2_sz_vy_plus", js_n_0_2_sz_vy_plus},
    {"js_S_0_2_intrinsic_sz_vy_plus", js_S_0_2_intrinsic_sz_vy_plus},
    {"js_S_0_2_extrinsic_sz_vy_plus", js_S_0_2_extrinsic_sz_vy_plus}
  };

  std::stringstream outpath;
  outpath << "wannier90_Hamiltonian_WSe2_test_out.json";

  std::ofstream fp_out(outpath.str());
  fp_out << j_out.dump();
  fp_out.close();
}