#include "TSOp.hh"
#include "gtest/gtest.h"
#include "rnmat.hh"
#include "adjtest.hh"
#include "functions.hh"

//#define VERBOSE

#define ATOL 1.e-5

namespace {

  void FwdEuler(RVL::LinearOp<float> const & A,
		RVL::Vector<float> const & x0,
		RVL::Vector<float> & x,
		float dt, float tmin, float tmax) {
    float t = tmin;
    x.copy(x0);
    // // cerr<<"FwdEuler initial state norm = "<<x.norm()<<"\n";
    while (t + 0.1*dt < tmax) {
      // x = (I + dt * A) x
      // cerr<<"FwdEuler t="<<t<<" -> t="<<t+dt<<"\n";
      A.applyOp(dt,x,1.0,x);
      // cerr<<"  new x.norm = "<<x.norm()<<"\n";
      t += dt;
    }
  }

  void BwdEuler(RVL::LinearOp<float> const & A,
		RVL::Vector<float> const & x0,
		RVL::Vector<float> & x,
		float dt, float tmin, float tmax) {
    float t = tmax;
    x.copy(x0);
    // cerr<<"FwdEuler initial state norm = "<<x.norm()<<"\n";
    while (t - 0.1*dt > tmin) {
      // x = (I + dt * A^T) x
      // cerr<<"FwdEuler t="<<t<<" -> t="<<t+dt<<"\n";
      A.applyAdjOp(dt,x,1.0,x);
      // cerr<<"  new x.norm = "<<x.norm()<<"\n";
      t -= dt;
    }
  }

  class Aapply0: public RVL::TSFO {
  private:
    RVL::GenMat<float> & A;
    mutable RVL::RnArray<float> z;
    mutable RVL::RVLLinCombObject<float> lco;
  public:
    Aapply0(RVL::GenMat<float> & _A, float _dt)
      : A(_A), z(A.getNCols()), lco(_dt,1.0)  {
      // cerr<<"Aapply0 dt="<<_dt<<endl;
    }
    void operator()(RVL::TSDC & y) const {
      try {
	// cerr<<"  --> Aapply0 eval factor \n";
	if ((y.getSize() != 2)) {
	  // cerr<<"  error\n";
	  RVL::RVLException e;
	  e<<"Error: Aapply0::operator()\n";
	  e<<"  takes only (control, state) pairs, size=2 \n";
	  e<<"  output vector y:\n";
	  RVL::Writeable const & yw
	    = dynamic_cast<RVL::Writeable const &>(y);
	  yw.write(e);
	  throw e;
	}
	A.getMatVec().setAdj(false);
	RVL::ProductDataContainer & pdcy =
	  dynamic_cast<RVL::ProductDataContainer &>(y[1]);
	RVL::RnArray<float> & rnout =
	  dynamic_cast<RVL::RnArray<float> &>(pdcy[0]);
	// cerr<<" --> MatVec\n";
	A.getMatVec()(z,rnout);
	lco(rnout,z);
	// cerr<<" exit\n";
      }
      catch (bad_cast) {
	// cerr<<"  Aapply0 bad cast\n";
	RVL::RVLException e;
	e<<"Error: Aapply0::operator()\n";
	e<<"  TSDC y does not enclose RnArray \n";
	throw e;
      }
      catch (RVL::RVLException & e) {
	e<<"\ncalled from Aapply0::operator()\n";
	throw e;
      }
    }
    
    std::string getName() const {
      std::string x = "Aapply0\n"; return x;
    }
  };

  class AapplyAdj0: public RVL::TSFO {
  private:
    RVL::GenMat<float> & A;
    mutable RVL::RnArray<float> z;
    mutable RVL::RVLLinCombObject<float> lco;
  public:
    AapplyAdj0(RVL::GenMat<float> & _A, float _dt)
      : A(_A), z(A.getNRows()), lco(_dt,1.0) {}
    
