#include "iwop.hh"
#include "blockop.hh"
#include "adjtest.hh"

#include "MPS_to_RHS.hh"
#include "MPS_Space_Examples.hh"
#include "MPS_conv.hh"

#define CSQ "csq"
#define BULKMOD "bulkmod"
#define BUOYANCY "buoyancy"
#define SOURCE_P "source_p"
#define SOURCE_V0 "source_v0"
#define SOURCE_V1 "source_v1"
#define DATA_P "data_p"
#define DATA_V0 "data_v0"
#define DATA_V1 "data_v1"

#ifndef MY_TOL
#define MY_TOL 1e-3
#endif

#define VERBOSE_MJB

using RVL::valparse;
using RVL::AssignFilename;
using RVL::OperatorEvaluation;

using TSOpt::IWaveLOVOp;
using TSOpt::IWaveSpace;
using TSOpt::CanScal_MPS_Space;
using TSOpt::CanVec_MPS_Space;
using TSOpt::ExVec_MPS_Space;
using TSOpt::Scal_MPS_Space;
using TSOpt::MPS_to_RHS;
using TSOpt::MPS_conv;


template< class T_MPS_Space >
void test_kernel(PARARRAY &pars, FILE &stream, string test_info){

  std::stringstream str;
  string message;
  
  string jobname = valparse<string>(pars,"jobname");

  try {
    str << "\n============================================================\n"
	<< "Test name: "<< jobname <<"\n"
	<< "Runnning test for MPS_conv.\n"
	<< test_info <<"\n";
#ifdef IWAVE_USE_MPI
    str << "rank = "<< retrieveGlobalRank() <<"\n";
#endif
    str << "============================================================\n";

    string driver_type = valparse<string>(pars,"driver");
    string src_type = valparse<string>(pars,"src_type");
    string data_type = valparse<string>(pars,"data_type");

    if( driver_type.compare("acd")!=0 &&
        driver_type.compare("asg")!=0 ){
      RVLException e;
      e << "Driver ("<< driver_type <<") is of unexpected type!\n";
      throw e;
    }
    if( src_type.compare("pressure")!=0 &&
	src_type.compare("velocity")!=0 ){
      RVLException e;
      e << "Source ("<< src_type <<") is of unexpected type!\n";
      throw e;
    }
    if( data_type.compare("pressure")!=0 &&
	data_type.compare("velocity")!=0 ){
      RVLException e;
      e << "Data ("<< data_type <<") is of unexpected type!\n";
      throw e;
    }


    TSOpt::MPS_KEYS mks;
    mks.appx_ord  = "appx_ord";
    mks.MPS_file  = "MPS_file";

    //setting mks.grid_file
    if( driver_type.compare("acd")==0 ){
      mks.grid_file = CSQ;
    }
    else{
      mks.grid_file = BULKMOD;
    }
    //setting mks.RHS_files
    if( src_type.compare("pressure")==0 ){
      mks.RHS_files.push_back(SOURCE_P);
    }
    else{
      mks.RHS_files.push_back(SOURCE_V0);
      mks.RHS_files.push_back(SOURCE_V1);
    }
    //setting mks.hdr_files and mks.G_files
    if( data_type.compare("pressure")==0 ){
      mks.hdr_files.push_back(DATA_P);
      mks.G_files.push_back("data_p_g");
    }
    else{
      mks.hdr_files.push_back(DATA_V0);
      mks.hdr_files.push_back(DATA_V1);
      mks.G_files.push_back("data_v0_g");
      mks.G_files.push_back("data_v1_g");
    }


    /////////////////
    // making data //
    /////////////////
#ifdef VERBOSE_MJB
    str << "\n[] Making data:\n";
#endif

    T_MPS_Space mps_sp(mks,pars);
    MPS_to_RHS m2r(mks,pars,mps_sp);
    IWaveLOVOp iwop(m2r.get_pars(),&stream);
    RVL::LinCompLOVOp<float> IWOP(m2r,iwop);

    //setting inputs
    Vector<ireal> x(IWOP.getDomain());
    Components<ireal> cx(x);
    Components<ireal> cx0(cx[0]);
    Components<ireal> cx1(cx[1]);

    if( driver_type.compare("acd")==0 ){
      AssignFilename af_csq(valparse<string>(pars,CSQ));
      cx0[0].eval(af_csq);
    }
    else{
      AssignFilename af_bmod(valparse<string>(pars,BULKMOD));
      AssignFilename af_buoy(valparse<string>(pars,BUOYANCY));
      cx0[0].eval(af_bmod);
      cx0[1].eval(af_buoy);
    }
    AssignFilename af_src(valparse<string>(pars,mks.MPS_file));
    cx1[0].eval(af_src);

    //setting outputs
    Vector<ireal> y(IWOP.getRange());
    Components<ireal> cy(y);

    if(data_type.compare("pressure")==0){
      AssignFilename af_data_p(valparse<string>(pars,DATA_P));
      cy[0].eval(af_data_p);
    }
    else{
      AssignFilename af_data_v0(valparse<string>(pars,DATA_V0));
      AssignFilename af_data_v1(valparse<string>(pars,DATA_V1));
      cy[0].eval(af_data_v0);
      cy[1].eval(af_data_v1);
    }

    OperatorEvaluation<ireal> opeval(IWOP,x);
    y.copy(opeval.getValue());

#ifdef IWAVE_USE_MPI
    MPI_Barrier(retrieveGlobalComm());
#endif
 
    ///////////////////////
    // making green data //
    ///////////////////////
#ifdef VERBOSE_MJB
    str << "\n[] Making Green data:\n";
#endif

    PARARRAY *pars_2 = ps_new();
    ps_copy(&pars_2,pars);
    
    //adding key-val pairs for Green data simulation
    vector<string> keys(0);
    vector<string> vals(0);

    if(data_type.compare("pressure")==0){
      keys.push_back(DATA_P);
      vals.push_back(valparse<string>(pars,"data_p_g"));
    }
    else{
      keys.push_back(DATA_V0);
      vals.push_back(valparse<string>(pars,"data_v0_g"));
      keys.push_back(DATA_V1);
      vals.push_back(valparse<string>(pars,"data_v1_g"));
    }
    if(src_type.compare("pressure")==0){
      keys.push_back(SOURCE_P);
      vals.push_back(valparse<string>(pars,"source_p_g"));
    }
    else{
      keys.push_back(SOURCE_V0);
      vals.push_back(valparse<string>(pars,"source_v0_g"));
      keys.push_back(SOURCE_V1);
      vals.push_back(valparse<string>(pars,"source_v1_g"));
    }
    keys.push_back("MPS_file");
    vals.push_back(valparse<string>(pars,"MPS_file_g"));

    TSOpt::add_to_pars(*pars_2,keys,vals);

    //Initializing 
    T_MPS_Space mps_sp_g(mks,*pars_2);
    MPS_to_RHS m2r_g(mks,*pars_2,mps_sp_g);
    IWaveLOVOp iwop_g(m2r_g.get_pars(),&stream);
    RVL::LinCompLOVOp<float> IWOP_g(m2r_g,iwop_g);
    
    //setting inputs
    Vector<ireal> x_g(IWOP_g.getDomain());
    Components<ireal> cx_g(x_g);
    Components<ireal> cx0_g(cx_g[0]);
    Components<ireal> cx1_g(cx_g[1]);

    if( driver_type.compare("acd")==0 ){
      AssignFilename af_csq(valparse<string>(pars,CSQ));
      cx0_g[0].eval(af_csq);      
    }
    else{
      AssignFilename af_bmod(valparse<string>(pars,BULKMOD));
      AssignFilename af_buoy(valparse<string>(pars,BUOYANCY));
      cx0_g[0].eval(af_bmod);
      cx0_g[1].eval(af_buoy);
    }
    AssignFilename af_src_g(valparse<string>(*pars_2,"MPS_file_g"));
    cx1_g[0].eval(af_src_g);

    //setting outputs
    Vector<ireal> y_g(IWOP_g.getRange());
    Components<ireal> cy_g(y_g);

    if( data_type.compare("pressure")==0 ){
      AssignFilename af_data_p_g(valparse<string>(*pars_2,"data_p_g"));
      cy_g[0].eval(af_data_p_g);
    }
    else{
      AssignFilename af_data_v0_g(valparse<string>(*pars_2,"data_v0_g"));
      AssignFilename af_data_v1_g(valparse<string>(*pars_2,"data_v1_g"));
      cy_g[0].eval(af_data_v0_g);
      cy_g[1].eval(af_data_v1_g);
    }

    OperatorEvaluation<ireal> opeval_g(IWOP_g,x_g);
    y_g.copy(opeval_g.getValue());

#ifdef IWAVE_USE_MPI
    MPI_Barrier(retrieveGlobalComm());
#endif

    /////////////////////////////////
    // making data via convolution //
    /////////////////////////////////
#ifdef VERBOSE_MJB
    str << "\n[] Making data via convolution:\n\n";
#endif

    //initialing green kernel
    Vector<ireal> k(iwop_g.getIWaveRange());
    Components<ireal> ck(k);
    
    if( data_type.compare("pressure")==0 ){
      AssignFilename af_k_p(valparse<string>(pars,"data_p_g"));
      ck[0].eval(af_k_p);    
    }
    else{
      AssignFilename af_k_v0(valparse<string>(pars,"data_v0_g"));
      AssignFilename af_k_v1(valparse<string>(pars,"data_v1_g"));
      ck[0].eval(af_k_v0);    
      ck[1].eval(af_k_v1);
    }

    //initialize
    MPS_conv conv( pars,
		   y_g,
		   m2r.get_MPS_Domain(),
		   iwop.getIWaveRange() );
    
    //setting inputs
    Vector<ireal> x_re(conv.getDomain());
    AssignFilename af_x_re(valparse<string>(pars,"MPS_file"));
    x_re.eval(af_x_re);

    //setting outputs
    Vector<ireal> y_re(conv.getRange());
    Components<ireal> cy_re(y_re);
   
    if( data_type.compare("pressure")==0 ){
      AssignFilename af_data_p_re(valparse<string>(pars,"data_p_re"));
      cy_re[0].eval(af_data_p_re);
    }
    else{
      AssignFilename af_data_v0_re(valparse<string>(pars,"data_v0_re"));
      AssignFilename af_data_v1_re(valparse<string>(pars,"data_v1_re"));
      cy_re[0].eval(af_data_v0_re);
      cy_re[1].eval(af_data_v1_re);
    }


    //evaluating operator 
    int MPS_idx = valparse<int>(pars,"MPS_idx");
    conv.set_MPS_idx(MPS_idx);
    OperatorEvaluation<ireal> opeval_conv(conv,x_re);
    y_re.copy(opeval_conv.getValue());

#ifdef IWAVE_USE_MPI
    MPI_Barrier(retrieveGlobalComm());
    for( int r=0; r<retrieveGlobalSize(); r++ ){
      if( retrieveGlobalRank()==r ){
	cerr <<"rank="<<r<<", after conv eval\n";
      }
      MPI_Barrier(retrieveGlobalComm());
    } 
#endif

    //////////////////
    // FORWARD TEST //
    //////////////////
    Vector<ireal> fwd_res(y_re);
    fwd_res.copy(y_re);
    fwd_res.linComb(-1,y);
    float norm_fwd_res = fwd_res.norm() / y.norm();

    bool fwd_test = (fabs(norm_fwd_res)<MY_TOL);
    if( fwd_test ) message="\n\033[1;32m[PASSED]\033[0m";
    else message="\n\033[1;31m[FAILED]\033[0m";
    
    message += " Fwd MPS_conv test\n";
    str  << message;
    if(!fwd_test){
      str << "|data - G*f|/|data| = " << norm_fwd_res << "\n"
	  << "|data| = " << y.norm() << "\n"
	  << "|G*f| = " << y_re.norm() << "\n";
    }

#ifdef VERBOSE_MJB
      str << "|data - G*f|/|data| = " << norm_fwd_res << "\n"
	  << "|data| = " << y.norm() << "\n"
	  << "|G*f| = " << y_re.norm() << "\n";
#endif

    
#ifdef IWAVE_USE_MPI
    MPI_Barrier(retrieveGlobalComm());
#endif


    //////////////////
    // Adjoint test //
    //////////////////
#ifdef VERBOSE_MJB
    str << "\n[] Running adjoint test on convolution op:\n\n";
#endif

    std::stringstream res;
    RVL::RVLRandomize<float> rnd(getpid(),-1.0,1.0);
    bool adj_test = RVL::AdjointTest<float>(conv,rnd,res);

    if(adj_test) message="\n\033[1;32m[PASSED]\033[0m";
    else message="\n\033[1;31m[FAILED]\033[0m";
    
    message += " Adjoint test for MPS_conv op\n";
    str  << message;
    if(!adj_test) str << res.str();

#ifdef VERBOSE_MJB
    str << res.str();
#endif

#ifdef IWAVE_USE_MPI
    MPI_Barrier(retrieveGlobalComm());
#endif
    
    /*
    if( driver_type.compare("acd")==0 ){

      ///////////////////////////
      // Adjoint test on LOVOp //
      ///////////////////////////
      str << "\n[] Adjoint test for LOVOp:\n";
    
      Vector<float> x0(IWOP.getProductDomain()[0]);
      AssignFilename af_csq(valparse<string>(pars,CSQ));
      x0.eval(af_csq);

      RVL::LinearRestrictOp<float> lin_IWOP(IWOP,x0);
      std::stringstream res_lin_IWOP;
      RVL::RVLRandomize<float> rnd_lin_IWOP(getpid(),-1.0,1.0);
      adj_test=RVL::AdjointTest<float>(lin_IWOP,rnd_lin_IWOP,res_lin_IWOP);
      str << res_lin_IWOP.str();

      if(adj_test) message="\n[PASSED]";
      else message="\n[FAILED]";

      message += " Adjoint test for LOVOp\n";
      str  << message;

#ifdef IWAVE_USE_MPI
    MPI_Barrier(retrieveGlobalComm());
#endif
      

      /////////////////////////
      // making adjoint data //
      /////////////////////////
      str << "\n[] Making adjoint data:\n";
      
      Vector<float> x1(IWOP.getProductDomain()[1]);
      AssignFilename af_src_adj(valparse<string>(pars,"MPS_adj"));
      x1.eval(af_src_adj);
      
      RVL::AdjLinearOp<float> adj_IWOP(lin_IWOP);
      OperatorEvaluation<float> opeval_adj_IWOP(adj_IWOP,y);
      x1.copy(opeval_adj_IWOP.getValue());

#ifdef IWAVE_USE_MPI
    MPI_Barrier(retrieveGlobalComm());
#endif
      
      /////////////////////////////////////////
      // making adjoint data via convolution //
      /////////////////////////////////////////
      str << "\n[] Making adjoint data via conv:\n\n";
      
      Vector<float> x1_re(IWOP.getProductDomain()[1]);
      AssignFilename af_src_adj_re(valparse<string>(pars,"MPS_adj_re"));    
      x1_re.eval(af_src_adj_re);
      
      RVL::AdjLinearOp<float> adj_conv(conv);
      OperatorEvaluation<float> opeval_adj_conv(adj_conv,y);
      x1_re.copy(opeval_adj_conv.getValue());
      
      //residual
      Vector<ireal> adj_res(x1_re.getSpace());
      adj_res.copy(x1_re);
      adj_res.linComb(-1,x1);
      float norm_adj_res = adj_res.norm() / x1.norm();
      str << "|x1 - adjG(x0)*y|/|x1| = " << norm_adj_res << "\n"
	  << "|x1| = " << x1.norm() << "\n"
	  << "|adjG(x0)*y| = " << x1_re.norm() << "\n";

      if( fabs(norm_adj_res)<MY_TOL ) message="\n[PASSED]";
      else message="\n[FAILED]";

      message += " Adj MPS_conv test\n"; 
      str  << message;
    }
    */

#ifdef IWAVE_USE_MPI
    MPI_Barrier(retrieveGlobalComm());
#endif

    ps_delete(&pars_2);

    // writing out
    string outfile = valparse<string>(pars,"outfile","");

#ifdef IWAVE_USE_MPI
    for( int r=0; r<retrieveGlobalSize(); r++ ){
      if( retrieveGlobalRank()==r ){
	cerr << "rank="<<r<<"\n";
#endif
	if(outfile.size()>0){
	  ofstream outf;
	  outf.open(outfile.c_str(),ios::app);
	  outf << str.str();
	  outf.close();
	}
	else{
	  cout << str.str();
	}
#ifdef IWAVE_USE_MPI
      }
      MPI_Barrier(retrieveGlobalComm());
    }
#endif
  }
  catch(RVLException &e){
    e << "ERROR from test_kernel()!\n"
      << "Dumping output stream:\n"
      << str.str();
    throw e;
  }
}