#include <vector>
#include <boost/numeric/ublas/traits.hpp>
#include <boost/bind.hpp>
#include "mismatch.H"

namespace vec {
  template<typename T1, typename T2=T1, typename ret_t=T1>
  struct plus {
    typedef T1 first_argument_type;
    typedef T2 second_argument_type;
    typedef ret_t result_type;
    ret_t operator() (const T1& x, const T2& y) {
      return x + y;
    }
  };

  template<typename T1, typename T2=T1, typename ret_t=T1>
  struct minus {
    typedef T1 first_argument_type;
    typedef T2 second_argument_type;
    typedef ret_t result_type;
    ret_t operator() (const T1& x, const T2& y) {
      return x - y;
    }
  };

  template<typename T1, typename T2=T1, typename ret_t=T1>
  struct multiplies {
    typedef T1 first_argument_type;
    typedef T2 second_argument_type;
    typedef ret_t result_type;
    ret_t operator() (const T1& x, const T2& y) {
      return x * y;
    }
  };

  template<typename T1, typename T2=T1, typename ret_t=T1>
  struct divides  {
    typedef T1 first_argument_type;
    typedef T2 second_argument_type;
    typedef ret_t result_type;
    ret_t operator() (const T1& x, const T2& y) {
      return x / y;
    }
  };

  template<typename T1, typename T2=T1>
  struct pluseq  {
    typedef T1 first_argument_type;
    typedef T2 second_argument_type;
    typedef T1& result_type;
    T1& operator() (T1& a, const T1& b) {
      a += b;
      return a;
    }
  };

  template<typename T1, typename T2=T1>
  struct minuseq  {
    typedef T1 first_argument_type;
    typedef T2 second_argument_type;
    typedef T1& result_type;
    T1& operator() (T1& a, const T1& b) {
      a -= b;
      return a;
    }
  };

  template<typename T1, typename T2=T1>
  struct multiplieseq  {
    typedef T1 first_argument_type;
    typedef T2 second_argument_type;
    typedef T1& result_type;
    T1& operator() (T1& a, const T1& b) {
      a *= b;
      return a;
    }
  };

  template<typename T1, typename T2=T1>
  struct divideseq  {
    typedef T1 first_argument_type;
    typedef T2 second_argument_type;
    typedef T1& result_type;
    T1& operator() (T1& a, const T1& b) {
      a /= b;
      return a;
    }
  };

  template<typename type1, typename type2>
  struct arith_traits;

  template<typename type1, typename type2>
  struct arith_traits<std::vector<type1>,std::vector<type2> > {
    typedef typename boost::numeric::ublas::promote_traits<type1, type2>::promote_type ret_type;
    typedef std::vector<ret_type> vec_ret_type;
  };

  template<typename type1, typename type2>
  struct arith_traits<std::vector<type1>,type2> {
    typedef typename boost::numeric::ublas::promote_traits<type1, type2>::promote_type ret_type;
    typedef std::vector<ret_type> vec_ret_type;
  };

  template<typename type1, typename type2, template <typename A, typename B, typename C> class func>
  struct vec_arith;

  template<typename type1, typename type2, template <typename A, typename B, typename C> class func>
  struct vec_arith<std::vector<type1>, std::vector<type2>, func> {
    typedef typename arith_traits<std::vector<type1>,std::vector<type2> >::vec_ret_type vec_ret_type;
    typedef typename arith_traits<std::vector<type1>,std::vector<type2> >::ret_type ret_type;
  
    vec_ret_type operator() (std::vector<type1> const& a, std::vector<type2> const& b) {
      if (a.size() != b.size()) throw size_mismatch();
      vec_ret_type ret (a.size());
      std::transform (a.begin(), a.end(), b.begin(), ret.begin(), func<type1,type2,ret_type>());
      return ret;
    }
  };

  template<typename type1, typename type2, template <typename A,typename B, typename C> class func>
  struct vec_arith<std::vector<type1>, type2, func> {
    typedef typename arith_traits<std::vector<type1>,type2>::vec_ret_type vec_ret_type;
    typedef typename arith_traits<std::vector<type1>,type2>::ret_type ret_type;

    vec_ret_type operator() (std::vector<type1>const& a, type2 b) {
      vec_ret_type ret (a.size());
      std::transform (a.begin(), a.end(), ret.begin(), boost::bind (func<type1,type2,ret_type>(), _1, b));
      return ret;
    }
  };