    void operator()(RVL::TSDC & y) const {
      try {
	if ((y.getSize() != 2)) {
	  RVL::RVLException e;
	  e<<"Error: AapplyAdj0::operator()\n";
	  e<<"  takes only (control, state) pairs, size=2 \n";
	  e<<"  output vector y:\n";
	  RVL::Writeable const & yw
	    = dynamic_cast<RVL::Writeable const &>(y);
	  yw.write(e);
	  throw e;
	}
	A.getMatVec().setAdj(true);
	RVL::ProductDataContainer & pdcy =
	  dynamic_cast<RVL::ProductDataContainer &>(y[1]);
	RVL::RnArray<float> & rnout =
	  dynamic_cast<RVL::RnArray<float> &>(pdcy[0]);
	A.getMatVec()(z,rnout);
	lco(rnout,z);
      }
      catch (bad_cast) {
	RVL::RVLException e;
	e<<"Error: AapplyAdj0::operator()\n";
	e<<"  TSDC does not enclose RnArray\n";
	throw e;
      }
      catch (RVL::RVLException & e) {
	e<<"\ncalled from AapplyAdj0::operator()\n";
	throw e;
      }
    }
    
    std::string getName() const {
      std::string x = "AapplyAdj0\n"; return x;
    }
    
  };

  // these two FOs should be set up to take tangent vectors,
  // i.e. pairs of (control,space) pairs
  
  class AapplyFwd1: public RVL::TSFO {
  public:
    void operator()(RVL::TSDC & y) const {}
    std::string getName() const {
      std::string x = "AapplyFwd1\n"; return x;
    }
  };
  
  class AapplyAdj1: public RVL::TSFO {
    void operator()(RVL::TSDC & y) const {}    
    std::string getName() const {
      std::string x = "AapplyAdj1\n"; return x;
    }
  };

  class ArandAccess: public RVL::TSRAFO<float> {
    void operator()(RVL::TSDC & y) const {}    
    std::string getName() const {
      std::string x = "ArandAccess\n"; return x;
    }
  };

  class AStep: public RVL::TSStep<float> {
  private:

    RVL::GenMat<float> & A;
    mutable Aapply0 fwd0;
    mutable AapplyAdj0 adj0;
    mutable AapplyFwd1 fwd1;
    mutable AapplyAdj1 adj1;
    mutable ArandAccess ar0;
    float dt;
    RVL::TSSpace<float> ctrlsp;
    RVL::TSSpace<float> statesp;
    RVL::TSSpace<float> pdom;

    RVL::TSFO & applyFO() const { // cerr<<"  AStep:: return applyFO\n";
      return fwd0; }
    RVL::TSFO & applyAdjFO() const { return adj0; }
    RVL::TSFO & applyTangentFO() const { return fwd1; }
    RVL::TSFO & applyAdjTangentFO() const { return adj1; }
    RVL::TSRAFO<float> & randAccessFO() const { return ar0; }
    
  protected:

    void stepTimeFwd() const { TSClock<float>::stepTime(+dt); }
    void stepTimeBwd() const { TSClock<float>::stepTime(-dt); }
    
    RVL::Operator<float> * clone() const {
      return new AStep(*this);
    }
    
  public:
    AStep(RVL::GenMat<float> & _A,
	  float _t, float _dt)
      : A(_A), fwd0(_A,_dt), adj0(_A,_dt), dt(_dt),
	ctrlsp(1), statesp(1), pdom(2) {
      ctrlsp.set(A.getDomain(),0);
      statesp.set(A.getDomain(),0);
      // cerr<<"in AStep constructor: ctrlsp\n";
      //      ctrlsp.write(cerr);
      pdom.set(ctrlsp,0);
      pdom.set(statesp,1);
      //      cerr<<"in AStep constructor: pdom\n";
      //      pdom.write(cerr);
      this->setTime(_t);
    }
    AStep(AStep const & s)
      : A(s.A), fwd0(A,s.dt), adj0(A,s.dt), dt(s.dt),
		ctrlsp(1), statesp(1), pdom(2) {
      ctrlsp.set(A.getDomain(),0);
      statesp.set(A.getDomain(),0);
      pdom.set(ctrlsp,0);
      pdom.set(statesp,1);
      this->setTime(s.getTime());
    }
    RVL::TSSpace<float> const & getTSDomain() const {
      return pdom;
    }
    RVL::Space<float> const & getDomain() const {
      return getTSDomain();
    }
    RVL::Space<float> const & getRange() const { return pdom; }
    
