#ifndef DIST_MATRIX_HPP
#define DIST_MATRIX_HPP

#include <ThrustTypes.hpp>
#include <DistObject.hpp>
#include <DistMatrixProto.hpp>

/**
 * Distributed matrix library.
 *
 * This class allows VV graphs to be manipulated like vectors and matrices, with
 * the computations being done on the GPU. It uses distributed object to move 
 * data back and forth that is stored in the properties of VVGraph vertices and edges.
 * or stored in the attributes keyed by vertices and edges. 
 */
namespace mgx
{

  // Forward class declarations
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT> class DVector;
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT> class DMatrix;
  
  // --- Intermediate classes for operators ---
  
  //
  // Define the operations. These will hold the operands in a structure, and the
  // work will be done when "=" is invoked.
  //
  
  // Vector-vector addition
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class VVAddOp {
  public:
    DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec1;
    DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec2;
  
    VVAddOp(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec1, 
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec2) : vec1(p_vec1), vec2(p_vec2) {}
  };

  // Vector-vector subtraction
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class VVSubOp {
  public:
    DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec1;
    DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec2;
  
    VVSubOp(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec1, 
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec2) : vec1(p_vec1), vec2(p_vec2) {}
  };
  
  // Vector-scalar multiplication
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class VSMultOp {
  public:
    DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec;
    ScalarT a;
  
    VSMultOp(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec, ScalarT p_a) 
      : vec(p_vec), a(p_a) {}
  };

  // Vector saxpy
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class VSaxpyOp {
  public:
    DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec1;
    DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec2;
    ScalarT a;
  
    VSaxpyOp(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec1,
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec2, ScalarT p_a)
       : vec1(p_vec1), vec2(p_vec2), a(p_a) {}
  };  

  // Matrix-scalar multiplication
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class MSMultOp {
  public:
    DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &mat;
    ScalarT a;
    MSMultOp(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat, ScalarT p_a) 
      : mat(p_mat), a(p_a) {}
  };

  // Matrix-scalar addition
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class MSAddOp {
  public:
    DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &mat;
    ScalarT a;
  
    MSAddOp(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat, double p_a)
      : mat(p_mat), a(p_a) {}
  }; 
  
  // Matrix-matrix addition
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class MMAddOp {
  public:
    DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &mat1;
    DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &mat2;
 
    MMAddOp(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat1, 
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat2) : mat1(p_mat1), mat2(p_mat2) {}
  };
  
  //

  // Matrix-matrix subtraction
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class MMSubOp {
  public:
    DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &mat1;
    DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &mat2;
  
    MMSubOp(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat1, 
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat2) : mat1(p_mat1), mat2(p_mat2) {}
  };
  
  // Matrix-vector multiplication
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class MVMultOp {
  public:
    DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &mat;
    DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec;
  
    MVMultOp(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat, 
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec) : mat(p_mat), vec(p_vec) {}
  };
  
  //
  // Define the operators to make nice syntax
  //
  
  // Vector-vector addition
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  const VVAddOp<DistNhbdT, MatrixT, VectorT, ScalarT> 
    operator+(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec1, 
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec2)
  {
    VVAddOp<DistNhbdT, MatrixT, VectorT, ScalarT> ans(p_vec1, p_vec2);
    return ans;
  }
  
  // Vector-vector subtraction
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  const VVSubOp<DistNhbdT, MatrixT, VectorT, ScalarT> 
    operator-(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec1, 
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec2)
  {
    VVSubOp<DistNhbdT, MatrixT, VectorT, ScalarT> ans(p_vec1, p_vec2);
    return ans;
  }
  
  // Vector-scalar multiplication
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  const VSMultOp<DistNhbdT, MatrixT, VectorT, ScalarT> 
    operator*(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec, ScalarT p_a)
  {
    VSMultOp<DistNhbdT, MatrixT, VectorT, ScalarT> ans(p_vec, p_a);
    return ans;
  }

