#include "iwopt.hh"

#define VERBOSE

namespace TSOpt {

  void IWaveOpt(PARARRAY pars,
		RVL::Operator<float> const & op,
		RVL::Vector<float> const & d,
		RVL::Vector<float> & m) {
    try {
      /*
      Components<float> cm(m);

      for (int i=0;i<cm.getSize();i++) {
#ifdef IWAVE_USE_MPI
	RVL::RVLMin<float> svmin;
	RVL::RVLMax<float> svmax;
	RVL::MPISerialFunctionObjectRedn<float,float> vmin(svmin);
	RVL::MPISerialFunctionObjectRedn<float,float> vmax(svmax);
	cm[i].eval(vmin);
	cm[i].eval(vmax);
	cerr<<"StdOLS bfore \n";
	cerr<<"  background model comp "<<i<<" min="<<vmin.getValue()<<" max="<<vmax.getValue()<<"\n";
	int nnn;
	cin>>nnn;
#else
	RVL::RVLMin<float> vmin;
	RVL::RVLMax<float> vmax;
	cm[i].eval(vmin);
	cm[i].eval(vmax);
	cerr<<"  background model comp "<<i<<" min="<<vmin.getValue()<<" max="<<vmax.getValue()<<"\n";
#endif
      }
      */
      /* output stream */
      std::stringstream outstr;
      outstr<<scientific;
      std::ostream * optr = &outstr;
      std::string outfile = RVL::valparse<std::string>(pars,"outfile","");
      if ((outfile.size()==0) && (retrieveRank()==0)) optr = &cout;
	
      /* switch on method */
      RVLAlg::Algorithm * alg = NULL;
      std::string optmethod = RVL::valparse<std::string>(pars,"OptMethod","lbfgs");
#ifdef VERBOSE
      cerr<<"IWaveOpt: optmethod="<<optmethod<<"\n";
#endif
      if (optmethod == "lbfgs") {
	RVL::StdLeastSquaresFcnlGN<float> f(op,d);
	alg = new RVLUmin::LBFGSBT<float>
	  (f, m,
	   RVL::valparse<float>(pars,"InvHessianScale",1.0f),
	   RVL::valparse<int>(pars,"MaxInvHessianUpdates",5),
	   RVL::valparse<int>(pars,"MaxSubSteps",10),
	   RVL::valparse<bool>(pars,"VerboseDisplay",true), 
	   RVL::valparse<float>(pars,"InitStepBound",1.0f), 
	   RVL::valparse<float>(pars,"MinDecrease",0.1f),
	   RVL::valparse<float>(pars,"GoodDecrease",0.9f), 
	   RVL::valparse<float>(pars,"StepDecrFactor",0.5f), 
	   RVL::valparse<float>(pars,"StepIncrFactor",1.8f),
	   RVL::valparse<float>(pars,"MaxFracDistToBdry",1.0), 
	   RVL::valparse<float>(pars,"MinStepTol",1.e-06),     // min step as frac of LS segment
	   RVL::valparse<int>(pars,"MaxSteps",10), 
	   RVL::valparse<float>(pars,"AbsGradThresh",0.0), 
	   RVL::valparse<float>(pars,"RelGradThresh",1.e-2), 
	   *optr);
      }
      else if (optmethod == "trcg") {
	RVL::ResidualOperator<float> rop(op,d);      
	RVLUmin::TRGNAlg<float, RVLUmin::CGNEPolicy<float> > * tralg = new RVLUmin::TRGNAlg<float, RVLUmin::CGNEPolicy<float> >
	  (rop, m,
	   RVL::valparse<int>(pars,"MaxSteps",10),             // _maxcount,
	   RVL::valparse<float>(pars,"ResidualTol",0.0f),       // _jtol,
	   RVL::valparse<float>(pars,"AbsGradThresh",0.0f),    // _agtol,
	   RVL::valparse<float>(pars,"RelGradThresh",1.0e-2),  // _rgtol,
	   RVL::valparse<float>(pars,"MinDecrease",0.1f),      // _eta1
	   RVL::valparse<float>(pars,"GoodDecrease",0.9f),     // _eta2
	   RVL::valparse<float>(pars,"StepDecrFactor",0.5f),   // _gamma1
	   RVL::valparse<float>(pars,"StepIncrFactor",1.8f),   // _gamma2
	   RVL::valparse<float>(pars,"MinStepTol",1.e-06),     // min step as frac
	   *optr);
	
	// assign CG params
	tralg->assign
	  (RVL::valparse<float>(pars,"CGNE_ResTol",0.0f),      // rtol
	   RVL::valparse<float>(pars,"CGNE_GradTol",0.001f),   // nrtol,
	   RVL::valparse<float>(pars,"InitStepBound",1.0f),    // Delta
	   RVL::valparse<int>(pars,"MaxSubSteps",10),          // maxcount
	   RVL::valparse<bool>(pars,"VerboseDisplay",true));   // verbose
	alg=tralg;
      }
      else {
	RVL::RVLException e;
	e<<"ERROR: IWaveOpt\n";
	e<<"  no optimization algorithm selected, or algoritm specified is not\n";
	e<<"  supported. Currently supported algorithms:\n";
	e<<"  OptMethod = lbfgs (Limited Memory Broyden-Fletcher-Goldfarb-Shanno)\n";
	e<<"  OptMethod = trcg (Trust Region Gauss-Newton with CG linear solver\n)";
	throw e;
      }

#ifdef VERBOSE
      cerr<<"IWaveOpt -> run\n";
#endif
      alg->run();
#ifdef VERBOSE
      cerr<<"IWaveOpt -> report\n";
#endif
      delete alg;
      /*
      for (int i=0; i<cm.getSize();i++) {
#ifdef IWAVE_USE_MPI
	RVL::RVLMin<float> svmin;
	RVL::RVLMax<float> svmax;
	RVL::MPISerialFunctionObjectRedn<float,float> vmin(svmin);
	RVL::MPISerialFunctionObjectRedn<float,float> vmax(svmax);
	cm[i].eval(vmin);
	cm[i].eval(vmax);
	cerr<<"StdOLS after\n";
	cerr<<"  background model comp "<<i<<" min="<<vmin.getValue()<<" max="<<vmax.getValue()<<"\n";
	int nnn;
	cin>>nnn;
#else
	RVL::RVLMin<float> vmin;
	RVL::RVLMax<float> vmax;
	cm[i].eval(vmin);
	cm[i].eval(vmax);
	cerr<<"  background model comp "<<i<<" min="<<vmin.getValue()<<" max="<<vmax.getValue()<<"\n";
#endif	
      }
      */
    
      if ((outfile.size()>0) && (retrieveRank()==0)) {
	std::ofstream outf(outfile.c_str());
        outf<<outstr.str();
        outf.close();
      }

      return;

    }
    catch (RVL::RVLException & e) {
      e<<"\ncalled from IWaveOpt\n";
      throw e;

    }
  }