    ostream & write(ostream & str) const {
      str<<"GenMat linear time step, \n";
      str<<"  t = "<<this->getTime()<<" dt = "<<dt<<"\n";
      str<<"  GenMat data member = \n";
      A.write(str);
      str<<"domain:\n";
      this->getDomain().write(str);
      return str;
    }
  };
  /*
  class ASample: public RVL::TSSample<float> {
  private:

    RVL::GenMat<float> & A;
    RVL::TSSpace<float> statesp;

  protected:

    void apply(RVL::Vector<float> const & x,
	       RVL::Vector<float> & y) const {
      if (start && (abs(this->getTime() - this->getMinTime()) <
		    ATOL*(this->getMaxTime() - this->getMinTime()))) {
	Components<float> cx(x);
	Components<float> cy(y);
	A.applyOp(cx[0],cy[0]);
      }
      if (!start && (abs(this->getTime() - this->getMaxTime()) <
		     ATOL*(this->getMaxTime() - this->getMinTime()))) {
	A.applyOp(cx[0],cy[0]);
      }
    }
    void applyAdj(RVL::Vector<float> const & x,
		  RVL::Vector<float> & y) const {
      if (start && (abs(this->getTime() - this->getMinTime()) <
		    ATOL*(this->getMaxTime() - this->getMinTime()))) {
	// cerr<<"ACopy::apply - start="<<start<<" copy\n";
	A.applyAdjOp(cx[0],cy[0]);
      }
      if (!start && (abs(this->getTime() - this->getMaxTime()) <
		     ATOL*(this->getMaxTime() - this->getMinTime()))) {
	// cerr<<"ACopy::apply - start="<<start<<" copy\n";	
	A.applyAdjOp(cx[0],cy[0]);
      }
    }
    
    RVL::Operator<float> * clone() const {
      return new ASample(*this);
    }
    
  public:
    ASample(RVL::GenMat<float> & _A,
	    float _t)
      : A(_A), statesp(1) {
      statesp.set(A.getDomain(),0);
      this->setTime(_t);
    }
    ASample(ASample const & s)
      : A(s.A), statesp(1) {
      statesp.set(A.getDomain(),0);
      this->setTime(s.getTime());
    }
    RVL::Space<float> const & getDomain() const {
      return statesp;
    }
    RVL::Space<float> const & getRange() const {
      return statesp;
    }
    ostream & write(ostream & str) const {
      str<<"GenMat snapshot sample  \n";
      str<<"  t = "<<this->getTime()<<" dt = "<<dt<<"\n";
      str<<"  GenMat data member = \n";
      A.write(str);
      str<<"domain:\n";
      this->getProductDomain().write(str);
      return str;
    }
  };
  */
  
