module vector

  use typedef, only: &
       real64

  implicit none

  private

  public :: &
       normp2, norm2, norm, &
       cross_product, dot_product, &
       double_cross_product, &
       cross_product_tensor_contract

  type, public::  vec
     real(kind=real64), dimension(1:3) :: &
          v

   contains

     procedure, private :: pos
     procedure, private :: add_vu
     generic, public :: operator (+) => pos, add_vu

     procedure, private :: neg
     procedure, private :: sub_vu
     generic, public :: operator (-) => neg, sub_vu

     procedure, private :: dot_product_v
     procedure, private :: mul_vs
     procedure, pass(u), private :: mul_su
     generic, public :: operator (*) => dot_product_v, mul_vs, mul_su

     procedure, private :: div_vs
     generic, public :: operator (/) => div_vs

     procedure, private :: cross_product_v
     generic, public :: operator (**) => cross_product_v

     procedure, private :: eq_vu
     procedure, private :: eq_vs
     procedure, pass(u), private :: eq_su
     generic, public :: operator (==) => eq_vu, eq_vs, eq_su

     procedure, private :: ne_vu
     procedure, private :: ne_vs
     procedure, pass(u), private :: ne_su
     generic, public :: operator (/=) => ne_vu, ne_vs, ne_su

     procedure, public  :: abs => norm_v
     procedure, public  :: norm2 => norm_v
     procedure, public  :: normp2_v
     procedure, public  :: normalized_v

     procedure, public  :: cross => cross_product_v
     procedure, public  :: dot => dot_product_v

  end type vec

  interface vec
     module procedure init_
     module procedure init_s
     module procedure init_v
     module procedure init_v3
  end interface vec

  ! interface assignment (=)
  !    module procedure assign_v
  !    module procedure assign_s
  ! end interface assignment (=)

  interface cross_product
     module procedure cross_product_v
     module procedure cross_product_a
  end interface cross_product

  interface double_cross_product
     module procedure double_cross_product_v
     module procedure double_cross_product_a
  end interface double_cross_product

  interface normp2
     module procedure normp2_v
     module procedure normp2_a
  end interface normp2

  interface norm2
     module procedure norm_v
  end interface norm2

  interface norm
     module procedure norm_a
  end interface norm

  interface dot_product
     module procedure dot_product_v
  end interface dot_product

  interface cross_product_tensor_contract
     module procedure cross_product_tensor_contract_a
  end interface cross_product_tensor_contract

