/*************************************************************************

Copyright Rice University, 2004-2015.
All rights reserved.

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, provided that the above copyright notice(s) and this
permission notice appear in all copies of the Software and that both the
above copyright notice(s) and this permission notice appear in supporting
documentation.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT OF THIRD PARTY
RIGHTS. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR HOLDERS INCLUDED IN THIS
NOTICE BE LIABLE FOR ANY CLAIM, OR ANY SPECIAL INDIRECT OR CONSEQUENTIAL
DAMAGES, OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR
PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF
THIS SOFTWARE.

Except as contained in this notice, the name of a copyright holder shall
not be used in advertising or otherwise to promote the sale, use or other
dealings in this Software without prior written authorization of the
copyright holder.

**************************************************************************/
#ifndef __RVL_TS
#define __RVL_TS

#include "space.hh"
#include "op.hh"
#include "productspace.hh"
#include "functions.hh"
#include "write.hh"

#define TOL 0.1

namespace RVL {

  template<typename T>
  class TSClock {
  private:
    mutable T t;
  public:
    TSClock(T _t = ScalarFieldTraits<T>::Zero()): t(_t) {}
    T const & getTime() const { return t; }
    void setTime(T _t) const { t = _t; }
    void stepTime(T dt) const { t += dt; }
  };

  template<typename T>
  class TSInterval {
  private:
    mutable T tmin;
    mutable T tmax;
  public:
    TSInterval(T _tmin = ScalarFieldTraits<T>::Zero(),
	       T _tmax = ScalarFieldTraits<T>::Zero())
      : tmin(_tmin), tmax(_tmax) {}
    TSInterval(TSInterval<T> const & x)
      : tmin(x.tmin), tmax(x.tmax) {}
    T const & getMinTime() const { return tmin;}
    T const & getMaxTime() const { return tmax;}
    void setMinTime(T _tmin) const { tmin=_tmin;}
    void setMaxTime(T _tmax) const { tmax=_tmax;}
  };

  template<typename T>
  class TSSample:
    public LinearOp<T>,
    public TSInterval<T>,
    public TSClock<T> {
  public:
    virtual void setTestTime() = 0;
  };

  // forward declaration
  class TSDC;
  
  class TSFO: public FunctionObject {
  public:
 
    // unary operator interface - overwrite mode 
    virtual void operator()(TSDC & y) const = 0;

  };

  template<typename T>
  class TSRAFO: public TSFO,
		public TSClock<T> {};

  // interface to intercept evaluation of tsfo's.
  // handles two cases:
  // 1) arguments are TSDCs (control,state);
  // 2) arguments are two pairs (TSDC, TSDC) ((control,state),(control,state))
  class TSDC: public StdProductDataContainer {
  public:
    void eval( FunctionObject & f,
	       std::vector<DataContainer const *> & x) {
      try {
	// cerr<<"TSDC::eval\n";
	// cerr<<"  FO = "<<f.getName()<<endl;
	TSFO * g = NULL;
	if ((x.size()==0) && (g = dynamic_cast<TSFO *>(&f))) {
	  (*g)(*this);
	}
	else {
	  // cerr<<"  -->StdProductDataContainer::eval\n";
	  StdProductDataContainer::eval(f,x);
	}
      }
      catch (RVLException & e) {
	e<<"\ncalled from TSDC::eval(FO,...)\n";
	throw e;
      }
    }

    ostream & write(ostream & str) const {
      str<<"TSDC: specialization of \n";
      StdProductDataContainer::write(str);
      return str;
    }
  };
  
  template<typename T>
  class TSSpace: public StdProductSpace<T> {
  protected:
    Space<T> * clone() const { return new TSSpace<T>(*this); }
  public:
    TSSpace(size_t nfac): StdProductSpace<T>(nfac) {}
    TSSpace(TSSpace<T> const & sp): StdProductSpace<T>(sp) {}
    