  class ACopy: public RVL::TSSample<float> {
  private:
    RVL::Space<float> const & dom;
    bool start;
  protected:
    void apply(RVL::Vector<float> const & x,
	       RVL::Vector<float> & y) const {
      if (start && (abs(this->getTime() - this->getMinTime()) <
		    ATOL*(this->getMaxTime() - this->getMinTime()))) {
	// cerr<<"ACopy::apply - start="<<start<<" copy\n";
	y.copy(x);
      }
      if (!start && (abs(this->getTime() - this->getMaxTime()) <
		     ATOL*(this->getMaxTime() - this->getMinTime()))) {
	// cerr<<"ACopy::apply - start="<<start<<" copy\n";
	y.copy(x);
      }
    }
    void applyAdj(RVL::Vector<float> const & x,
		  RVL::Vector<float> & y) const {
      if (start && (abs(this->getTime() - this->getMinTime()) <
		    ATOL*(this->getMaxTime() - this->getMinTime()))) {
	// cerr<<"ACopy::apply - start="<<start<<" copy\n";
	y.copy(x);
      }
      if (!start && (abs(this->getTime() - this->getMaxTime()) <
		     ATOL*(this->getMaxTime() - this->getMinTime()))) {
	// cerr<<"ACopy::apply - start="<<start<<" copy\n";	
	y.copy(x);
      }
    }
    RVL::Operator<float> * clone() const {
      return new ACopy(*this);
    }
  public:
    ACopy(RVL::Space<float> const & _dom, bool _start,
	  float _t=RVL::ScalarFieldTraits<float>::Zero(),
	  float _tmin=RVL::ScalarFieldTraits<float>::Zero(),
	  float _tmax=RVL::ScalarFieldTraits<float>::Zero())
      : dom(_dom), start(_start) {
      this->setTime(_t);
      this->setMinTime(_tmin);
      this->setMaxTime(_tmax);
    }
    ACopy(ACopy const & a): dom(a.dom), start(a.start) {
      this->setTime(a.getTime());
      this->setMinTime(a.getMinTime());
      this->setMaxTime(a.getMaxTime());
    }
    RVL::Space<float> const & getDomain() const { return dom; }
    RVL::Space<float> const & getRange() const { return dom; }
    void setTestTime() {
      if (start) this->setTime(this->getMinTime());
      else this->setTime(this->getMaxTime());
    }
    ostream & write(ostream & str) const {
      str<<"ACopy Time Sample Sampler\n";
      str<<"  time = "<<this->getTime()<<"\n";
      str<<"  min time = "<<this->getMinTime()<<"\n";
      str<<"  max time = "<<this->getMaxTime()<<"\n";
      if (start) str<<"  copy at min time\n";
      else str<<"  copy at max time\n"; 
      str<<"  domain = \n";
      dom.write(str);
    }
  };

  /** assumes multiple components, each of which is a product
      itself */
  template<typename T> 
  void copySubComponent(RVL::Vector<T> const & x,
		     RVL::Vector<T> & v,
			int i, int j) {
    try {
      RVL::Components<T> cv(v);
      RVL::Components<T> ccv(cv[i]);
      ccv[j].copy(x);
    }
    catch (RVL::RVLException & e) {
      e<<"\ncalled from copySubComponent\n";
      throw e;
    }
  }
  
  class TSOpTest : public ::testing::Test {
  protected:
    
    int n;
    float tmax;
    float tmin;
    float dt;
    RVL::RnSpace<float> dom;
    RVL::GenMat<float> A;
    AStep stp;
    RVL::Vector<float> x0;
    RVL::Vector<float> x0b;
    
