! (C) Alexander Heger (2022), All rights reserved.

module quaternion

  use, intrinsic :: ieee_arithmetic, only: &
       ieee_value, ieee_signaling_nan

  use typedef, only: &
       real64

  use vector, only: &
       vec, cross_product

  implicit none
  private

  public :: &
       assignment(=), &
       real, aimag, conjg, &
       log, exp, sqrt, abs, inv, &
       normalized, norm2, normp2

  real(kind=real64), parameter :: &
       lim_prec = 1.d-17, &
       lim_exp = 1.d-300

  type, public::  quat
     real(kind=real64), dimension(0:3) :: &
          f
   contains

     procedure, private :: pos
     procedure, private :: add_qp
     procedure, private :: add_qs
     procedure, pass(p), private :: add_sp
     generic, public :: operator (+) => pos, add_qp, add_qs, add_sp

     procedure, private :: neg
     procedure, private :: sub_qp
     procedure, private :: sub_qs
     procedure, pass(p), private :: sub_sp
     generic, public :: operator (-) => neg, sub_qp, sub_qs, sub_sp

     procedure, private :: mul_qp
     procedure, private :: mul_qs
     procedure, pass(p), private :: mul_sp
     generic, public :: operator (*) => mul_qp, mul_qs, mul_sp

     procedure, private :: div_qp
     procedure, private :: div_qs
     procedure, pass(p), private :: div_sp
     generic, public :: operator (/) => div_qp, div_qs, div_sp

     procedure, private :: eq_qp
     procedure, private :: eq_qs
     procedure, pass(p), private :: eq_sp
     generic, public :: operator (==) => eq_qp, eq_qs, eq_sp

     procedure, private :: ne_qp
     procedure, private :: ne_qs
     procedure, pass(p), private :: ne_sp
     generic, public :: operator (/=) => ne_qp, ne_qs, ne_sp

     procedure, private :: pow_qi
     procedure, private :: pow_qs
     procedure, private :: pow_qp
     procedure, pass(p), private :: pow_sp
     generic,   public  :: operator(**) => pow_qi, pow_qs, pow_qp, pow_sp

     procedure, public  :: abs => norm
     procedure, public  :: norm2 => norm
     procedure, public  :: normp2
     procedure, public  :: normalized
     procedure, public  :: conjg => conjugate
     procedure, public  :: real => scalar
     procedure, public  :: aimag => vectorial_vec

     procedure, public  :: inv => inverse
     procedure, public  :: exp => exp_q
     procedure, public  :: log => log_q
     procedure, public  :: sqrt => sqrt_q

  end type quat

  interface quat
     module procedure init_
     module procedure init_f
     module procedure init_s
     module procedure init_v3
     module procedure init_f4
  end interface quat

  interface assignment (=)
     module procedure assign_q
     module procedure assign_s
     module procedure assign_f
  end interface assignment (=)

  interface abs
    procedure norm
 end interface abs

  interface conjg
     module procedure conjugate
  end interface conjg

  interface real
    module procedure scalar
  end interface real

  interface aimag
    module procedure vectorial_vec
  end interface aimag

  interface exp
    module procedure exp_q
  end interface exp

  interface log
    module procedure log_q
  end interface log

  interface sqrt
    module procedure sqrt_q
 end interface sqrt

  interface inv
    module procedure inverse
 end interface inv

  interface norm2
    module procedure norm
 end interface norm2