    /** overrides method of StdProductDataContainer class.  */
    DataContainer * buildDataContainer() const {
      if (StdProductSpace<T>::init()) {
	TSDC * d = 
	  new TSDC();
	for (size_t i=0; i<StdProductSpace<T>::getSize(); i++) {
	  SpaceDCF<T> f((*this)[i]);
	  d->push(f);
	}
	// cerr<<"*** from TSSpace::buildDataContainer\n";
	//	d->write(cerr);
	return d;
      }
      else {
	RVLException e;
	e<<"ERROR: TSSpace::buildDataContainer\n";
	e<<"  called on uninitialized StdProductSpace object\n";
	throw e;
      }
    }
    
    ostream & write(ostream & str) const {
      str<<"TSSpace: product space with TSDC components\n";
      ProductSpace<T>::write(str);
      return str;
    }
  };

  // forward declarations of friendship
  template<typename T>
  class TimeStepOp;

  template<typename T>
  class LinRestrictTSStep;
  
  template<typename T>
  class TSStep:
    public Operator<T>,
    public TSClock<T> {
    
    friend class TimeStepOp<T>;
    friend class LinRestrictTSStep<T>;
    
  private:
    virtual TSFO & applyFO() const = 0;
    virtual TSFO & applyAdjFO() const = 0;
    virtual TSFO & applyTangentFO() const = 0;
    virtual TSFO & applyAdjTangentFO() const = 0;
    virtual TSRAFO<T> & randAccessFO() const = 0;
    
  protected:

    // here x and y are both supposed to be control-state vectors
    void apply(Vector<T> const & x,
	       Vector<T> & y) const {
      try {
	//	cerr<<"  AStep::apply -> copy\n";
	if (&y != &x) y.copy(x);
	// cerr<<"  AStep::apply -> eval(applyFO()) FO = "<<applyFO().getName()<<"\n";
	//	cerr<<"  input vector:\n";
	//	y.write(cerr);
	y.eval(applyFO());
	// cerr<<"  AStep::apply -> stepTimeFwd()\n";	
	this->stepTimeFwd();       
      }
      catch (RVLException & e) {
	e<<"\ncalled from TSStep::apply\n";
	throw e;
      }
    }

    /** Adjoint of apply in state variable. Precond: x, y are members of
	TSSpace with two components; applyFO and applyAdjFO copy first 
	components (control) and do something linear to second (state)
    */
    void applyAdj(Vector<T> const & x,
		  Vector<T> & y) const {
      try {
	stepTimeBwd();
	if (&y != &x) y.copy(x);
	y.eval(applyAdjFO());
      }
      catch (RVLException & e) {
	e<<"\ncalled from TSStep::apply\n";
      }
    }

    /** tangent map: (c,s,dc,ds) -> (c,F[c,s],dc,DF[c,s][dc,ds]) 
	combine args into product space: (x,dx) = ((c,s),(dc,ds))
	ydy = (y,dy) = (s,ds)

	only TimeStepOp accesses the Tangent methods, so sanity checking
	deferred to there.
    */
    void applyTangent(Vector<T> const & xdx,
		      Vector<T> & ydy) const {
      try {
	if (&ydy != &xdx) ydy.copy(xdx);
	ydy.eval(applyTangentFO());
	stepTimeFwd();
      }
      catch (RVLException & e) {
	e<<"\ncalled from TSStep::apply\n";
      }
    }

    /** adjoint tangent map: (c,s,dc,ds) -> (c,H[c]^{-1}s,dc + ((DH[c]
	*)H[c]^{-1}s)^T dc, H[c]^{T} ds) cannot implement H[c]^{-1} in general,
	since H[c] need not be invertible. Instead, assert that (c,s)
	is value of base trajectory at t+dt, and interpret
	(c,H[c]^{-1}s) as state at time t. Therefore this computation
        requires two stages: (1) time reverse, t -> t-dt, and extraction
        of state - method for random access to tracjectory, TSFO/TSClock;
	and stationary adjoint tangent map, implementing 
	(c,s,dc,ds) -> (c,s,dc + ((DH[c]*)s)^T dc, H[c]^{T} ds)
    */
    void applyAdjTangent(Vector<T> const & xdx,
		      Vector<T> & ydy) const {
      try {
	if (&ydy != &xdx) ydy.copy(xdx);
	stepTimeBwd();
	randAccessFO().setTime(this->getTime());
	Components<T> cydy(ydy);
	cydy[0].eval(randAccessFO());
	ydy.eval(applyAdjTangentFO());
      }
      catch (RVLException & e) {
	e<<"\ncalled from TSStep::apply\n";
      }
    }
    
    /** implement via tangent map 
	x   = (c,s)^T
	dx  = (dc,ds)^T
	dy  = (dc,DF[c,s][dc,ds]) = (dc, (DH[c]dc)s + H[c]ds)
	= M[c,s][dc ds]^T

	M[c,s] = [ I            0   ]
                 [ (DH[c] *)s  H[c] ]

        this method and applyAdjDeriv are useful mostly for test purposes

    */
    void applyDeriv(Vector<T> const & x,
		    Vector<T> const & dx,
		    Vector<T> & dy) const {
      try {
	StdProductSpace<T> TX(this->getDomain(), this->getDomain());
	Vector<T> xdx(TX);
	Vector<T> ydy(TX);
	Components<T> cxdx(xdx);
	Components<T> cydy(ydy);
	cxdx[0].copy(x);
	cxdx[1].copy(dx);
	applyTangent(xdx,ydy);
	dy.copy(cydy[1]);
	//	stepTimeFwd(); already by applyTangent
      }
      catch (RVLException & e) {
	e<<"\ncalled from TSStep::applyDeriv\n";
      }
    }
    
    /** 
	x = (c,s)^T
	dx  = (dc,ds)^T
	dy  = M[c,s]^Tdx
	M[c,s]^T = [ I           ((DH[c] *)s)^T ]
                   [ 0                   H[c]^T ]
    */
    void applyAdjDeriv(Vector<T> const & x,
		       Vector<T> const & dx,
		       Vector<T> & dy) const {
      try {
	StdProductSpace<T> TX(this->getDomain(), this->getDomain());
	Vector<T> xdx(TX);
	Vector<T> ydy(TX);
	Components<T> cxdx(xdx);
	Components<T> cydy(ydy);
	cxdx[0].copy(x);
	cxdx[1].copy(dx);
	applyAdjTangent(xdx,ydy);
	dy.copy(cydy[1]);
	// bwd time step applied by applyAdjTangent
      }
      catch (RVLException & e) {
	e<<"\ncalled from TSStep::applyAdjDeriv\n";
      }
    }

    // should be implemented in terms of TSClock::stepTime
    virtual void stepTimeFwd() const = 0;
    virtual void stepTimeBwd() const = 0;

  public:

    TSSpace<T> const & getTSDomain() const {
      try {
	TSSpace<T> const & pdom =
	  dynamic_cast<TSSpace<T> const &>(this->getDomain());
	return pdom;
      }
      catch (bad_cast) {
	RVLException e;
	e<<"Error: TSStep::getTSDomain\n";
	e<<"  it aint a PD\n";
	this->getDomain().write(e);
	throw e;
      }
    }
    /*    ProductSpace<T> const & getProductDomain() const {
      return this->getTSDomain();
    }
    */
  };