  // Vector saxpy
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  const VSaxpyOp<DistNhbdT, MatrixT, VectorT, ScalarT> 
    operator+(const VSMultOp<DistNhbdT, MatrixT, VectorT, ScalarT> &p_VSMultOp, 
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec)
  {
    VSaxpyOp<DistNhbdT, MatrixT, VectorT, ScalarT> ans(p_VSMultOp.vec, p_vec, p_VSMultOp.a);
    return ans;
  }
  
  // Matrix-matrix addition
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  const MMAddOp<DistNhbdT, MatrixT, VectorT, ScalarT> 
    operator+(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat1, 
      DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat2)
  {
    MMAddOp<DistNhbdT, MatrixT, VectorT, ScalarT> ans(p_mat1, p_mat2);
    return ans;
  }
  
  // Matrix-matrix subtraction
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  const MMSubOp<DistNhbdT, MatrixT, VectorT, ScalarT> 
    operator-(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat1, 
      DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat2)
  {
    MMSubOp<DistNhbdT, MatrixT, VectorT, ScalarT> ans(p_mat1, p_mat2);
    return ans;
  }
  
  // Matrix-scalar multiplication
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  const MSMultOp<DistNhbdT, MatrixT, VectorT, ScalarT> 
    operator*(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat, ScalarT p_a)
  {
    MSMultOp<DistNhbdT, MatrixT, VectorT, ScalarT> ans(p_mat, p_a);
    return ans;
  }
  
  // Matrix-vector multiplication
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  const MVMultOp<DistNhbdT, MatrixT, VectorT, ScalarT> 
    operator*(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &p_mat, 
      DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &p_vec)
  {
    MVMultOp<DistNhbdT, MatrixT, VectorT, ScalarT> ans(p_mat, p_vec);
    return ans;
  }
  
  //
  // Classes for distributed matrix and vector
  //
  // These allow VV graphs to be used like matrices and vectors
  //
  
  //
  // Distributed Vector class
  //
  template<typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT> 
  class DVector
  {
  public:
    // Type of neighborhood object
    typedef DistNhbdT DNhbd;
    // Type of distributed vector
    typedef DistVertex<DistNhbdT, VectorT> DVertex;
    // Type of a vertex
    typedef typename DNhbd::Vertex Vertex;
  
  private:
    DistNhbdT &_nhbd;
    DistVertex<DistNhbdT, VectorT> &_vobj;
  
  public:
    // Constructor for vector objects
    DVector(DistVertex<DistNhbdT, VectorT> &vobj) : 
      _nhbd(vobj.nhbd()), _vobj(vobj) {}
  
    ~DVector() {}
  
    // Return neighborhood object
    DistNhbdT &nhbd() { return _nhbd; }
  
    // Return vertex object
    DistVertex<DistNhbdT, VectorT> &vobj() { return _vobj; }
  
    // Allocate storage for a vector on GPU
    void alloc()
    {
      vobj().alloc();
    }
    
    // Send vector data to GPU
    void write()
    {
      vobj().write();
    }
    
    // Copy vector data from GPU
    void read()
    {
      vobj().read();
    }
    
    // Vector-vector assignment
    void operator=(DVector &v)
    {
      copyGPU(v.vobj().data(), vobj().data());
    }
  
    // Vector-scalar assignment
    void operator=(ScalarT a)
    {
      fillGPU(a, vobj().data());
    }
  
    // Vector-matrix assignment (assigns diagonal)
    void operator=(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &m)
    {
      copyGPU(m.vobj().data(), vobj().data());
    }
  
    // Vector-scalar multiply
    void operator=(VSMultOp<DistNhbdT, MatrixT, VectorT, ScalarT>  op)
    {
      multGPU(op.vec.vobj().data(), op.a, vobj().data());
    }

    // Vector-vector in place addition
    void operator+=(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec)
    {
      addGPU(vobj().data(), vec.vobj().data(), vobj().data());
    }
  
