!===============================================================================
! Module for generic views on arrays and scalars
! Last edited: Oct 31, 2021 (AM)
!
! TODO:
! - Use PCT
!===============================================================================
module moa_view_types
    use iso_c_binding
    use moa_basic_view_types

    implicit none

    private
    public :: moa_view_type, operator(//), size, shape, rank

    !
    ! NOTE: Intel Fortran 2021.04 does not support overloading //
    ! Then a user-defined operator can be used instead:
    !
    ! interface operator(.cat.)
    !     module procedure catenate_array_array, catenate_array_view, catenate_view_array, catenate_view_view
    ! end interface



    !
    ! Type for defining "views" on arrays - to allow for a#b#c: pointers to other views
    !
    type :: moa_view_type
        type(moa_basic_view)         :: left_array
        type(moa_basic_view)         :: right_array
        type(moa_view_type), pointer :: left_view   => null()
        type(moa_view_type), pointer :: right_view  => null()
    contains
         generic           :: elem        => elem_single, elem_ndim
         procedure         :: elem_single => get_elem_single
         procedure         :: elem_ndim   => get_elem_ndim
    end type moa_view_type

    interface operator(//)
        module procedure catenate_array_array, catenate_array_view, catenate_view_array, catenate_view_view
    end interface

    interface size
        module procedure size_view
    end interface

    interface shape
        module procedure shape_view
    end interface

    interface rank
        module procedure rank_view
    end interface
contains
function catenate_array_array( array1, array2 ) result(new_view)
    integer, dimension(..), target, intent(in) :: array1
    integer, dimension(..), target, intent(in) :: array2
    class(moa_view_type), allocatable          :: new_view

    call check_shapes( shape(array1), shape(array2) )

    allocate( new_view )

    call new_view%left_array%point_to(  array1 )
    call new_view%right_array%point_to( array2 )
end function catenate_array_array

function catenate_view_array( view1, array2 ) result(new_view)
    class(moa_view_type), target, intent(in)   :: view1
    integer, dimension(..), target, intent(in) :: array2
    class(moa_view_type), allocatable          :: new_view

    call check_shapes( shape(view1), shape(array2) )

    allocate( new_view )

    new_view%left_view   => view1
    call new_view%right_array%point_to( array2 )
end function catenate_view_array

function catenate_array_view( array1, view2 ) result(new_view)
    integer, dimension(..), target, intent(in) :: array1
    class(moa_view_type), target, intent(in)   :: view2
    class(moa_view_type), allocatable          :: new_view

    call check_shapes( shape(array1), shape(view2) )

    allocate( new_view )

    call new_view%left_array%point_to(  array1 )
    new_view%right_view => view2
end function catenate_array_view

function catenate_view_view( view1, view2 ) result(new_view)
    class(moa_view_type), target, intent(in) :: view1
    class(moa_view_type), target, intent(in) :: view2
    class(moa_view_type), allocatable        :: new_view

    call check_shapes( shape(view1), shape(view2) )

    allocate( new_view )

    new_view%left_view  => view1
    new_view%right_view => view2
end function catenate_view_view

recursive subroutine get_pointer(view, rnk, idx, elem, found)
    class(moa_view_type), intent(inout)  :: view
    integer, intent(in)                  :: rnk
    integer, dimension(:), intent(inout) :: idx
    integer, pointer                     :: elem
    logical, intent(out)                 :: found

    integer                              :: sz
    integer, dimension(rnk)              :: shp

    found = .false.

!!    write(*,*) 'Index, rank: ',idx, ' - ', rnk

    if ( associated(view%left_array%array) ) then
        shp = view%left_array%shp
        sz  = shp(1)
!!        write(*,*) 'left array', shp, '-- ', idx
        if ( idx(1) <= sz ) then
            found =  .true.
            elem  => view%left_array%elem(idx)
        else
            idx(1) = idx(1) - sz
            if ( associated(view%right_array%array) ) then
                shp = view%right_array%shp
                sz  = shp(1)
                if ( idx(1) <= sz ) then
                    found =  .true.
                    elem  => view%right_array%elem(idx)
                else
                    idx(1) = idx(1) - sz
                    return
                endif
            else
                call get_pointer( view%right_view, rnk, idx, elem, found )
            endif
        endif
    else
!!        write(*,*) 'left view'
        call get_pointer(view%left_view, rnk, idx, elem, found )

        if ( .not. found ) then
            if ( associated(view%right_array%array) ) then
                shp = view%right_array%shp
                sz  = shp(1)
!!                write(*,*) 'right array', shp, ' - ', idx
                if ( idx(1) <= sz ) then
                    found =  .true.
                    elem  => view%right_array%elem(idx)
                else
                    !write(*,*) 'Element not found - index: ', idx
                    error stop
                endif
            else
!!                write(*,*) 'right view'
                call get_pointer( view%right_view, rnk, idx, elem, found )
            endif
        endif
    endif
end subroutine get_pointer

function get_elem_single(view, i) result(elem)
    class(moa_view_type), intent(inout) :: view
    integer, intent(in)                 :: i
    integer, pointer                    :: elem

    integer, dimension(1)               :: inew

    inew =  i
    elem => get_elem_ndim( view, inew )
end function

function get_elem_ndim(view, i) result(elem)
    class(moa_view_type), intent(inout) :: view
    integer, dimension(:), intent(in)   :: i
    integer, pointer                    :: elem

    integer                             :: rnk
    integer, dimension(size(i))         :: inew
    logical                             :: found

    rnk  = rank(view)
    inew = i

    call get_pointer( view, rnk, inew, elem, found )
end function

integer function size_view( view )
    class(moa_view_type), intent(in) :: view

    size_view = product(shape(view))
end function size_view

recursive function shape_view( view )
    class(moa_view_type), intent(in)   :: view

    integer, dimension(:), allocatable :: shape_view
    integer, dimension(:), allocatable :: shp

    allocate( shape_view(0) )

    if ( associated(view%left_array%array) ) then
        if ( size(view%right_array%shp) > 0 ) then
            shape_view = view%left_array%shp
        else
            shape_view = [1]
        endif
    elseif ( associated(view%left_view) ) then
        shape_view = shape(view%left_view)
    endif

    if ( associated(view%right_array%array) ) then
        if ( size(view%right_array%shp) > 0 ) then
            shp            = shape(view%right_array)
            shape_view(1)  = shape_view(1) + view%right_array%shp(1)
        else
            shape_view = [shape_view(1) + 1]
        endif
    elseif ( associated(view%right_view) ) then
        shp           = shape(view%right_view)
        shape_view(1) = shape_view(1) + shp(1)
    endif
end function shape_view

recursive integer function rank_view( view )
    type(moa_view_type), intent(in) :: view

    integer                         :: rank_left, rank_right

    if ( associated(view%left_array%array) ) then
        rank_left = size(view%left_array%shp)
    else
        rank_left = rank(view%left_view)
    endif

    if ( associated(view%right_array%array) ) then
        rank_right = size(view%right_array%shp)
    else
        rank_right = rank(view%right_view)
    endif

    rank_view = max( 1, rank_left, rank_right )  ! Hm, not actually enough

end function rank_view

!subroutine check_indices( view, idx )
!    class(moa_view_type), intent(in)   :: view
!    integer, dimension(:), intent(in)  :: idx
!
!    if ( size(idx) /= rank(view) ) then
!        write(*,*) 'Index has the wrong shape - number of dimensions expected:', rank(view)
!        error stop
!    endif
!
!    if ( any( idx < 1 ) .or. any( idx > shape(view) ) ) then
!        write(*,*) 'Index out of bounds - index should be within the shape:', shape(view)
!        error stop
!    endif
!end subroutine check_indices

subroutine check_shapes( shp1, shp2 )
    integer, dimension(:), intent(in) :: shp1
    integer, dimension(:), intent(in) :: shp2

    integer                           :: i, offset1, offset2, first, last

    if ( abs( size(shp1) - size(shp2) ) > 1 ) then
        write(*,*) 'Incorrect combination of arrays - ranks differ more than 1'
        error stop
    endif

    offset1 = 0
    offset2 = 0
    first   = 2
    last    = max( size(shp1), size(shp2) )

    if ( size(shp1) > size(shp2) ) then
        offset2 = -1
    endif

    if ( size(shp1) < size(shp2) ) then
        offset1 = -1
    endif

    do i = first,last
        if ( shp1(i+offset1) /= shp2(i+offset2) ) then
            write(*,*) 'Incorrect combination of arrays - shapes differ after first dimension'
            error stop
        endif
    enddo
end subroutine check_shapes

end module moa_view_types

program test_cat
    use moa_view_types
    implicit none

    type(moa_view_type)        :: x, y, z
    integer, dimension(10)  :: a, b, c
    integer, dimension(2,5) :: aa, bb
    integer                 :: i, j
    integer, pointer        :: el

    integer, dimension(1)  :: shp

    !     1   2   3   4    5    6    7    8    9   10
    a = [ 2,  3,  5,  7,  11,  13,  17,  19,  23,  29]
    b = [-2, -3, -5, -7, -11, -13, -17, -19, -23, -29]
    c = [31, 37, 41, 43,  47,  53,  59,  61,  67,  71]

    x = a // b
    y = x // c

    !j = y%elem(21)

    write(*,*) 'Shape x:', shape(x)
    write(*,*) 'Shape x%left_array:', shape(x%left_array)
    write(*,*) 'Shape x%left_array:', x%left_array%shp

    !
    ! Hm, with x(i), i = 1,..., 30 we should run out of bounds. No error message though!
    !
    do i = 1,30
        !write(*,*) i, y%elem(i)
        el => y%elem(i)
        write(*,*) i, el
    enddo

    !
    ! Scalars do not work yet: checking of the index needs to be repaired
    !

    write(*,*) 'Scalars:'
    i = 12
    j = 34
    z = i // j
    write(*,*) 'Shape:', shape(z)
!!    el => z%elem(1)
!!    write(*,*) el
!!    el => z%elem(2)
!!    write(*,*) el
!!    write(*,*) z%elem(1), z%elem(2)
!!    z%elem(1) = 3
!!    write(*,*) z%elem(1), z%elem(2)

    write(*,*) 'Shape y: ', shape(y)
!!    write(*,*) 'Shape z: ', shape(z)
!!
!!    write(*,*) 'Rank z: ', rank(z)

    ! This is not a good idea: it stores temporary arrays!
    ! z = reshape(a,[2,5]) // reshape(a,[2,5])

    aa = reshape(a,[2,5])
    bb = reshape(b,[2,5])

    z = aa // bb

    write(*,*) 'Rank z:  ', rank(z)
    write(*,*) 'Shape z: ', shape(z)

    el => z%elem([2,2])
    write(*,*) 'z(2,2):  ', el
!    write(*,*) 'z(2,2):  ', z%elem([2,2])
end program test_cat