  template<typename T>
  class LinRestrictTSStep: public LinearOp<T> {
  private:
    TSStep<T> const & step;
    mutable Vector<T> xbuf; mutable Components<T> cxbuf;
    mutable Vector<T> ybuf; mutable Components<T> cybuf;    
    
  protected:

    void apply(const Vector<T> & x,
	       Vector<T> & y) const {
      try {
	// input x0 is already stored in comp[0]
	cxbuf[1].copy(x);
	step.apply(xbuf,ybuf);
	y.copy(cybuf[1]);
      }
      catch (RVLException & e) {
	e<<"\ncalled from LinRestrictTSStep::apply\n";
	throw e;
      }
    }
    
    virtual void applyAdj(const Vector<T> & x,
			  Vector<T> & y) const {
      try {
	// input x0 is already stored in comp[0]
	cxbuf[1].copy(x);
	step.applyAdj(xbuf,ybuf);
	y.copy(cybuf[1]);
      }
      catch (RVLException & e) {
	e<<"\ncalled from LinRestrictTSStep::applyAdj\n";
	throw e;
      }
    }

    Operator<T> * clone() const {
      return new LinRestrictTSStep<T>(*this);
    }

  public:

    LinRestrictTSStep(TSStep<T> const & _step,
		      Vector<T> const & x0)
      : step(_step),
	xbuf(step.getTSDomain()),
	cxbuf(xbuf),
	ybuf(step.getTSDomain()),
      	cybuf(ybuf) {
      try {
	// this catches any mismatch between x0 and
	// dom[0] - other mismatches caught by LinearOp interface
	cxbuf[0].copy(x0);
      }
      catch (RVLException & e) {
	e<<"\ncalled from LinRestrictTSStep constructor\n";
	e<<"input x0:\n";
	x0.write(e);
	e<<"output comp[0]:\n";
	cxbuf[0].write(e);
	throw e;
      }
    }

