! ========================================================
! How to define your own operators in F90
! ========================================================

! --------------------------------------------------------
! (1a) Write an ordinary matrix-vector multiply function
! --------------------------------------------------------
  function MatVecMult(A, v) result (w)
   implicit none

   real, dimension(:,:), intent(in) :: A
   real, dimension(:), intent(in) :: v
   real, dimension( SIZE(A,1) ) :: w

   integer :: i, j
   integer :: N

   N = size(v)

   w = 0.0       !! clear whole vector
   DO i = 1, N
      w = w + v(i) * A( :, i )
   END DO
  end function

! --------------------------------------------------------
! (1b) Write another ordinary function: vector * matrix multiply 
! --------------------------------------------------------
  function VecMatMult(v, A) result (w)
   implicit none

   real, dimension(:), intent(in) :: v
   real, dimension(:,:), intent(in) :: A
   real, dimension( SIZE(A,2) ) :: w

   integer :: i, N

   N = size(v)

   w = 0.0       !! clear whole vector
   DO i = 1, N
      w = w + v(i) * A(i,:)
   END DO
  end function

! ================================================================

  PROGRAM Main
   implicit none

! -------------------------------------------------------------------------
! (2) This is how you OVERLOAD the operator *
! -------------------------------------------------------------------------
   interface operator (*)
    function MatVecMult(A, v) result (w)     	!! Function 1
     real, dimension(:,:), intent(in) :: A
     real, dimension(:), intent(in) :: v
     real, dimension( SIZE(A,1) ) :: w
    end function

    function VecMatMult(v, A) result (w)	!! Function 2
     real, dimension(:), intent(in) :: v
     real, dimension(:,:), intent(in) :: A
     real, dimension( SIZE(A,2) ) :: w
    end function
   end interface 

   real, dimension( 3, 3 ) :: A
   real, dimension( 3 ) :: v1, v2
   integer :: i, j

   CALL random_number( A )
   CALL random_number( v1 )

   v2 = A * v1

   DO i = 1, SIZE(A(:, 1))
      WRITE (6, '("[")', ADVANCE="NO")
      DO j = 1, SIZE(A(1, :))
         WRITE (6, '(2X F6.4)', ADVANCE="NO")   A(i,j)
      END DO
      WRITE (6, '("]")', ADVANCE="NO")

      WRITE (6, '(6X "[" F6.4 "]")', ADVANCE="NO")   v1(i)
      WRITE (6, '(6X "[" F6.4 "]")', ADVANCE="NO")   v2(i)
      print *
   END DO
   print *

   v2 = v1 * A

   DO i = 1, SIZE(A(:, 1))
      IF (i == 2) THEN
	 WRITE (6, '("[" F6.4 "  " F6.4 "  " F6.4"]")', ADVANCE="NO")   v1
      ELSE
	 WRITE (6, '("                        ")', ADVANCE="NO")
      END IF

      WRITE (6, '("  [")', ADVANCE="NO")
      DO j = 1, SIZE(A(1, :))
         WRITE (6, '(2X F6.4)', ADVANCE="NO")   A(i,j)
      END DO
      WRITE (6, '("]")', ADVANCE="NO")

      IF (i == 2) THEN
         WRITE (6, '(6X "[" F6.4 "  " F6.4 "  " F6.4"]")', ADVANCE="NO")   v2
      ELSE
         WRITE (6, "(12X)", ADVANCE="NO")
      END IF
      print *
   END DO
   print *

  end program