    // Vector-vector in place subtraction
    void operator-=(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec)
    {
      subGPU(vobj().data(), vec.vobj().data(), vobj().data());
    }
  
    // Vector-scalar in place multiplication
    void operator*=(ScalarT a)
    {
      multGPU(vobj().data(), a, vobj().data());
    }
  
    // Vector-scalar in place division
    void operator/=(ScalarT a)
    {
      if(a != 0)
        multGPU(vobj().data(), ScalarT(1)/a, vobj().data());
    }
  
    // Vector-vector addition
    void operator=(VVAddOp<DistNhbdT, MatrixT, VectorT, ScalarT>  op)
    {
      addGPU(op.vec1.vobj().data(), op.vec2.vobj().data(), vobj().data());
    }

    // Vector-vector subtraction
    void operator=(VVSubOp<DistNhbdT, MatrixT, VectorT, ScalarT>  op)
    {
      subtractGPU(op.vec1.vobj().data(), op.vec2.vobj().data(), vobj().data());
    }
  
    // Vector-vector multiplication
    const ScalarT operator*(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &vec)
    {
      ScalarT r;
      multGPU(vobj().data(), vec.vobj().data(), r);
      return(r);
    }

    // Vector saxpy
    void operator=(VSaxpyOp<DistNhbdT, MatrixT, VectorT, ScalarT>  op)
    {
      saxpyGPU(op.vec1.vobj().data(), op.vec2.vobj().data(), op.a, vobj().data());
    }
   
    // Matrix-vector multiplication
    void operator=(MVMultOp<DistNhbdT, MatrixT, VectorT, ScalarT>  op)
    {
      multGPU(nhbd().data(), op.mat.vobj().data(), op.mat.eobj().data(), op.vec.vobj().data(), vobj().data());
    }
  };
  