    LinRestrictTSStep(LinRestrictTSStep<T> const & a)
      : step(a.step),
	xbuf(step.getProductDomain()),
	cxbuf(xbuf),
	ybuf(step.getProductDomain()),
	cybuf(ybuf) {	
      try {
	cxbuf[0].copy(a.cxbuf[0]);
      }
      catch (RVLException & e) {
	e<<"\ncalled from LinRestrictTSStep constructor\n";
	throw e;
      }
    }

    Space<T> const & getDomain() const { return step.getProductDomain()[1]; }
    Space<T> const & getRange() const { return step.getProductDomain()[1]; }

    ostream & write(ostream & str) const {
      str<<"LinRestrictTSStep linear op: restriction to comp 1\n";
      str<<"  of TSStep operator\n";
      step.write(str);
      str<<"  with comp 0 restricted to\n";
      cxbuf[0].write(str);
      return str;
    }
  };
  
  template<typename T>
  class TimeStepOp: public LinOpValOp<T> {
  private:
    TSStep<T> & step;
    TSSample<T> & src;
    TSSample<T> & data;
    StdProductSpace<T> pdom;   // product domain of LOVOp
    StdProductSpace<T> tdom;   // ctrl-state domain of step tangent map
    // for tests
    mutable RVLRandomize<T> rnd;
    mutable ofstream report;
    bool testflag;

  protected:

    /** x0 = control, x1 = src, y = output */
    void apply0(Vector<T> const & x0,
		Vector<T> const & x1,
		Vector<T> & y) const {
      try {
	Vector<T> ctrlstate(step.getProductDomain());
	Components<T> comp_ctrlstate(ctrlstate);
	// implicit sanity test
	comp_ctrlstate[0].copy(x0);
	comp_ctrlstate[1].zero();
	y.zero();
	T t = min(src.getMinTime(), data.getMinTime());
	step.setTime(t);
	T tend = data.getMaxTime();
	T tlast = t;
	//	cerr<<"******* fwd: before call to LinRestrictOp constr\n";
	//	cerr<<"******* x0:\n";
	//	x0.write(cerr);
	//	LinRestrictOp<T> linstep(step,x0);
	while (t + TOL*(t-tlast) < tend) {
	  src.setTime(t);
	  src.applyOp(x1,comp_ctrlstate[1]);
	  //	  cerr<<"TimeStepOp::apply -> linstep.apply state.norm="<<state.norm()<<"\n";
	  //	  cerr<<"TimeStepOp::apply -> linstep.apply time = "<< t <<"\n";
	  //	  linstep.applyOp(state,state);
	  step.apply(ctrlstate,ctrlstate);
	  // cerr<<"TimeStepOp::apply <- linstep.apply state.norm="<<state.norm()<<"\n";
	  // cerr<<"TimeStepOp t="<<t<<" -> t=";
	  tlast = t;
	  t = step.getTime();
	  // cerr<<t<<"\n";
	  data.setTime(t);
	  //	data.applyPlusOp(state,y);
	  data.applyOp(comp_ctrlstate[1],y);
	  // cerr<<"  state.norm="<<state.norm()<<"\n";
	}
      }
      catch (RVLException & e) {
	e<<"\ncalled from TimeStepOp::apply0\n";
	throw e;
      }
    }