  void IWaveResMod(PARARRAY pars,
		   IWaveFWIOp const & op,
		   RVL::Vector<float> const d,
		   RVL::Vector<float> const r) {

    try {

      // create model vector, optionally assign components to
      // archival files
      RVL::Vector<float> m(op.getTransformOp().getRange());
      Components<float> cm(m);
      std::vector<std::string> mkeys = op.getDomainKeys();

      size_t flag = 0;
      for (int i=0; i<mkeys.size(); i++) {
	// chop off "_est"
	std::string mdl = mkeys[i].substr(0,mkeys[i].size()-4);
	// tack on "_inv", check to see if file is specified
	std::string ename = RVL::valparse<std::string>(pars,mdl+"_est","");
	std::string fname = RVL::valparse<std::string>(pars,mdl+"_inv","");
	//	cerr<<"IWaveResMod:\n";
	//	cerr<<"  model index="<<i<<" key="<<mkeys[i]<<" root="<<mdl<<" inv="<<fname<<"\n";
	// if so, connect to vector
	if ((ename.size() > 0) && (fname.size() > 0)) {
	  RVL::AssignFilename fn(fname);
	  cm[i].eval(fn);
	  flag += fname.size();
	  /*
#ifdef IWAVE_USE_MPI
	  RVL::RVLMin<float> svmin;
	  RVL::RVLMax<float> svmax;
	  RVL::MPISerialFunctionObjectRedn<float,float> vmin(svmin);
	  RVL::MPISerialFunctionObjectRedn<float,float> vmax(svmax);
	  cm[i].eval(vmin);
	  cm[i].eval(vmax);
	  cerr<<"  background model comp "<<i<<" min="<<vmin.getValue()<<" max="<<vmax.getValue()<<"\n";
	  int nnn;
	  cin>>nnn;
#else
	  RVL::RVLMin<float> vmin;
	  RVL::RVLMax<float> vmax;
	  cm[i].eval(vmin);
	  cm[i].eval(vmax);
	  cerr<<"  background model comp "<<i<<" min="<<vmin.getValue()<<" max="<<vmax.getValue()<<"\n";
#endif	
	  */
	}
      }
      // don't bother if the user hasn't asked for any model files
      if (flag > 0) {
	RVL::OperatorEvaluation<float> transeval(op.getTransformOp(),r);
	m.copy(transeval.getValue());
      }

      // create residual vector, optionally assign components to
      // archival files
      RVL::Vector<float> res(op.getModelingOp().getRange());
      Components<float> cres(res);
      std::vector<std::string> ekeys = op.getRangeKeys();

      size_t elag = 0;
      for (int i=0; i<ekeys.size(); i++) {
	std::string fname = RVL::valparse<std::string>(pars,ekeys[i]+"_res","");
	if (fname.size() > 0) {
	  RVL::AssignFilename fn(fname);
	  cres[i].eval(fn);
	  elag += fname.size();
	}
      }
      // don't bother if the user hasn't asked for any residual files
      if (elag > 0) {
	if (flag > 0) {
	  RVL::OperatorEvaluation<float> modeleval(op.getModelingOp(),m);
	  res.copy(modeleval.getValue());
	}
	else {
	  RVL::OperatorEvaluation<float> modeleval(op,r);
	  res.copy(modeleval.getValue());
	}
      }
      res.linComb(-1.0f,d);
    }
    catch (RVL::RVLException & e) {
      e<<"\ncalled from IWaveResMod\n";
      throw e;
    }
  }