  /**
   * \class DMatrix DistMatrix.hpp <DistMatrix.hpp>
   *
   * Distributed matrix class. This class allows the graph to be used like a matrix, with
   * the operations performed on the GPU. It uses \class DistObject.
   */
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  class DMatrix 
  {
  public:
    // Type of neighborhood object
    typedef DistNhbdT DNhbd;
    // Type of distributed vector
    typedef DistVertex<DistNhbdT, MatrixT> DVertex;
    // Type of distributed vector
    typedef DistEdge<DistNhbdT, MatrixT> DEdge;
    // Type of a vertex
    typedef typename DNhbd::Vertex Vertex;
    // Type of an edge
    typedef typename DNhbd::Edge Edge;
  
  private:
    DNhbd &_nhbd;
    DVertex &_vobj;
    DEdge &_eobj;
  
  public:
    // Constructor for matrix objects
    DMatrix(DVertex &vobj, DEdge &eobj) : 
          _nhbd(vobj.nhbd()), _vobj(vobj), _eobj(eobj)  {}
  
    ~DMatrix() {}
  
    // Return neighborhood object
    DistNhbdT &nhbd() { return _nhbd; };
  
    // Return vertex object
    DistVertex<DistNhbdT, MatrixT> &vobj() { return _vobj; };
  
    // Return edge object
    DistEdge<DistNhbdT, MatrixT> &eobj() { return _eobj; };
  
    // Allocate storage for a matrix on GPU
    void alloc()
    {
      vobj().alloc();
      eobj().alloc();
    }
  
    // Send matrix data to GPU
    void write()
    {
      vobj().write();
      eobj().write();
    }
  
    // Copy matrix data from GPU
    void read()
    {
    }
    
    // Matrix assigment
    void operator=(DMatrix &m)
    {
      copyGPU(m.vobj().data(), vobj().data());
      copyGPU(m.eobj().data(), eobj().data());
    }
  
    // Matrix-scalar assignment
    void operator=(ScalarT a)
    {
      fillGPU(a, vobj().data());
      fillGPU(a, eobj().data());
    }
  
    // Matrix-vector assignment (assigns to diagonal)
    void operator=(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &v)
    {
      copyGPU(v.vobj().data(), vobj().data());
    }
  
    // Matrix-scalar multiply
    void operator=(MSMultOp<DistNhbdT, MatrixT, VectorT, ScalarT>  op)
    {
      multGPU(op.mat.vobj().data(), op.a, vobj().data());
      multGPU(op.mat.eobj().data(), op.a, eobj().data());
    }
  
    // Matrix-matrix in place addition
    void operator+=(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &mat)
    {
      addGPU(vobj().data(), mat.vobj().data(), vobj().data());
      addGPU(eobj().data(), mat.eobj().data(), eobj().data());
    }
  
    // Matrix-matrix in place subtraction
    void operator-=(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &mat)
    {
      subGPU(vobj().data(), mat.vobj().data(), vobj().data());
      subGPU(eobj().data(), mat.eobj().data(), eobj().data());
    }
  
    // Matrix-scalar in place multiplication
    void operator*=(ScalarT a)
    {
      multGPU(vobj().data(), a, vobj().data());
      multGPU(eobj().data(), a, eobj().data());
    }
  
    // Matrix-scalar in place division
    void operator/=(ScalarT a)
    {
      multGPU(vobj().data(), ScalarT(1)/a, vobj().data());
      multGPU(eobj().data(), ScalarT(1)/a, eobj().data());
    }
  
    // Matrix-matrix addition
    void operator=(MMAddOp<DistNhbdT, MatrixT, VectorT, ScalarT>  op)
    {
      addGPU(op.mat1.vobj().data(), op.mat2.vobj().data(), vobj().data());
      addGPU(op.mat1.vobj().data(), op.mat2.vobj().data(), eobj().data());
    }
  
    // Matrix-matrix subtraction
    void operator=(MMSubOp<DistNhbdT, MatrixT, VectorT, ScalarT>  op)
    {
      subtractGPU(op.mat1.vobj().data(), op.mat2.vobj().data(), vobj().data());
      subtractGPU(op.mat1.eobj().data(), op.mat2.eobj().data(), eobj().data());
    }

    // Add to diagonal
    double addToDiagonal(double a)
    {
      addToDiagGPU(a, vobj().data());
      return a;
    }
  };

  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  double maxDiagonal(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &m)
  {
    ScalarT r;
    maxGPU(m.vobj().data(), r);
    return r;
  }

  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  double maxOffDiagonal(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &m)
  {
    ScalarT r;
    maxGPU(m.eobj().data(), r);
    return r;
  }

  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  double maxEntry(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &m)
  {
    ScalarT rv, re;
    maxGPU(m.vobj().data(), rv);
    maxGPU(m.eobj().data(), re);
    return max(rv, re);
  }

  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  double minDiagonal(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &m)
  {
    ScalarT r;
    minGPU(m.vobj().data(), r);
    return r;
  }

  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  double minOffDiagonal(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &m)
  {
    ScalarT r;
    minGPU(m.eobj().data(), r);
    return r;
  }

  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  double minEntry(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &m)
  {
    ScalarT rv, re;
    minVGPU(m.vobj().data(), rv);
    minVGPU(m.eobj().data(), re);
    return min(rv, re);
  }

  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  ScalarT norm(DVector<DistNhbdT, MatrixT, VectorT, ScalarT> &v)
  {
    ScalarT r;
    r = v * v;
    return pow(r, 0.5);
  }
  
  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  ScalarT norm(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &m)
  {
    ScalarT rv, re;
    multGPU(m.vobj().data(), m.vobj().data(), rv);
    multGPU(m.eobj().data(), m.eobj().data(), re);
    return pow(rv + re, 0.5);
  }

  template <typename DistNhbdT, typename MatrixT, typename VectorT, typename ScalarT>
  void jacobiPreCond(DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &m, DMatrix<DistNhbdT, MatrixT, VectorT, ScalarT> &a)
  {
    m = 0;
    jacobiPreCondGPU(m.nhbd().data(), m.vobj().data(), m.eobj().data(), a.vobj().data(), a.eobj().data());
  }
}
  
#endif