    void applyAdj0(const Vector<T> & x0,
		   const Vector<T> & x1, 
		   Vector<T> & y) const {
      try {
	Vector<T> ctrlstate(step.getProductDomain());
	Components<T> comp_ctrlstate(ctrlstate);
	comp_ctrlstate[0].copy(x0);
	comp_ctrlstate[1].zero();

	T t = data.getMaxTime();
	step.setTime(t);
	T tbeg = min(src.getMinTime(), data.getMinTime());
	T tnext = t;
	//	cerr<<"******* adj: before call to LienarRestrictOp constr\n";
	//	cerr<<"******* x0:\n";
	//	x0.write(cerr);
	while (t - TOL*(tnext-t) > tbeg) {
	  data.setTime(t);
	  data.applyAdjOp(x1,comp_ctrlstate[1]);
	  step.applyAdj(ctrlstate,ctrlstate);
	  tnext = t;
	  t = step.getTime();
	  src.setTime(t);
	  src.applyAdjOp(comp_ctrlstate[1],y);
	}
      }	
      catch (RVLException & e) {
	e<<"\ncalled from TimeStepOp::applyAdj0\n";
	throw e;
      }
    }

    // here x0 is control, x1 is external source, y is external data
    // state does not appear in arg list!
    /** partial deriv wrt control. Dynamics:
	s_{n+1} = H[c]s_n + src_n
	differentiated wrt c:
	ds_{n+1} = (DH[c]dc)s_n + H[c]ds_n
	Two options:

	(1) use LinOpValOp methods:
	ds = H[c]ds
	ds += (DH[c]dc)s ropeval.getDeriv().applyPlusOp(s,ds), 
	rop = RestrictOp(step,s) = (c \rightarrow H[c]s)
	however you slice it, this requires two passes through ds data

	(2) use Op methods
	view op as F[c,s] = H[c]s
	DF[c,s](dc,ds) = (DH[c]dc)s + H[c]ds
	this preserves opportunities for loop fusion, but because (c,s) is
	updated, must be careful that aux data does not get recomputed every 
	step

	Also note that F[c,s_n] = s_{n+1}, that is step map is from (c,s) -> s.

    */
    void applyPartialDeriv0(const Vector<T> & x0,
			    const Vector<T> & x1,
			    const Vector<T> & dx0,
			    Vector<T> & dy) const {
      try {

	Vector<T> tangstate(tdom);
	Components<T> comp_tangstate(tangstate);
	Vector<T> & ctrlstate  = comp_tangstate[0];
	Vector<T> & dctrlstate = comp_tangstate[1];
	Components<T> comp_ctrlstate(ctrlstate);
	comp_ctrlstate[0].copy(x0);
	comp_ctrlstate[1].zero();
	Components<T> comp_dctrlstate(dctrlstate);
	comp_dctrlstate[0].copy(dx0);
	comp_dctrlstate[1].zero();

	dy.zero();

	//	cerr<<"sadfasdflj\n";
	//	OperatorEvaluation<T> stepeval(step,ctrlstate);
	
	T t = min(src.getMinTime(), data.getMinTime());
	step.setTime(t);
	T tend = data.getMaxTime();
	T tlast = t;
	while (t + TOL*(t-tlast) < tend) {
	  src.setTime(t);
	  src.applyOp(x1,comp_ctrlstate[1]);
	  step.applyTangent(tangstate,tangstate);	  
	  tlast = t;
	  t = step.getTime();
	  data.setTime(t);
	  data.applyOp(comp_dctrlstate[1],dy);
	}
      }
      catch (RVLException & e) {
	e<<"\ncalled from TimeStepOp::applyPartialDeriv0\n";
	throw e;
      }
    }
    
    void applyAdjPartialDeriv0(const Vector<T> & x0,
			       const Vector<T> & x1,
			       const Vector<T> & dy,
			       Vector<T> & dx0) const {}
    