contains

  !-----------------------------------------------------------------------
  ! intrinsic functions
  !-----------------------------------------------------------------------

  elemental pure function exp_q(q) result(r)
    class(quat), intent(in) :: &
         q
    type(quat) :: &
         r

    real(kind=real64) :: &
         t, s, c

    ! t = norm2(q%f(1:3))
    t = sqrt(q%f(1)**2 + q%f(2)**2 + q%f(3)**2)

    if (t < lim_prec) then
       s = 1.d0
       c = 1.d0
    else
       s = sin(t)/t
       c = cos(t)
    endif


    r%f(0) = c
    r%f(1:3) = q%f(1:3) * s

    r%f = r%f * exp(q%f(0))

  end function exp_q


  elemental pure function log_q(q) result(r)
    class(quat), intent(in) :: &
         q
    type(quat) :: &
         r

    real(kind=real64) :: &
         t, u, a

    !u = norm2(q%f)
    !t = norm2(q%f(1:3))

    u = sqrt(q%f(0)**2 + q%f(1)**2 + q%f(2)**2 + q%f(3)**2)
    t = sqrt(q%f(1)**2 + q%f(2)**2 + q%f(3)**2)

    if (u == 0.d0) then
       r%f = ieee_value(1.d0,ieee_signaling_nan)
       return
    endif

    r%f(0) = log(u)
    a = q%f(0) / u
    if ( t < lim_exp) then
       r%f(1:3) = 0.d0
    else
       r%f(1:3) = q%f(1:3) * (acos(a) / t)
    end if

  end function log_q


  elemental pure function sqrt_q(q) result(r)
    class(quat), intent(in) :: &
         q
    type(quat) :: &
         r

    r =  q ** (0.5d0)

  end function sqrt_q


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

  pure function init_f(f) result(r)
    real(kind=real64), dimension(0:), intent(in) :: &
         f
    type(quat) :: &
         r

    integer :: &
         n

    n = size(f)

    if (n == 4) then
       r%f    = f
    else if (n == 3) then
       r%f(0)   = 0.d0
       r%f(1:3) = f(0:2)
    else
       error stop '[init_f] invalid quad size'
    endif

  end function init_f


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

    r%f(0)   = s
    r%f(1:3) = 0.d0
  end function init_s


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

    r%f(0) = 0.d0
    r%f(1) = v1
    r%f(2) = v2
    r%f(3) = v3
  end function init_v3


  function init_f4(f0, f1, f2, f3) result(r)
    real(kind=real64), intent(in) :: &
         f0, f1, f2, f3
    type(quat) :: &
         r

    r%f(0) = f0
    r%f(1) = f1
    r%f(2) = f2
    r%f(3) = f3
  end function init_f4


  function init_() result(r)
    type(quat) :: &
         r

    r%f(:) = 0.d0
  end function init_


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

  elemental pure subroutine assign_s(r, s)
    type(quat), intent(out) :: &
         r
    real(kind=real64), intent(in) :: &
         s
    r%f(0) = s
    r%f(1:3) = 0.d0
  end subroutine assign_s


  pure subroutine assign_f(r, f)
    type(quat), intent(out) :: &
         r
    real(kind=real64), dimension(0:), intent(in) :: &
         f

    integer :: &
         n

    n = size(f)

    if (n == 4) then
       r%f = f
    else if (n == 3) then
       r%f(0)   = 0.d0
       r%f(1:3) = f(0:2)
    else
       error stop '[assign_f] invalid quad size'
    endif

  end subroutine assign_f


  elemental pure subroutine assign_q(r, p)
    type(quat), intent(out) :: &
         r
    type(quat), intent(in) :: &
         p

    r%f = p%f
  end subroutine assign_q

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

  elemental pure function conjugate(q) result(r)
    class(quat), intent(in) :: &
         q
    type(quat) :: &
         r

    r%f(0)   =  q%f(0)
    r%f(1:3) = -q%f(1:3)

  end function conjugate

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

  elemental pure function norm(q) result(s)
    class(quat), intent(in) :: &
         q
    real(kind=real64) :: &
         s

    ! s = norm2(q%f)
    s = sqrt(q%f(0)**2 + q%f(1)**2 + q%f(2)**2 + q%f(3)**2)

  end function norm

  elemental pure function normp2(q) result(s)
    class(quat), intent(in) :: &
         q
    real(kind=real64) :: &
         s

    ! s = dot_product(q%f, q%f)
    ! s = sum(q%f**2)
    s = q%f(0)**2 + q%f(1)**2 + q%f(2)**2 + q%f(3)**2

  end function normp2

  elemental pure function normalized(q) result(r)
    class(quat), intent(in) :: &
         q
    type(quat) :: &
         r

    real(kind=real64) :: &
         u

    ! u = norm2(q%f)
    u = sqrt(q%f(0)**2 + q%f(1)**2 + q%f(2)**2 + q%f(3)**2)

    if (u == 0.d0) then
       r%f = 0.d0
       return
    endif

    r%f = q%f / u

  end function normalized

  !-----------------------------------------------------------------------
  ! operator real function
  !-----------------------------------------------------------------------

  elemental pure function scalar(q) result(s)
    class(quat), intent(in) :: &
         q
    real(kind=real64) :: &
         s

    s = q%f(0)

  end function scalar

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

  pure function vectorial(q) result(v)
    class(quat), intent(in) :: &
         q
    real(kind=real64), dimension(3) :: &
         v

    v(1:3) = q%f(1:3)

  end function vectorial

  elemental pure function vectorial_vec(q) result(v)
    class(quat), intent(in) :: &
         q
    type(vec) :: &
         v

    v%v(1:3) = q%f(1:3)

  end function vectorial_vec

  !-----------------------------------------------------------------------
  ! inverse function
  !-----------------------------------------------------------------------

  elemental pure function inverse(p) result(r)
    class(quat), intent(in) :: &
         p
    type(quat) :: &
         r

    real(kind=real64) :: &
         u2

    r%f(0)   =  p%f(0)
    r%f(1:3) = -p%f(1:3)

    ! u2 = dot_product(p%f, p%f)
    ! u2 = sum(p%f**2)
    u2 = p%f(0)**2 + p%f(1)**2 + p%f(2)**2 + p%f(3)**2

    if (u2 == 0.d0) then
       r%f = ieee_value(1.d0,ieee_signaling_nan)
       return
    endif

    r%f = r%f / u2

  end function inverse

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

  elemental pure function add_qp(q, p) result(r)
    class(quat), intent(in) :: &
         q, p
    type(quat) :: &
         r

    r%f = q%f + p%f
  end function add_qp

  elemental pure function add_qs(q, s) result(r)
    class(quat), intent(in) :: &
         q
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    r%f(0)   = q%f(0) + s
    r%f(1:3) = q%f(1:3)

  end function add_qs

  elemental pure function add_sp(s, p) result(r)
    class(quat), intent(in) :: &
         p
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    r%f(0)   = p%f(0) + s
    r%f(1:3) = p%f(1:3)

  end function add_sp

  elemental pure function pos(q) result(r)
    class(quat), intent(in) :: &
         q
    type(quat) :: &
         r

    r%f = q%f

  end function pos

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

  elemental pure function sub_qp(q, p) result(r)
    class(quat), intent(in) :: &
         q, p
    type(quat) :: &
         r

    r%f = q%f - p%f

  end function sub_qp

  elemental pure function sub_qs(q, s) result(r)
    class(quat), intent(in) :: &
         q
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    r%f(0)   = q%f(0) - s
    r%f(1:3) = q%f(1:3)

  end function sub_qs

  elemental pure function sub_sp(s, p) result(r)
    class(quat), intent(in) :: &
         p
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    r%f(0)   = s - p%f(0)
    r%f(1:3) =     p%f(1:3)

  end function sub_sp

  elemental pure function neg(q) result(r)
    class(quat), intent(in) :: &
         q
    type(quat) :: &
         r

    r%f = -q%f

  end function neg

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

  elemental pure function mul_qs(q, s) result(r)
    class(quat), intent(in) :: &
         q
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    r%f = q%f * s

  end function mul_qs

  elemental pure function mul_sp(s, p) result(r)
    class(quat), intent(in) :: &
         p
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    r%f = p%f * s

  end function mul_sp

  elemental pure function mul_qp(q, p) result(r)
    class(quat), intent(in) :: &
         q, p
    type(quat) :: &
         r

    ! r%f(0) = q%f(0) * p%f(0) - dot_product(q%f(1:3), p%f(1:3))
    ! r%f(1:3) = q%f(0) * p%f(1:3) + p%f(0) * q%f(1:3) + &
    !      cross_product(q%f(1:3), p%f(1:3))

    r%f(0) = q%f(0) * p%f(0) - (q%f(1) * p%f(1) + q%f(2) * p%f(2) + q%f(3) * p%f(3))
    ! r%f(1:3) = q%f(0) * p%f(1:3) + p%f(0) * q%f(1:3) + &
    !      (/q%f(2) * p%f(3) - q%f(3) * p%f(2), &
    !        q%f(3) * p%f(1) - q%f(1) * p%f(3), &
    !        q%f(1) * p%f(2) - q%f(2) * p%f(1)/)

    r%f(1:3) = q%f(0) * p%f(1:3) + p%f(0) * q%f(1:3)
    r%f(1) = r%f(1) + q%f(2) * p%f(3) - q%f(3) * p%f(2)
    r%f(2) = r%f(2) + q%f(3) * p%f(1) - q%f(1) * p%f(3)
    r%f(3) = r%f(3) + q%f(1) * p%f(2) - q%f(2) * p%f(1)

    ! r%f(1) = q%f(0) * p%f(1) + p%f(0) * q%f(1) + q%f(2) * p%f(3) - q%f(3) * p%f(2)
    ! r%f(2) = q%f(0) * p%f(2) + p%f(0) * q%f(2) + q%f(3) * p%f(1) - q%f(1) * p%f(3)
    ! r%f(3) = q%f(0) * p%f(3) + p%f(0) * q%f(3) + q%f(1) * p%f(2) - q%f(2) * p%f(1)

  end function mul_qp

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

  elemental pure function div_qs(q, s) result(r)
    class(quat), intent(in) :: &
         q
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    r%f = q%f * (1.d0 / s)

  end function div_qs


  elemental pure function div_sp(s, p) result(r)
    class(quat), intent(in) :: &
         p
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    real(kind=real64) :: &
         u2

    r%f(0)   =  p%f(0)
    r%f(1:3) = -p%f(1:3)

    ! u2 = dot_product(p%f, p%f)
    ! u2 = sum(p%f**2)
    u2 = p%f(0)**2 + p%f(1)**2 + p%f(2)**2 + p%f(3)**2

    if (u2 == 0.d0) then
       r%f = ieee_value(1.d0,ieee_signaling_nan)
       return
    endif

    r%f = r%f * (s / u2)

  end function div_sp


  elemental pure function div_qp(q, p) result(r)
    class(quat), intent(in) :: &
         q, p
    type(quat) :: &
         r

    real(kind=real64) :: &
         t

    ! t = dot_product(p%f, p%f)
    ! t = sum(p%f**2)
    t = p%f(0)**2 + p%f(1)**2 + p%f(2)**2 + p%f(3)**2

    if (t == 0.d0) then
       r%f = ieee_value(1.d0,ieee_signaling_nan)
       return
    endif

    t = 1.d0 / t

    ! r%f(0)   =  p%f(0)
    ! r%f(1:3) = -p%f(1:3)
    ! r = (r * q) * t

    ! r%f(0) = (q%f(0) * p%f(0) + dot_product(q%f(1:3), p%f(1:3))) * t
    ! r%f(1:3) = t * &
    !       (-cross_product(q%f(1:3), p%f(1:3)) &
    !        -q%f(0) * p%f(1:3) + p%f(0) * q%f(1:3))

    r%f(0) = q%f(0) * p%f(0) + q%f(1) * p%f(1) + q%f(2) * p%f(2) + q%f(3) * p%f(3)
    r%f(1:3) =  p%f(0) * q%f(1:3) - q%f(0) * p%f(1:3)
    r%f(1) = r%f(1) + q%f(3) * p%f(2) - q%f(2) * p%f(3)
    r%f(2) = r%f(2) + q%f(1) * p%f(3) - q%f(3) * p%f(1)
    r%f(3) = r%f(3) + q%f(2) * p%f(1) - q%f(1) * p%f(2)
    r%f = r%f * t

  end function div_qp

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

  elemental pure function eq_qs(q, s) result(l)
    class(quat), intent(in) :: &
         q
    real(kind=real64), intent(in) :: &
         s
    logical :: &
         l

    l = (s == q%f(0)) .and. all(q%f(1:3) == 0.d0)

  end function eq_qs


  elemental pure function eq_sp(s, p) result(l)
    class(quat), intent(in) :: &
         p
    real(kind=real64), intent(in) :: &
         s
    logical :: &
         l

    l = (s == p%f(0)) .and. all(p%f(1:3) == 0.d0)

  end function eq_sp


  elemental pure function eq_qp(q, p) result(l)
    class(quat), intent(in) :: &
         q, p
    logical :: &
         l

    l = all(q%f == p%f)

  end function eq_qp

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

  elemental pure function ne_qs(q, s) result(l)
    class(quat), intent(in) :: &
         q
    real(kind=real64), intent(in) :: &
         s
    logical :: &
         l

    l = (s /= q%f(0)) .or. any(q%f(1:3) /= 0.d0)

  end function ne_qs


  elemental pure function ne_sp(s, p) result(l)
    class(quat), intent(in) :: &
         p
    real(kind=real64), intent(in) :: &
         s
    logical :: &
         l

    l = (s /= p%f(0)) .or. any(p%f(1:3) /= 0.d0)

  end function ne_sp


  elemental pure function ne_qp(q, p) result(l)
    class(quat), intent(in) :: &
         q, p
    logical :: &
         l

    l = any(q%f /= p%f)

  end function ne_qp

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

  elemental pure function pow_qi(q, i) result(r)
    class(quat), intent(in) :: &
         q
    integer, intent(in) :: &
         i
    type(quat) :: &
         r

    select case(i)
       case(0)
          r%f = (/1.d0, 0.d0, 0.d0, 0.d0/)
       case(1)
          r = q
       case (2)
          r = q * q
       case (3)
          r = q * q * q
       case (4)
          r = q * q
          r = r * r
       case (5)
          r = q * q
          r = r * r * q
       case (6)
          r = q * q
          r = r * r * r
       case (-1)
          r = q%inv()
       case (-2)
          r = q%inv()
          r = r * r
       case (-3)
          r = q%inv()
          r = r * r * r
       case (-4)
          r = q%inv()
          r = r * r
          r = r * r
       case (-6)
          r = q%inv()
          r = r * r
          r = r * r * r
       case default
          r = pow_qs(q, dble(i))
       end select

  end function pow_qi


  elemental pure function pow_sp(s, p) result(r)
    class(quat), intent(in) :: &
         p
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    r = exp(log(s) * p)

  end function pow_sp


  elemental pure function pow_qs_(q, s) result(r)
    class(quat), intent(in) :: &
         q
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    r = exp(log(q) * s)

  end function pow_qs_


  elemental pure function pow_qs(q, s) result(r)
    class(quat), intent(in) :: &
         q
    real(kind=real64), intent(in) :: &
         s
    type(quat) :: &
         r

    real(kind=real64) :: &
         u,t,a,p

    if (s == 0) then
       r%f = (/1.d0, 0.d0, 0.d0, 0.d0/)
       return
    endif

    ! u = norm2(q%f)
    ! t = norm2(q%f(1:3))

    u = sqrt(q%f(0)**2 + q%f(1)**2 + q%f(2)**2 + q%f(3)**2)
    t = sqrt(q%f(1)**2 + q%f(2)**2 + q%f(3)**2)

    if (u == 0.d0) then
       if (s < 0.d0) then
          r%f = ieee_value(1.d0,ieee_signaling_nan)
          return
       endif
       a = 0.d0
    else
       a = exp(log(u) * s)
    endif
    if (t > lim_exp) then
       r%f(1:3) = q%f(1:3) / t
    else
       r%f(1:3) = 0.d0
    endif

    p = atan2(t, q%f(0)) * s
    r%f(0)   =  a * cos(p)
    r%f(1:3) = (a * sin(p)) * r%f(1:3)

  end function pow_qs


  elemental pure function pow_qp(q, p) result(r)
    class(quat), intent(in) :: &
         q, p
    type(quat) :: &
         r

    r = exp(log(q) * p)

  end function pow_qp

end module quaternion