contains

  !=======================================================================
  ! helper functions
  !=======================================================================

  !-----------------------------------------------------------------------
  ! cross product
  !-----------------------------------------------------------------------

  pure function cross_product_a(x, y) result(z)

    use typedef, only: &
         real64

    implicit none

    real(kind=real64), dimension(3), intent(in) :: &
         x, y
    real(kind=real64), dimension(3) :: &
         z

    z(1) = x(2) * y(3) - x(3) * y(2)
    z(2) = x(3) * y(1) - x(1) * y(3)
    z(3) = x(1) * y(2) - x(2) * y(1)

  end function cross_product_a


  elemental pure function cross_product_v(x, y) result(z)

    implicit none

    class(vec), intent(in) :: &
         x, y
    type(vec) :: &
         z

    z%v(1) = x%v(2) * y%v(3) - x%v(3) * y%v(2)
    z%v(2) = x%v(3) * y%v(1) - x%v(1) * y%v(3)
    z%v(3) = x%v(1) * y%v(2) - x%v(2) * y%v(1)

  end function cross_product_v


  pure function cross_product_tensor_contract_a(p, q) result(v)

    implicit none

    real(kind=real64), dimension(3,3), intent(in) :: &
         p, q
    real(kind=real64), dimension(3) :: &
         v

    v(1) = sum(p(2,:) * q(3,:) - p(3,:) * q(2,:))
    v(2) = sum(p(3,:) * q(1,:) - p(1,:) * q(3,:))
    v(3) = sum(p(1,:) * q(2,:) - p(2,:) * q(1,:))

  end function cross_product_tensor_contract_a


  !-----------------------------------------------------------------------
  ! double cross product
  !-----------------------------------------------------------------------

  pure function double_cross_product_a(x, y, x2) result(z)

    use typedef, only: &
         real64

    implicit none

    real(kind=real64), dimension(3), intent(in) :: &
         x, y
    real(kind=real64), intent (in), optional :: &
         x2
    real(kind=real64), dimension(3) :: &
         z
    real(kind=real64) :: &
         xx, xy

    if (present(x2)) then
       xx = x2
    else
       ! xx = dot_product(x(:), x(:))
       xx = x(1) * x(1) + x(2) * x(2) + x(3) * x(3)
    endif

    ! xy = dot_product(x(:), y(:))
    xy = x(1) * y(1) + x(2) * y(2) + x(3) * y(3)

    z(:) = x(:) * xy - y(:) * xx

  end function double_cross_product_a


  elemental pure function double_cross_product_v(x, y, x2) result(z)

    use typedef, only: &
         real64

    implicit none

    type(vec), intent(in) :: &
         x, y
    real(kind=real64), intent (in), optional :: &
         x2
    type(vec) :: &
         z
    real(kind=real64) :: &
         xx, xy

    if (present(x2)) then
       xx = x2
    else
       ! xx = dot_product(x(:), x(:))
       xx = x%v(1) * x%v(1) + x%v(2) * x%v(2) + x%v(3) * x%v(3)
    endif

    ! xy = dot_product(x(:), y(:))
    xy = x%v(1) * y%v(1) + x%v(2) * y%v(2) + x%v(3) * y%v(3)

    z%v(:) = x%v(:) * xy - y%v(:) * xx

  end function double_cross_product_v

  !-----------------------------------------------------------------------
  ! norm
  !-----------------------------------------------------------------------

  pure function norm_a(x) result(t)

    use typedef, only: &
         real64

    implicit none

    real(kind=real64), dimension(3), intent(in) :: &
         x
    real(kind=real64) :: &
         t

    t = sqrt(x(1) * x(1) + x(2) * x(2) + x(3) * x(3))

  end function norm_a


  pure function norm_v(x) result(t)

    implicit none

    class(vec), intent(in) :: &
         x
    real(kind=real64) :: &
         t

    ! y = dot_product(x(:), x(:))
    t = sqrt(x%v(1) * x%v(1) + x%v(2) * x%v(2) + x%v(3) * x%v(3))

  end function norm_v

  !-----------------------------------------------------------------------
  ! normp2
  !-----------------------------------------------------------------------

  pure function normp2_a(x) result(t)

    use typedef, only: &
         real64

    implicit none

    real(kind=real64), dimension(3), intent(in) :: &
         x
    real(kind=real64) :: &
         t

    t = x(1) * x(1) + x(2) * x(2) + x(3) * x(3)

  end function normp2_a


  pure function normp2_v(x) result(t)

    implicit none

    class(vec), intent(in) :: &
         x
    real(kind=real64) :: &
         t

    ! y = dot_product(x(:), x(:))
    t = x%v(1) * x%v(1) + x%v(2) * x%v(2) + x%v(3) * x%v(3)

  end function normp2_v

  !-----------------------------------------------------------------------
  ! normalized
  !-----------------------------------------------------------------------

  elemental pure function normalized_v(v) result(w)
    class(vec), intent(in) :: &
         v
    type(vec) :: &
         w

    real(kind=real64) :: &
         u

    ! u = norm2(v%v)
    u = sqrt(v%v(1) * v%v(1) + v%v(2) * v%v(2) + v%v(3) * v%v(3))

    if (u == 0.d0) then
       w%v = 0.d0
       return
    endif

    w%v = v%v / u

  end function normalized_v

  !-----------------------------------------------------------------------
  ! init functions
  !-----------------------------------------------------------------------

  function init_() result(w)
    type(vec) :: &
         w

    w%v = 0.d0
  end function init_

  pure function init_s(s) result(w)
    real(kind=real64), intent(in) :: &
         s
    type(vec) :: &
         w

    w%v = s
  end function init_s

  pure function init_v(v) result(w)
    type(vec), intent(in) :: &
         v
    type(vec) :: &
         w

    w%v = v%v
  end function init_v

  function init_v3(v1, v2, v3) result(w)
    real(kind=real64), intent(in) :: &
         v1, v2, v3
    type(vec) :: &
         w

    w%v(1) = v1
    w%v(2) = v2
    w%v(3) = v3
  end function init_v3

  !-----------------------------------------------------------------------
  ! assignmnt functions
  !-----------------------------------------------------------------------

  ! elemental pure subroutine assign_s(w, s)
  !   type(vec), intent(out) :: &
  !        w
  !   real(kind=real64), intent(in) :: &
  !        s

  !   w%v = s
  ! end subroutine assign_s


  ! elemental pure subroutine assign_v(w, v)
  !   type(vec), intent(out) :: &
  !        w
  !   type(vec), intent(in) :: &
  !        v

  !   w%v = v%v
  ! end subroutine assign_v

  !-----------------------------------------------------------------------
  ! operator + functions
  !-----------------------------------------------------------------------

  elemental pure function add_vu(v, u) result(w)
    class(vec), intent(in) :: &
         v, u
    type(vec) :: &
         w

    w%v = v%v + u%v
  end function add_vu


  elemental pure function pos(v) result(w)
    class(vec), intent(in) :: &
         v
    type(vec) :: &
         w

    w%v = v%v

  end function pos

  !-----------------------------------------------------------------------
  ! operator - functions
  !-----------------------------------------------------------------------

  elemental pure function sub_vu(v, u) result(w)
    class(vec), intent(in) :: &
         v, u
    type(vec) :: &
         w

    w%v = v%v - u%v
  end function sub_vu


  elemental pure function neg(v) result(w)
    class(vec), intent(in) :: &
         v
    type(vec) :: &
         w

    w%v = -v%v

  end function neg

  !-----------------------------------------------------------------------
  ! operator * functions
  !-----------------------------------------------------------------------

  elemental pure function mul_vs(v, s) result(w)
    class(vec), intent(in) :: &
         v
    real(kind=real64), intent(in) :: &
         s
    type(vec) :: &
         w

    w%v = v%v * s

  end function mul_vs

  elemental pure function mul_su(s, u) result(w)
    class(vec), intent(in) :: &
         u
    real(kind=real64), intent(in) :: &
         s
    type(vec) :: &
         w

    w%v = u%v * s

  end function mul_su

  elemental pure function dot_product_v(v, u) result(s)
    class(vec), intent(in) :: &
         v, u
    real(kind=real64) :: &
         s

    s = v%v(1) * u%v(1) + v%v(2) * u%v(2) + v%v(3) * u%v(3)

  end function dot_product_v

  !-----------------------------------------------------------------------
  ! operator / functions
  !-----------------------------------------------------------------------

  elemental pure function div_vs(v, s) result(w)
    class(vec), intent(in) :: &
         v
    real(kind=real64), intent(in) :: &
         s
    type(vec) :: &
         w

    w%v = v%v * (1.d0 / s)

  end function div_vs

  !-----------------------------------------------------------------------
  ! operator == functions
  !-----------------------------------------------------------------------

  elemental pure function eq_vs(v, s) result(l)
    class(vec), intent(in) :: &
         v
    real(kind=real64), intent(in) :: &
         s
    logical :: &
         l

    l = all(v%v(1:3) == s)

  end function eq_vs


  elemental pure function eq_su(s, u) result(l)
    class(vec), intent(in) :: &
         u
    real(kind=real64), intent(in) :: &
         s
    logical :: &
         l

    l = all(u%v(1:3) == s)

  end function eq_su


  elemental pure function eq_vu(v, u) result(l)
    class(vec), intent(in) :: &
         v, u
    logical :: &
         l

    l = all(v%v == u%v)

  end function eq_vu

  !-----------------------------------------------------------------------
  ! operator /= functions
  !-----------------------------------------------------------------------

  elemental pure function ne_vs(v, s) result(l)
    class(vec), intent(in) :: &
         v
    real(kind=real64), intent(in) :: &
         s
    logical :: &
         l

    l = any(v%v(1:3) /= s)

  end function ne_vs


  elemental pure function ne_su(s, u) result(l)
    class(vec), intent(in) :: &
         u
    real(kind=real64), intent(in) :: &
         s
    logical :: &
         l

    l = any(u%v(1:3) /= s)

  end function ne_su


  elemental pure function ne_vu(v, u) result(l)
    class(vec), intent(in) :: &
         v, u
    logical :: &
         l

    l = any(v%v /= u%v)

  end function ne_vu

end module vector