    OperatorProductDomain<T> * clonePD() const {
      return new TimeStepOp(*this);
    }
    
  public:
    TimeStepOp(TSStep<T> & _step,
	       TSSample<T> & _src, 
	       TSSample<T> & _data,
	       bool _testflag = false)
      : step(_step), src(_src), data(_data),
	pdom((step.getProductDomain())[0],src.getDomain()),
	tdom(step.getDomain(),step.getDomain()),
	rnd(getpid(),-ScalarFieldTraits<T>::One(),
	    ScalarFieldTraits<T>::One()),
	testflag(_testflag) {
      if (testflag) {
	std::stringstream tmp;
	tmp << getpid();
	std::string fn = "TimeStepOp_TestReportPID" + tmp.str() + ".txt";
	report.open(fn,std::ofstream::app);
      }
    }
    TimeStepOp(TimeStepOp<T> const & op)
      : step(op.step), src(op.src), data(op.data),
	pdom((step.getProductDomain())[0],src.getDomain()),
	tdom(step.getDomain(),step.getDomain()),
	rnd(getpid(),-ScalarFieldTraits<T>::One(),
	    ScalarFieldTraits<T>::One()),
	testflag(op.testflag) {
      if (testflag) {
	std::stringstream tmp;
	tmp << getpid();
	std::string fn = "TimeStepOp_TestReportPID" + tmp.str() + ".txt";
	report.open(fn,std::ofstream::app);
      }
    }
    
    ~TimeStepOp() { report.close(); }

    ProductSpace<float> const & getProductDomain() const { return pdom; }
    Space<float> const & getRange() const { return data.getRange(); }

    // adjoint tests - x0 is arg so that any necessary constraints
    // may be obeyed - design assumes that 
    bool testSrcAdj() {
      if (testflag) {
	report<<"\n\n*** TimeStepOp: Src Adjoint Test\n\n";
	src.setTestTime();
	return AdjointTest(src,rnd,report);
      }
      else {
	RVLException e;
	e<<"Error: TimeStepOp::testSrcAdj\n";
	e<<"  attempt to perform test with testflag not set\n";
	e<<"  testflag must be set on construction\n";
	throw e;
      }
    }

    bool testDataAdj() {
      if (testflag) {
	report<<"\n\n*** TimeStepOp: Data Adjoint Test\n\n";
	data.setTestTime();
	return AdjointTest(data,rnd,report);
      }
      else {
	RVLException e;
	e<<"Error: TimeStepOp::testSrcAdj\n";
	e<<"  attempt to perform test with testflag not set\n";
	e<<"  testflag must be set on construction\n";
	throw e;
      }
    }

    bool testStepAdj0(Vector<T> const & x0) {
      if (testflag) {
	report<<"\n\n*** TimeStepOp: Step Adjoint Test Order 0\n\n";
	//	cerr<<"******* before call to LinRestrictOp constr\n";
	//	cerr<<"******* x0:\n";
	//	x0.write(cerr);

	LinRestrictTSStep<T> linstep(step,x0);
	return AdjointTest(linstep,rnd,report);
      }
      else {
	RVLException e;
	e<<"Error: TimeStepOp::testStepAdj0\n";
	e<<"  attempt to perform test with testflag not set\n";
	e<<"  testflag must be set on construction\n";
	throw e;
      }
    }

    bool testAll(Vector<T> const & x0) {
      if (testflag) {
	return (this->testSrcAdj() &&
		this->testDataAdj() &&
		this->testStepAdj0(x0));
      }
      else {
	RVLException e;
	e<<"Error: TimeStepOp::testAll\n";
	e<<"  attempt to perform test with testflag not set\n";
	e<<"  testflag must be set on construction\n";
	throw e;
      }
    }
    
    ostream & write(ostream & str) const {
      str<<"TimeStepOp: \n";
      str<<"*** time step \n";
      step.write(str);
      str<<"*** source sampler (input) \n";
      src.write(str);
      str<<"*** data sampler (output)\n";
      data.write(str);
      return str;
    }
  };
    
}

#endif