  int IWaveOpApply(int argc, char ** argv,
		   void (*invfcn)(PARARRAY, FILE *)) {

    try {

      PARARRAY * pars = NULL;
      FILE * stream = NULL;
      TSOpt::IWaveEnvironment(argc, argv, 0, &pars, &stream);

      if (retrieveGlobalRank()==0 && argc<2) {
	pagedoc();
	exit(0);
      }

#ifdef IWAVE_USE_MPI
      if (retrieveGroupID() == MPI_UNDEFINED) {
	fprintf(stream,"NOTE: finalize MPI, cleanup, exit\n");
	fflush(stream);
      }
      else {
#endif

	invfcn(*pars,stream);
	
#ifdef IWAVE_USE_MPI
	/* end nontriv comm branch */
      }
      MPI_Barrier(MPI_COMM_WORLD);
#endif
      ps_delete(&pars);
      fclose(stream);
      
      return 0;
    }
    catch (RVL::RVLException & e) {
      e<<"\ncalled from IWaveOpApply\n";
      throw e;
    }
  }

  void StraightOLS(PARARRAY pars, FILE * stream) {
    try {
#ifdef VERBOSE
      cerr<<"StraightOLS -> IWaveFWIOp constructor\n";
#endif
      TSOpt::IWaveFWIOp op(pars, stream);
#ifdef VERBOSE
      cerr<<"StraightOLS -> domain vector\n";
#endif
      RVL::Vector<float> r(op.getDomain());
#ifdef VERBOSE
      cerr<<"StraightOLS -> range vector\n";
#endif
      RVL::Vector<ireal> d(op.getRange());
#ifdef VERBOSE
      for (int i=0;i<op.getDomainKeys().size(); i++)
	cerr<<"StraightOLS: domain key "<<i<<" = "<<op.getDomainKeys()[i]<<"\n";
      cerr<<"StraightOLS->load domain vector\n";
#endif
      TSOpt::IWaveLoad(pars,r,op.getDomainKeys());
#ifdef VERBOSE
      for (int i=0;i<op.getRangeKeys().size(); i++)
	cerr<<"StraightOLS: range key "<<i<<" = "<<op.getRangeKeys()[i]<<"\n";
      cerr<<"StraightOLS->load range vector\n";
#endif      
      TSOpt::IWaveLoad(pars,d,op.getRangeKeys());
#ifdef VERBOSE
      cerr<<"StraightOLS: optimize\n";
#endif
      TSOpt::IWaveOpt(pars, op, d, r);
#ifdef VERBOSE
      cerr<<"StraightOLS: cleanup\n";
#endif      
      TSOpt::IWaveResMod(pars,op,d,r);
#ifdef VERBOSE
      cerr<<"StraightOLS exit\n";
#endif
    }
    catch (RVL::RVLException & e) {
      e<<"\ncalled from TSOpt::StraightOLS\n";
      throw e;
    }
  }
}