    TSOpTest()
      :n(10), tmin(0.0), tmax(1.0), dt(0.01),
       dom(n), A(dom,dom), stp(A,tmin,dt),
       x0(dom), x0b(dom) {
      RVL::RVLRandomize<float> rnd(getpid(),-1.0,1.0);
      A.eval(rnd);
      x0.eval(rnd);
      x0b.eval(rnd);
    }
    TSOpTest(TSOpTest const & tt)
      :n(tt.n), tmin(tt.tmin), tmax(tt.tmax), dt(tt.dt),
       A(tt.A), stp(tt.stp), x0(tt.x0), x0b(tt.x0b)    {}
    
  };
  /*
  TEST_F(TSOpTest, explicit_adjoint) {
    RVL::Vector<float> x1(dom); x1.zero();
    RVL::Vector<float> x2(dom); x2.zero();
    FwdEuler(A,x0,x1,dt,tmin,tmax);
    BwdEuler(A,x0b,x2,dt,tmin,tmax);
    float ipdom = (x0.inner(x2))/(x0.norm() * x0b.norm());
    float iprng = (x1.inner(x0b))/(x0.norm() * x0b.norm());
    EXPECT_LT(abs(ipdom-iprng),ATOL);
  }
  */
  TEST_F(TSOpTest, onestep_value_test_matrix_fwd) {
    // cerr<<">>>> onestep_value_test_matrix_fwd\n";
    // cerr<<"\n  state vec construct\n";
    RVL::Vector<float> x1(dom);
    // cerr<<"\n  state vec zero\n";
    x1.zero();
    // cerr<<"\n  fwd euler\n";
    FwdEuler(A,x0,x1,dt,0.0f,dt);

    // cerr<<"\n  step dom vec construct\n";
    RVL::Vector<float> cs(stp.getTSDomain());
    // cerr<<"\n  step dom vec components\n";
    copySubComponent(x0,cs,0,0);
    copySubComponent(x0,cs,1,0);
    // cerr<<"\n  opeval\n";
    RVL::OperatorEvaluation<float> stpeval(stp,cs);
    // cerr<<"\n  evaluate, components of output\n";
    RVL::Components<float> out(stpeval.getValue());
    // cerr<<"\n  components of components\n";
    RVL::Components<float> xout(out[1]);
    // cerr<<"\n  linComb\n";
    xout[0].linComb(-1.0,x1);
    
    EXPECT_LT(xout[0].norm()/x1.norm(),1.e-5);
  }

  TEST_F(TSOpTest, final_value_test_matrix_fwd) {
    RVL::Vector<float> x1(dom); x1.zero();
    RVL::Vector<float> tsx2((stp.getTSDomain())[1]); tsx2.zero();
    FwdEuler(A,x0,x1,dt,tmin,tmax);

    RVL::Vector<float> cs(stp.getTSDomain());
    copySubComponent(x0,cs,0,0);
    copySubComponent(x0,cs,1,0);
    
    ACopy src((stp.getTSDomain())[1],true,tmin,tmin,tmax);
    ACopy data((stp.getTSDomain())[1],false,tmin,tmin,tmax);
    RVL::Components<float> ccs(cs);
    RVL::TimeStepOp<float> op(stp,src,data);
    RVL::LinearRestrictOp<float> lop(op,ccs[0]);
    lop.applyOp(ccs[1],tsx2);
    RVL::Components<float> ctsx2(tsx2);
    ctsx2[0].linComb(-1.0,x1);
    EXPECT_LT(ctsx2[0].norm()/x1.norm(),1.e-5);
  }

  TEST_F(TSOpTest, final_value_test_matrix_bwd) {
    RVL::Vector<float> x1(dom); x1.zero();
    RVL::Vector<float> tsx2((stp.getTSDomain())[1]); tsx2.zero();
    BwdEuler(A,x0,x1,dt,tmin,tmax);

    RVL::Vector<float> cs(stp.getTSDomain());
    copySubComponent(x0,cs,0,0);
    copySubComponent(x0,cs,1,0);
    
    ACopy src((stp.getTSDomain())[1],true,tmin,tmin,tmax);
    ACopy data((stp.getTSDomain())[1],false,tmin,tmin,tmax);
    RVL::Components<float> ccs(cs);
    RVL::TimeStepOp<float> op(stp,src,data);
    RVL::LinearRestrictOp<float> lop(op,ccs[0]);
    lop.applyAdjOp(ccs[1],tsx2);
    RVL::Components<float> ctsx2(tsx2);
    ctsx2[0].linComb(-1.0,x1);
    EXPECT_LT(ctsx2[0].norm()/x1.norm(),1.e-5);
  }
  
