// @HEADER
// *****************************************************************************
//               Rapid Optimization Library (ROL) Package
//
// Copyright 2014 NTESS and the ROL contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

/*! \file  example_01.cpp
    \brief Shows how to solve the optimal control of Helmholtz problem.
*/

#include "Teuchos_Comm.hpp"
#include "Teuchos_Time.hpp"
#include "ROL_GlobalMPISession.hpp"

#include "Tpetra_Core.hpp"
#include "Tpetra_Version.hpp"

#include <iostream>
#include <algorithm>

#include "ROL_Stream.hpp"
#include "ROL_Solver.hpp"
#include "ROL_UserInputGenerator.hpp"
#include "ROL_ScalarLinearConstraint.hpp"
#include "ROL_Bounds.hpp"
#include "ROL_StdTeuchosBatchManager.hpp"
#include "ROL_Reduced_Objective_SimOpt.hpp"
#include "ROL_PrimalDualRisk.hpp"

#include "../../../../TOOLS/linearpdeconstraintK.hpp"
#include "../../../../TOOLS/pdeobjectiveK.hpp"
#include "../../../../TOOLS/pdevectorK.hpp"
#include "../../../../TOOLS/meshreaderK.hpp"

#include "ROL_OED_Factory.hpp"
#include "ROL_OED_StdMomentOperator.hpp"
#include "../../../utilities/OED_SplitComm.hpp"

#include "../src/pde_helmholtzK.hpp"
#include "../src/obj_helmholtzK.hpp"
#include "../src/noise.hpp"
#include "model.hpp"

template<typename Real>
void solve(const ROL::Ptr<ROL::OED::Factory<Real>>    &factory,
           const ROL::Ptr<ROL::SampleGenerator<Real>> &dsampler,
           const ROL::Ptr<ROL::SampleGenerator<Real>> &osampler,
           ROL::ParameterList                         &list,
           const std::string                          &dtype,
           std::ostream                               &stream     = std::cout,
           const bool                                  useUIG     = true,
           const bool                                  checkDeriv = false) {
  int nsampd = dsampler->numGlobalSamples();
  if (useUIG) {
    factory->setDesign(1.0/static_cast<Real>(nsampd));
  }
  else {
    std::stringstream dfile;
    dfile << dtype << "_optimal_design.txt";
    int err = factory->loadDesign(dfile.str(),2,nsampd);
    stream << "  Factory::loadDesign exited with error " << err << std::endl;
  }
  stream << "  Solve " << dtype << "-optimal design problem." << std::endl;
  std::clock_t timer = std::clock();
  ROL::Ptr<ROL::Problem<Real>> problem = factory->get(list,osampler);
  problem->setProjectionAlgorithm(list);
  problem->finalize(false,true,stream);
  std::string type = list.sublist("OED").get("Optimality Type","I");
  bool usePD = list.sublist("OED").sublist("R-Optimality").get("Use Primal-Dual Algorithm",false);
  if (type=="R" && usePD) {
    ROL::PrimalDualRisk<Real> solver(problem,osampler,list);
    if (checkDeriv) {
      factory->check(stream);
      //solver.check(stream);
      factory->reset();
    }
    solver.run(stream);
  }
  else {
    ROL::Solver<Real> solver(problem,list);
    if (checkDeriv) {
      factory->check(stream);
      problem->check(true,stream);
      factory->reset();
    }
    solver.solve(stream);
  }
  stream << "  " << dtype << "-optimal design time:      "
    << static_cast<Real>(std::clock()-timer)/static_cast<Real>(CLOCKS_PER_SEC)
    << " seconds" << std::endl;
  factory->profile(stream);
  factory->reset();
  std::stringstream dname, pname;
  dname << dtype << "_optimal_design";
  factory->printDesign(dname.str());
  pname << dtype << "_optimal_prediction_variance_samples";
  factory->printPredictionVariance(osampler,pname.str());
}

using RealT = double;
using DeviceT = Kokkos::HostSpace;

