program test_duvenaud_msgpass_layer
  use athena, only: &
       duvenaud_msgpass_layer_type, &
       base_layer_type, &
       learnable_layer_type
  use athena__duvenaud_msgpass_layer, only: read_duvenaud_msgpass_layer
  use graphstruc, only: graph_type
  implicit none

  class(base_layer_type), allocatable :: msgpass_layer
  class(base_layer_type), allocatable :: read_layer
  type(graph_type), dimension(1) :: graph
  integer, parameter :: num_vertex_features = 10, num_time_steps = 3, num_outputs = 5
  integer, parameter :: num_edge_features = 1
  integer, parameter :: num_vertices = 6, num_edges = 8
  integer :: unit
  real, allocatable, dimension(:,:) :: vertex_features, output
  real, parameter :: tol = 1.E-7
  logical :: success = .true.

  real, allocatable, dimension(:) :: params
  real, allocatable, dimension(:,:) :: outputs


!!!-----------------------------------------------------------------------------
!!! set up message passing layer
!!!-----------------------------------------------------------------------------
  msgpass_layer = duvenaud_msgpass_layer_type( &
       num_vertex_features = [ num_vertex_features ], &
       num_edge_features = [ num_edge_features ], &
       num_time_steps = num_time_steps, &
       max_vertex_degree = 6, &
       num_outputs = num_outputs &
  )

  !! check layer name
  if(.not. msgpass_layer%name .eq. 'duvenaud')then
     success = .false.
     write(0,*) 'duvenaud layer has wrong name'
  end if

!!!-----------------------------------------------------------------------------
!!! check layer type
!!!-----------------------------------------------------------------------------
  select type(msgpass_layer)
  type is(duvenaud_msgpass_layer_type)
     !! check number of vertex features
     if(msgpass_layer%num_vertex_features(1) .ne. num_vertex_features)then
        success = .false.
        write(0,*) 'duvenaud layer has wrong input features'
     end if

     !! check number of outputs
     if(msgpass_layer%num_outputs .ne. num_outputs)then
        success = .false.
        write(0,*) 'duvenaud layer has wrong output features'
     end if

     !! check number of time steps
     if(msgpass_layer%num_time_steps .ne. num_time_steps)then
        success = .false.
        write(0,*) 'duvenaud layer has wrong number of time steps'
     end if
  end select

!!!-----------------------------------------------------------------------------
!!! set up a simple graph for testing
!!!-----------------------------------------------------------------------------
  ! Create a simple graph with 6 vertices and some edges
  call graph(1)%set_num_vertices(num_vertices, num_vertex_features)
  call graph(1)%set_num_edges(num_edges)

  ! Set up some dummy vertex features
  graph(1)%is_sparse = .true.
  allocate(graph(1)%vertex_features(num_vertex_features, num_vertices))
  graph(1)%vertex_features = 1.0

  ! Set up edge connectivity (creating a simple connected graph)
  block
    integer, allocatable :: index_list(:, :)
    allocate(index_list(2, num_edges))
    index_list(:, 1) = [1, 2]
    index_list(:, 2) = [1, 3]
    index_list(:, 3) = [2, 3]
    index_list(:, 4) = [2, 4]
    index_list(:, 5) = [3, 5]
    index_list(:, 6) = [4, 5]
    index_list(:, 7) = [4, 6]
    index_list(:, 8) = [5, 6]

    ! Generate adjacency matrix
    call graph(1)%generate_adjacency(index_list)
    deallocate(index_list)

    ! Set edge weights
    allocate(graph(1)%edge_weights(num_edges))
    graph(1)%edge_weights = 1.0

    allocate(graph(1)%edge_features(num_edge_features, num_edges), source = 0.0)
  end block

  ! Set the graph for the layer
  call msgpass_layer%set_graph(graph)

  ! ! Initialise the layer
  ! call msgpass_layer%init([num_vertex_features, num_vertices])

!!!-----------------------------------------------------------------------------
!!! check handling of layer parameters, gradients, and outputs
!!!-----------------------------------------------------------------------------
  select type(msgpass_layer)
  class is(learnable_layer_type)
     !! check layer parameter handling
     params = msgpass_layer%get_params()
     if(size(params) .eq. 0)then
        success = .false.
        write(0,*) 'duvenaud layer has wrong number of parameters'
     end if
     params = 1.E0
     call msgpass_layer%set_params(params)
     params = msgpass_layer%get_params()
     if(any(abs(params - 1.E0).gt.tol))then
        success = .false.
        write(0,*) 'duvenaud layer has wrong parameters'
     end if

     !! check layer gradient handling
     params = msgpass_layer%get_gradients()
     if(size(params) .eq. 0)then
        success = .false.
        write(0,*) 'duvenaud layer has wrong number of gradients'
     end if
     params = 1.E0
     call msgpass_layer%set_gradients(params)
     params = msgpass_layer%get_gradients()
     if(any(abs(params - 1.E0).gt.tol))then
        success = .false.
        write(0,*) 'duvenaud layer has wrong gradients'
     end if
     call msgpass_layer%set_gradients(10.E0)
     params = msgpass_layer%get_gradients()
     if(any(abs(params - 10.E0).gt.tol))then
        success = .false.
        write(0,*) 'duvenaud layer has wrong gradients'
     end if
  end select


!!!-----------------------------------------------------------------------------
!!! Test file I/O operations
!!!-----------------------------------------------------------------------------
  write(*,*) "Testing file I/O operations..."

  ! Create a temporary file for testing
  open(newunit=unit, file='test_duvenaud_msgpass_layer.tmp', &
       status='replace', action='write')

  ! Write layer to file
  write(unit,'("DUVENAUD_MSGPASS")')
  call msgpass_layer%print_to_unit(unit)
  write(unit,'("END DUVENAUD_MSGPASS")')
  close(unit)

  ! Read layer from file
  open(newunit=unit, file='test_duvenaud_msgpass_layer.tmp', &
       status='old', action='read')
  read(unit,*) ! Skip first line
  read_layer = read_duvenaud_msgpass_layer(unit)
  close(unit)

  ! Check that read layer has correct properties
  select type(read_layer)
  type is (duvenaud_msgpass_layer_type)
     if (.not. read_layer%name .eq. 'duvenaud') then
        success = .false.
        write(0,*) 'read duvenaud layer has wrong name'
     end if
  class default
     success = .false.
     write(0,*) 'read layer is not duvenaud_msgpass_layer_type'
  end select

  ! Clean up temporary file
  open(newunit=unit, file='test_duvenaud_msgpass_layer.tmp', status='old')
  close(unit, status='delete')


!!!-----------------------------------------------------------------------------
!!! check for any failed tests
!!!-----------------------------------------------------------------------------
  write(*,*) "----------------------------------------"
  if(success)then
     write(*,*) 'test_duvenaud_msgpass_layer passed all tests'
  else
     write(0,*) 'test_duvenaud_msgpass_layer failed one or more tests'
     stop 1
  end if

end program test_duvenaud_msgpass_layer