  // Dispatch for out-of-place ops
  template<typename type1, typename type2>
  inline typename arith_traits<std::vector<type1>,std::vector<type2> >::vec_ret_type operator+ (std::vector<type1> const& a, std::vector<type2> const& b) {
    return vec_arith<std::vector<type1>,std::vector<type2>,vec::plus>() (a, b);
  }

  template<typename type1, typename type2>
  inline typename arith_traits<std::vector<type1>,type2>::vec_ret_type operator+ (std::vector<type1> const& a, type2 const& b) {
    return vec_arith<std::vector<type1>,type2,vec::plus>() (a, b);
  }

  template<typename type1, typename type2>
  inline typename arith_traits<std::vector<type1>,std::vector<type2> >::vec_ret_type operator- (std::vector<type1> const& a, std::vector<type2> const& b) {
    return vec_arith<std::vector<type1>,std::vector<type2>,vec::minus>() (a, b);
  }

  template<typename type1, typename type2>
  inline typename arith_traits<std::vector<type1>,type2>::vec_ret_type operator- (std::vector<type1> const& a, type2 const& b) {
    return vec_arith<std::vector<type1>,type2,vec::minus>() (a, b);
  }

  template<typename type1, typename type2>
  inline typename arith_traits<std::vector<type1>,std::vector<type2> >::vec_ret_type operator* (std::vector<type1> const& a, std::vector<type2> const& b) {
    return vec_arith<std::vector<type1>,std::vector<type2>,vec::multiplies>() (a, b);
  }

  template<typename type1, typename type2>
  inline typename arith_traits<std::vector<type1>,type2>::vec_ret_type operator* (std::vector<type1> const& a, type2 const& b) {
    return vec_arith<std::vector<type1>,type2,vec::multiplies>() (a, b);
  }

  template<typename type1, typename type2>
  inline typename arith_traits<std::vector<type1>,std::vector<type2> >::vec_ret_type operator/ (std::vector<type1> const& a, std::vector<type2> const& b) {
    return vec_arith<std::vector<type1>,std::vector<type2>,vec::divides>() (a, b);
  }

  template<typename type1, typename type2>
  inline typename arith_traits<std::vector<type1>,type2>::vec_ret_type operator/ (std::vector<type1> const& a, type2 const& b) {
    return vec_arith<std::vector<type1>,type2,vec::divides>() (a, b);
  }


  // Define in-place ops for vectors
  template<typename type1, typename type2, template <typename A, typename B> class func>
  struct vector_assign_arith;

  template<typename type1, typename type2, template <typename A, typename B> class func>
  struct vector_assign_arith<std::vector<type1>, std::vector<type2>, func> {
    typedef std::vector<type1> ret_type;
    ret_type& operator() (std::vector<type1>& a, std::vector<type2>const& b) {
      if (a.size() != b.size()) throw size_mismatch();
      typename std::vector<type1>::iterator i = a.begin();
      typename std::vector<type2>::const_iterator j = b.begin();
      for (; i != a.end(); i++, j++)
	func<type1,type2>() (*i, *j);

      return a;
    }
  };

  template<typename type1, typename type2, template <typename A, typename B> class func>
  struct vector_assign_arith<std::vector<type1>, type2, func> {
    typedef std::vector<type1> ret_type;
    ret_type& operator() (std::vector<type1>& a, type2 const& b) {
      typename std::vector<type1>::iterator i = a.begin();
      for (; i != a.end(); i++)
	func<type1,type2>() (*i, b);

      return a;
    }
  };

  // Dispatch for in-place ops
  template<typename type1, typename type2>
  inline std::vector<type1>& operator+= (std::vector<type1> & a, type2 const& b) {
    return vector_assign_arith<std::vector<type1>,type2,vec::pluseq>() (a, b);
  }

  template<typename type1, typename type2>
  inline std::vector<type1>& operator-= (std::vector<type1> & a, type2 const& b) {
    return vector_assign_arith<std::vector<type1>,type2,vec::minuseq>() (a, b);
  }

  template<typename type1, typename type2>
  inline std::vector<type1>& operator*= (std::vector<type1> & a, type2 const& b) {
    return vector_assign_arith<std::vector<type1>,type2,vec::multiplieseq>() (a, b);
  }

  template<typename type1, typename type2>
  inline std::vector<type1>& operator/= (std::vector<type1> & a, type2 const& b) {
    return vector_assign_arith<std::vector<type1>,type2,vec::divideseq>() (a, b);
  }

}