int main(int argc, char *argv[]) {
//  feenableexcept(FE_DIVBYZERO | FE_INVALID | FE_OVERFLOW);

  /*** Initialize communicator. ***/
  ROL::GlobalMPISession mpiSession(&argc, &argv);
  Kokkos::ScopeGuard kokkosScope (argc, argv);
  ROL::Ptr<const Teuchos::Comm<int>> dcomm, ocomm;
#ifdef HAVE_MPI
  OED_SplitComm<int,Teuchos::MpiComm<int>> splitcomm(1);
  dcomm = splitcomm.getDesignComm();
  ocomm = splitcomm.getSampleComm();
#else
  dcomm = ROL::makePtr<Teuchos::SerialComm<int>>();
  ocomm = Tpetra::getDefaultComm();
#endif
  // This little trick lets us print to std::cout only if a (dummy) command-line argument is provided.
  const int drank = dcomm->getRank();
  const int orank = ocomm->getRank();
  bool iprint = (argc > 1) && (drank == 0) && (orank == 0);
  iprint = ((argc > 2) ? true : iprint);
  auto outStream = ROL::makeStreamPtr( std::cout, iprint );

  int errorFlag  = 0;

  // *** Example body.
  try {
    RealT tol(1e-8);// one(1);

    /*************************************************************************/
    /******* BUILD LINEAR REGRESSION MODEL BASED ON HELMHOLTZ ****************/
    /*************************************************************************/
    /*** Read in XML input ***/
    auto parlist = ROL::getParametersFromXmlFile("input.xml");
    int   numSpeakers = 8;
    bool   checkDeriv = parlist->sublist("Problem").get("Check Derivatives",  false);
    int     verbosity = parlist->sublist("General").get("Print Verbosity",0);
    bool      alPrint = parlist->sublist("Step").sublist("Augmented Lagrangian").get("Print Intermediate Optimization History",false);
    verbosity         = (iprint ? verbosity : 0);
    alPrint           = (iprint ? alPrint : false);
    parlist->sublist("General").set("Print Verbosity",verbosity);
    parlist->sublist("Step").sublist("Augmented Lagrangian").set("Print Intermediate Optimization History",alPrint);

    // Initialize PDE describing Helmholtz equation.
    auto meshMgr = ROL::makePtr<MeshReader<RealT,DeviceT>>(*parlist);
    auto pde = ROL::makePtr<PDE_Helmholtz_OCT<RealT,DeviceT>>(*parlist);
    auto con = ROL::makePtr<Linear_PDE_Constraint<RealT,DeviceT>>(pde,meshMgr,dcomm,*parlist,*outStream);
    auto assembler = con->getAssembler();
    con->printMeshData(*outStream);

    // Create state vector.
    auto u_ptr = assembler->createStateVector();
    auto r_ptr = assembler->createResidualVector();
    auto up = ROL::makePtr<PDE_PrimalSimVector<RealT,DeviceT>>(u_ptr,pde,assembler,*parlist);
    auto rp = ROL::makePtr<PDE_DualSimVector<RealT,DeviceT>>(r_ptr,pde,assembler,*parlist);
    auto zp = ROL::makePtr<ROL::StdVector<RealT>>(numSpeakers);

    // Create observation function
    auto qoi = ROL::makePtr<QoI_Helmholtz_Observation<RealT,DeviceT>>(pde->getFE(),*parlist);
    auto obs = ROL::makePtr<PDE_Objective<RealT,DeviceT>>(qoi,assembler);

    // Run derivative checks
    if ( checkDeriv ) {
      up->randomize();  zp->randomize();
      auto dup = up->clone(); dup->randomize();
      auto dzp = zp->clone(); dzp->randomize();
      ROL::Vector_SimOpt<RealT> uz(up,zp), duz(dup,dzp);
      std::vector<RealT> param(2,0);
      obs->setParameter(param);
      *outStream << "\n\nCheck Jacobian_1 of PDE Constraint\n";
      con->checkApplyJacobian_1(*up,*zp,*dup,*rp,true,*outStream);
      *outStream << "\n\nCheck Jacobian_2 of PDE Constraint\n";
      con->checkApplyJacobian_2(*up,*zp,*dzp,*rp,true,*outStream);
      *outStream << "\n\nCheck Gradient of Full Objective Function\n";
      obs->checkGradient(uz,duz,true,*outStream);
      *outStream << "\n\nCheck Hessian of Full Objective Function\n";
      obs->checkHessVec(uz,duz,true,*outStream);
    }

    // Compute states for factor decompostion and build regression model
    *outStream << std::endl << "Begin: PDE Solves" << std::endl;
    std::vector<ROL::Ptr<ROL::Vector<RealT>>> state(numSpeakers);
    for (int i = 0; i < numSpeakers; ++i) {
      // Solve real part for speaker i
      std::clock_t time_r = std::clock();
      con->solve(*rp,*up,*zp->basis(i),tol);
      state[i] = up->clone(); state[i]->set(*up);
      std::stringstream real;
      real << "state_" << i << ".txt";
      con->outputTpetraVector(u_ptr,real.str());
      *outStream
        << "  Solve " << i << " time:      "
        << static_cast<RealT>(std::clock()-time_r)/static_cast<RealT>(CLOCKS_PER_SEC)
        << std::endl;
    }
    *outStream << "End: PDE Solves" << std::endl << std::endl;

    // Build nonlinear model and parameter vector
    auto theta = ROL::makePtr<ROL::StdVector<RealT>>(numSpeakers,1);
    auto model = ROL::makePtr<ROL::Reduced_Objective_SimOpt<RealT>>(obs,con,up,theta,up,false);

    /*************************************************************************/
    /******* BUILD EXPERIMENTAL DESIGN PROBLEM AND SOLVE *********************/
    /*************************************************************************/
    // Build samplers for experiment space
    std::string pd = parlist->sublist("Problem").sublist("Design").get("Points File",          "points_5000.txt");
    std::string wd = parlist->sublist("Problem").sublist("Design").get("Weights File",         "weights_5000.txt");
    int nsampd     = parlist->sublist("Problem").sublist("Design").get("Number of Samples",    5000);
    std::string po = parlist->sublist("Problem").sublist("Objective").get("Points File",       "points_200.txt");
    std::string wo = parlist->sublist("Problem").sublist("Objective").get("Weights File",      "weights_200.txt");
    int nsampo     = parlist->sublist("Problem").sublist("Objective").get("Number of Samples", 200);
    *outStream << std::endl;
    *outStream << "  Design Points File:     " << pd << std::endl;
    *outStream << "  Design Weights File:    " << wd << std::endl;
    *outStream << "  Objective Points File:  " << po << std::endl;
    *outStream << "  Objective Weights File: " << wo << std::endl;
    *outStream << std::endl;
    auto dbman  = ROL::makePtr<ROL::StdTeuchosBatchManager<RealT,int>>(dcomm);
    auto obman  = ROL::makePtr<ROL::StdTeuchosBatchManager<RealT,int>>(ocomm);
    auto dsampler  = ROL::makePtr<ROL::UserInputGenerator<RealT>>(pd,wd,nsampd,2,dbman);
    auto osampler  = ROL::makePtr<ROL::UserInputGenerator<RealT>>(po,wo,nsampo,2,obman);

    // Build OED problem factory.
    bool useUIG, homNoise;
    std::string regType, dtype;
    homNoise = parlist->sublist("Problem").get("Homoscedastic Noise",true);
    regType  = parlist->sublist("Problem").get("Regression Type","Least Squares"); 
    auto type     = ROL::OED::StringToRegressionType(regType);
    auto noise    = ROL::makePtr<Helmholtz_Noise<RealT>>();
    auto cov      = ROL::makePtr<ROL::OED::StdMomentOperator<RealT>>(type,homNoise,noise);
    auto factory  = ROL::makePtr<ROL::OED::Factory<RealT>>(model,dsampler,theta,cov,*parlist);
    obman->barrier();

    // Print factors
    bool printFactors = parlist->sublist("Problem").get("Print Factors",false);
    if (printFactors) {
      std::stringstream factName_d, noisName_d, factName_o, noisName_o;
      std::ofstream factFile_d, noisFile_d, factFile_o, noisFile_o;
      ROL::Ptr<ROL::Vector<RealT>> Fp = zp->dual().clone();
      std::vector<RealT> pt;

      factName_d << "factors_des_" << dsampler->batchID() << ".txt";
      noisName_d << "noise_des_" << dsampler->batchID() << ".txt";
      factFile_d.open(factName_d.str());
      noisFile_d.open(noisName_d.str());
      factFile_d << std::scientific << std::setprecision(15);
      noisFile_d << std::scientific << std::setprecision(15);
      for (int i = 0; i < dsampler->numMySamples(); ++i) {
        pt = dsampler->getMyPoint(i);
        factory->getFactors()->evaluate(*Fp,pt);
        for (int j = 0; j < numSpeakers; ++j) {
          factFile_d << std::right << std::setw(25)
                     << (*ROL::dynamicPtrCast<ROL::StdVector<RealT>>(Fp)->getVector())[j];
        }
        factFile_d << std::endl;
        for (int j = 0; j < 2; ++j)
          noisFile_d << std::right << std::setw(25) << pt[j];
        noisFile_d << std::right << std::setw(25) << noise->evaluate(pt) << std::endl;
      }
      factFile_d.close();
      noisFile_d.close();

      factName_o << "factors_opt_" << osampler->batchID() << ".txt";
      noisName_o << "noise_opt_" << osampler->batchID() << ".txt";
      factFile_o.open(factName_o.str());
      noisFile_o.open(noisName_o.str());
      factFile_o << std::scientific << std::setprecision(15);
      noisFile_o << std::scientific << std::setprecision(15);
      for (int i = 0; i < osampler->numMySamples(); ++i) {
        pt = osampler->getMyPoint(i);
        factory->getFactors()->evaluate(*Fp,pt);
        for (int j = 0; j < numSpeakers; ++j) {
          factFile_o << std::right << std::setw(25)
                     << (*ROL::dynamicPtrCast<ROL::StdVector<RealT>>(Fp)->getVector())[j];
        }
        factFile_o << std::endl;
        for (int j = 0; j < 2; ++j)
          noisFile_o << std::right << std::setw(25) << pt[j];
        noisFile_o << std::right << std::setw(25) << noise->evaluate(pt) << std::endl;
      }
      factFile_o.close();
      noisFile_o.close();
    }

    // Solve OED problem
    dtype = parlist->sublist("OED").get("Optimality Type","A");
    useUIG = parlist->sublist("Problem").get("Uniform Initial Guess", false);
    solve<RealT>(factory,dsampler,osampler,*parlist,dtype,*outStream,useUIG,checkDeriv);

    // Get a summary from the time monitor.
    Teuchos::TimeMonitor::summarize();
  }
  catch (std::logic_error &err) {
    std::cout << "Design MPI Rank = " << drank
              << "  Objective MPI Rank = " << orank
              << std::endl << err.what() << std::endl; 
    errorFlag = -1000;
  }; // end try

  if (errorFlag != 0)
    std::cout << "End Result: TEST FAILED\n";
  else
    std::cout << "End Result: TEST PASSED\n";

  return 0;
}
