// cudalattice.h -- lattice forward/backward functions CUDA execution
//
// F. Seide, V-hansu

#pragma once
#include <stdexcept>            // (for NOCUDA version only)
#include "latticestorage.h"        // basic data types for storing lattices
#define __kernel_emulation__            // allow the compilation of CUDA kernels on the CPU
#include "latticefunctionskernels.h"    // for data types passed to kernel functions
#undef __kernel_emulation__

#include "Matrix.h"

namespace msra { namespace cuda {

// ---------------------------------------------------------------------------
// how to add vector types here:
//  - for types with methods that operate on the vector, follow the example of 'somedatavector'
//  - for types without extra methods, follow the example of 'ushortvector'
//  - the total set of classes that need to be defined is:
//     - XXXvector       --the interface that the calling code sees (in this header file)
//     - XXXvectorimpl   --glue for each method: switch CUDA device, then call kernel (in XXXvectorops, cudalattice.cpp)
//     - XXXvectorops    --contains the actual kernel functions (cudalatticeops.h, cudalatticeops.cu.h)
//  - comments:
//     - follow precisely this pattern
//     - for non-vector elements you need to pass by reference, you can still use this with a size-1 vector
//     - and for functions that have no 'this', use a size-0 vector
//     - note that you cannot add data members to XXXvector, as there is no mechanism to sync between GPU and CPU; only vector elements are synced
//     - remember to Get Latest in both projects to see the change, for cudalattice.h and latticestorage.h
// ---------------------------------------------------------------------------


// a vector type; use this as a basetype
// Note that this will only be instantiated for types known inside this lib, and a newvector<> function must be exported for each.
template<typename ELEMTYPE> struct/*interface*/ vectorbase
{
    virtual ~vectorbase() { }
    virtual void allocate(size_t n) = 0;
    virtual size_t size() const throw() = 0;
    virtual void assign(const ELEMTYPE * p, size_t n, bool synchronize) = 0;
    template<class VECTOR> void assign(const VECTOR & v, bool synchronize) { allocate(v.size()); if (!v.empty()) assign(&v[0], v.size(), synchronize); }
    virtual void fetch(ELEMTYPE * p, size_t n, bool synchronize) const = 0;
    template<class VECTOR> void fetch(VECTOR & v, bool synchronize) const { v.resize(size()); if (!v.empty()) fetch(&v[0], v.size(), synchronize); }
    typedef ELEMTYPE elemtype;
};

// ---------------------------------------------------------------------------
// vectors of these with custom functions
// The XXXvector classes must derive from vectorbase<XXX>.
// ---------------------------------------------------------------------------

struct somedatavector : public vectorbase<msra::lattices::somedata>
{
    // must implement all members of vectorbase<>, and can add operations here
    virtual int somedataoperation (size_t arg) = 0;
};

typedef vectorbase<unsigned short> ushortvector;
typedef vectorbase<float> floatvector;
typedef vectorbase<double> doublevector;
typedef vectorbase<unsigned int> uintvector;
typedef vectorbase<size_t> sizetvector;
typedef vectorbase<msra::lattices::lrhmmdef> lrhmmdefvector;
typedef vectorbase<msra::lattices::lr3transP> lr3transPvector;
typedef vectorbase<msra::lattices::nodeinfo> nodeinfovector;
typedef vectorbase<msra::lattices::edgeinfowithscores> edgeinfowithscoresvector;
typedef vectorbase<msra::lattices::aligninfo> aligninfovector;

struct latticefunctions : public vectorbase<msra::lattices::empty> 
{ 
    virtual void edgealignment (const lrhmmdefvector & hmms, const lr3transPvector & transPs, const size_t spalignunitid, 
                                const size_t silalignunitid, const Microsoft::MSR::CNTK::Matrix<float>& logLLs, const nodeinfovector & nodes, 
                                const edgeinfowithscoresvector & edges, const aligninfovector & aligns,
                                const uintvector & alignoffsets, ushortvector & backptrstorage, const sizetvector & backptroffsets,
                                ushortvector & alignresult, floatvector & edgeacscores) = 0;        // output
    virtual void forwardbackwardlattice (const size_t * batchsizeforward, const size_t * batchsizebackward, 
                                         const size_t numlaunchforward, const size_t numlaunchbackward,
                                         const size_t spalignunitid, const size_t silalignunitid,
                                         const floatvector & edgeacscores, const edgeinfowithscoresvector & edges, 
                                         const nodeinfovector & nodes, const aligninfovector & aligns, 
                                         const ushortvector & alignoutput, const uintvector & alignoffsets,
                                         doublevector & logpps, doublevector & logalphas, doublevector & logbetas,
                                         const float lmf, const float wp, const float amf, const float boostingfactor, const bool returnEframescorrect, 
                                         const ushortvector & uids, const ushortvector & senone2classmap, 
                                         doublevector & logaccalphas, doublevector & logaccbetas,
                                         doublevector & logframescorrectedge, doublevector & logEframescorrect, 
                                         doublevector & Eframescorrectbuf, double & logEframescorrecttotal, double & totalfwscore) = 0;
    virtual void sMBRerrorsignal (const ushortvector & alignstateids, const uintvector & alignoffsets,
                                  const edgeinfowithscoresvector & edges, const nodeinfovector & nodes, 
                                  const doublevector & logpps, const float amf, const doublevector & logEframescorrect, 
                                  const double logEframescorrecttotal, Microsoft::MSR::CNTK::Matrix<float>& dengammas, Microsoft::MSR::CNTK::Matrix<float>& dengammasbuf) = 0;
    virtual void mmierrorsignal (const ushortvector & alignstateids, const uintvector & alignoffsets,
                                 const edgeinfowithscoresvector & edges, const nodeinfovector & nodes, 
                                 const doublevector & logpps, Microsoft::MSR::CNTK::Matrix<float>& dengammas) = 0;
    virtual void stateposteriors (const ushortvector & alignstateids, const uintvector & alignoffsets,
                                  const edgeinfowithscoresvector & edges, const nodeinfovector & nodes, 
                                  const doublevector & logqs, Microsoft::MSR::CNTK::Matrix<float>& logacc) = 0;
                                           
};

// ---------------------------------------------------------------------------
// factor methods
// ---------------------------------------------------------------------------

somedatavector * newsomedatavector(size_t deviceid);
ushortvector * newushortvector(size_t deviceid);
uintvector * newuintvector(size_t deviceid);
floatvector * newfloatvector(size_t deviceid);
doublevector * newdoublevector(size_t deviceid);
sizetvector * newsizetvector(size_t deviceid);
latticefunctions * newlatticefunctions(size_t deviceid);
lrhmmdefvector * newlrhmmdefvector(size_t deviceid);
lr3transPvector * newlr3transPvector(size_t deviceid);
nodeinfovector * newnodeinfovector(size_t deviceid);
edgeinfowithscoresvector * newedgeinfovector(size_t deviceid);
aligninfovector * newaligninfovector(size_t deviceid);

};};