  TEST_F(TSOpTest, src_adjoint) {
    ACopy src((stp.getTSDomain())[1],true,tmin,tmin,tmax);
    src.setTestTime();
    RVL::RVLRandomize<float> rnd(getpid(),-1.0,1.0);
    std::stringstream jnk;
    EXPECT_EQ(1,RVL::AdjointTest<float>(src,rnd,jnk));
  }

  TEST_F(TSOpTest, data_adjoint) {
    ACopy data((stp.getTSDomain())[1],true,tmin,tmin,tmax);
    data.setTestTime();
    RVL::RVLRandomize<float> rnd(getpid(),-1.0,1.0);
    std::stringstream jnk;
    EXPECT_EQ(1,RVL::AdjointTest<float>(data,rnd,jnk));
  }

  TEST_F(TSOpTest, step_adjoint_order0) {
    RVL::Vector<float> cs(stp.getTSDomain());
    copySubComponent(x0,cs,0,0);
    RVL::Components<float> ccs(cs);
    RVL::LinRestrictTSStep<float> lop(stp,ccs[0]);
    RVL::RVLRandomize<float> rnd(getpid(),-1.0,1.0);
    std::stringstream jnk;
    bool res = RVL::AdjointTest<float>(lop,rnd,jnk);
    EXPECT_EQ(1,res);
  }

  TEST_F(TSOpTest, timestep_src_adjoint) {
    ACopy src((stp.getTSDomain())[1],true,tmin,tmin,tmax);
    ACopy data((stp.getTSDomain())[1],false,tmin,tmin,tmax);
    RVL::TimeStepOp<float> op(stp,src,data,true);
    EXPECT_EQ(1,op.testSrcAdj());
  }

  TEST_F(TSOpTest, timestep_data_adjoint) {
    ACopy src((stp.getTSDomain())[1],true,tmin,tmin,tmax);
    ACopy data((stp.getTSDomain())[1],false,tmin,tmin,tmax);
    RVL::TimeStepOp<float> op(stp,src,data,true);
    EXPECT_EQ(1,op.testDataAdj());
  }

  TEST_F(TSOpTest, timestep_adjoint_order0) {
    RVL::Vector<float> cs(stp.getTSDomain());
    copySubComponent(x0,cs,0,0);
    RVL::Components<float> ccs(cs);
    ACopy src((stp.getTSDomain())[1],true,tmin,tmin,tmax);
    ACopy data((stp.getTSDomain())[1],false,tmin,tmin,tmax);
    RVL::TimeStepOp<float> op(stp,src,data,true);
    EXPECT_EQ(1,op.testStepAdj0(ccs[0]));
  }

  TEST_F(TSOpTest, timestep_test_all) {
    RVL::Vector<float> cs(stp.getTSDomain());
    copySubComponent(x0,cs,0,0);
    RVL::Components<float> ccs(cs);
    ACopy src((stp.getTSDomain())[1],true,tmin,tmin,tmax);
    ACopy data((stp.getTSDomain())[1],false,tmin,tmin,tmax);
    RVL::TimeStepOp<float> op(stp,src,data,true);
    EXPECT_EQ(1,op.testAll(ccs[0]));
  }

  TEST_F(TSOpTest, timestep_test_adjoint) {
    RVL::Vector<float> cs(stp.getTSDomain());
    copySubComponent(x0,cs,0,0);
    RVL::Components<float> ccs(cs);
    ACopy src((stp.getTSDomain())[1],true,tmin,tmin,tmax);
    ACopy data((stp.getTSDomain())[1],false,tmin,tmin,tmax);
    RVL::TimeStepOp<float> op(stp,src,data,true);
    RVL::LinearRestrictOp<float> lop(op,ccs[0]);
    RVL::RVLRandomize<float> rnd(getpid(),-1.0f,1.0f);
    std::stringstream jnk;
    EXPECT_EQ(1,RVL::AdjointTest<float>(lop,rnd,jnk));
  }

} // namespace