compile Module Subroutine

  module subroutine compile( &
       this, optimiser, loss_method, accuracy_method, &
       metrics, batch_size, verbose &
  )
    !! Compile the network
    implicit none

    ! Arguments
    class(network_type), intent(inout) :: this
    !! Instance of network
    class(base_optimiser_type), optional, intent(in) :: optimiser
    !! Optimiser to use for training
    class(*), optional, intent(in) :: loss_method
    !! Loss method
    character(*), optional, intent(in) :: accuracy_method
    !! Accuracy method
    class(*), dimension(..), optional, intent(in) :: metrics
    !! Metrics
    integer, optional, intent(in) :: batch_size
    !! Batch size
    integer, optional, intent(in) :: verbose
    !! Verbosity level

    ! Local variables
    integer :: i, j, k, child_id, parent_id, layer_id, num_inputs, input_rank
    !! Loop index
    integer :: parent_vertex, vertex_idx
    !! Vertex indices
    integer :: layer_rank, parent_rank, operator
    !! Ranks of layers
    integer :: verbose_ = 0
    !! Verbosity level
    logical :: use_graph_input = .false.
    !! Boolean whether to use graph input
    logical :: l_flatten_child, l_set_input_shape
    !! Booleans whether to flatten child or set input shape
    integer, dimension(:), allocatable :: input_shape, &
         child_vertices, parent_vertices, output_ranks, parent_ids
    !! Shapes of the input and output of the layers
    integer, dimension(:,:), allocatable :: merge_shape
    !! Shapes of the inputs to merge layers
    class(base_layer_type), allocatable :: &
         t_input_layer, t_flatten_layer, t_merge_layer
    !! Temporary input, flatten, and merge layers


    !---------------------------------------------------------------------------
    ! Initialise optional arguments
    !---------------------------------------------------------------------------
    if(present(verbose)) verbose_ = verbose


    !---------------------------------------------------------------------------
    ! Initialise metrics
    !---------------------------------------------------------------------------
    if(present(metrics)) call this%set_metrics(metrics)


    !---------------------------------------------------------------------------
    ! Initialise loss and accuracy methods
    !---------------------------------------------------------------------------
    if(present(loss_method)) call this%set_loss(loss_method, verbose_)
    if(present(accuracy_method)) &
         call this%set_accuracy(accuracy_method, verbose_)


    !---------------------------------------------------------------------------
    ! Check for input layers at root vertices
    !---------------------------------------------------------------------------
    this%auto_graph%directed = .true.
    call this%build_root_vertices()
    do i = 1, size(this%root_vertices)
       layer_id = this%auto_graph%vertex(this%root_vertices(i))%id
       if(.not.allocated(this%model(layer_id)%layer%input_shape))then
          call stop_program("input_shape of first layer not defined")
          return
       end if
       use_graph_input = .false.
       select type( root => this%model(layer_id)%layer)
       class is(input_layer_type)
          cycle
       class is(learnable_layer_type)
          input_shape = root%input_shape
          use_graph_input = root%use_graph_input
       class default
          input_shape = root%input_shape
       end select
       t_input_layer = input_layer_type(&
            input_shape = input_shape, &
            index = i, &
            use_graph_input = use_graph_input, &
            verbose=verbose_ &
       )
       call this%add( &
            t_input_layer, output_list = [ this%model(layer_id)%layer%id ] &
       )
       ! NEED TO CALL layer%init?
       deallocate(input_shape)
       deallocate(t_input_layer)
       this%root_vertices(i) = this%num_layers
       if(i.eq.1)then
          do j = 1, this%auto_graph%num_edges
             if(this%auto_graph%edge(j)%index(1).eq.0) &
                  this%auto_graph%edge(j)%index(1) = this%num_layers
          end do
       end if
    end do
    call this%auto_graph%generate_adjacency()


    !---------------------------------------------------------------------------
    ! Identify whether input is graph type
    !---------------------------------------------------------------------------
    if( &
         this%model( &
              this%auto_graph%vertex(this%root_vertices(1))%id &
         )%layer%use_graph_input &
    )then
       this%use_graph_input = .true.
    else
       this%use_graph_input = .false.
    end if


    !---------------------------------------------------------------------------
    ! Check for zero input rank layers
    !---------------------------------------------------------------------------
    do i = 1, size(this%auto_graph%vertex, dim = 1)
       layer_id = this%auto_graph%vertex(i)%id
       if(this%model(layer_id)%layer%input_rank.eq.0)then
          parent_ids = pack( &
               [ ( &
                    this%auto_graph%vertex(j)%id, &
                    j = 1, size(this%auto_graph%adjacency(:,i)) &
               ) ], &
               this%auto_graph%adjacency(:,i) .ne. 0 &
          )
          if(size(parent_ids).eq.0) cycle
          output_ranks = [ ( this%model(parent_ids(j))%layer%output_rank, &
               j=1,size(parent_ids) ) ]
          if(any(output_ranks.ne.output_ranks(1)))then
             write(0,*) output_ranks
             call stop_program( &
                  "input rank of layer "//trim(this%model(layer_id)%layer%name) // &
                  " is zero, but multiple parents with different output ranks" &
             )
             return
          end if
          input_rank = this%model(parent_ids(1))%layer%output_rank
          call this%model(layer_id)%layer%set_rank( &
               input_rank = input_rank, &
               output_rank = input_rank &
          )
       end if
    end do


    !---------------------------------------------------------------------------
    ! Check for required flatten layers
    !---------------------------------------------------------------------------
    i = 0
    flatten_loop: do
       i = i + 1
       if(i.gt.this%auto_graph%num_vertices) exit flatten_loop
       layer_id = this%auto_graph%vertex(i)%id

       ! get all child vertices
       child_vertices = pack( &
            [(j, j=1,size(this%auto_graph%adjacency(i,:)))], &
            this%auto_graph%adjacency(i,:) .ne. 0 &
       )
       child_loop: do j = 1, size(child_vertices)
          ! Get layer ID (needed for add() function's output_list parameter)
          child_id = this%auto_graph%vertex(child_vertices(j))%id
          if(trim(this%model(layer_id)%layer%type).eq."flat") cycle child_loop
          if( this%model(layer_id)%layer%output_rank .eq. &
               this%model(child_id)%layer%input_rank ) cycle child_loop
          if(this%model(layer_id)%layer%output_rank.eq.0) cycle child_loop

          ! get all parent vertices of the child vertex
          parent_vertices = pack( &
               [(k, k=1,size(this%auto_graph%adjacency(:,child_vertices(j))))], &
               this%auto_graph%adjacency(:,child_vertices(j)) .ne. 0 &
          )
          l_flatten_child = .true.
          do k = 1, size(parent_vertices)
             parent_id = this%auto_graph%vertex(parent_vertices(k))%id
             !check if ranks match, rather than input and output shapes
             if( this%model(layer_id)%layer%output_rank .ne. &
                  this%model(parent_id)%layer%input_rank &
             ) l_flatten_child = .false.
          end do
          t_flatten_layer = flatten_layer_type( &
               input_rank = this%model(layer_id)%layer%output_rank &
          )

          if(l_flatten_child)then
             ! add flatten layer in the place of the child layer
             operator = this%auto_graph%edge( &
                  this%auto_graph%adjacency(parent_vertices(1),child_vertices(j)) &
             )%id
             call this%auto_graph%remove_edges( &
                  indices = [ &
                       this%auto_graph%adjacency( &
                            parent_vertices(:),child_vertices(j) &
                       ) &
                  ] &
             )
             call this%add( &
                  t_flatten_layer, &
                  input_list=[parent_vertices(:)], output_list=[child_id], &
                  operator=operator &
             )
          else
             ! add flatten layer between the current layer and the child layer
             operator = this%auto_graph%edge( &
                  this%auto_graph%adjacency(i,child_vertices(j)) &
             )%id
             call this%auto_graph%remove_edges( &
                  indices = [this%auto_graph%adjacency(i,child_vertices(j))] &
             )
             call this%add( &
                  t_flatten_layer, input_list = [layer_id], &
                  output_list = [child_id], &
                  operator=operator &
             )
          end if
          deallocate(t_flatten_layer)
          deallocate(child_vertices)
          cycle flatten_loop
       end do child_loop
       deallocate(child_vertices)
    end do flatten_loop
    call this%build_vertex_order()


    !---------------------------------------------------------------------------
    ! Check for required merge layers
    !---------------------------------------------------------------------------
    i = 0
    merge_loop: do
       i = i + 1
       if(i.gt.this%auto_graph%num_vertices) exit merge_loop
       layer_id = this%auto_graph%vertex(i)%id
       if(this%model(layer_id)%layer%type.eq."merg") cycle merge_loop

       ! get all child vertices
       parent_vertices = pack( &
            [(j, j=1,size(this%auto_graph%adjacency(:,i)))], &
            this%auto_graph%adjacency(:,i) .ne. 0 &
       )
       if(size(parent_vertices).le.1) cycle merge_loop

       ! get edge id for merge layer
       operator = this%auto_graph%edge( &
            this%auto_graph%adjacency(parent_vertices(1),i) &
       )%id

       ! remove edges from parents to this layer
       do j = 1, size(parent_vertices)
          call this%auto_graph%remove_edges( &
               indices = [this%auto_graph%adjacency(parent_vertices(j),i)] &
          )
       end do
       parent_ids = &
            [ ( &
                 this%auto_graph%vertex(parent_vertices(k))%id, &
                 k = 1, size(parent_vertices) &
            ) ]
       select case(operator)
       case(1) ! concatenate
          t_merge_layer = concat_layer_type( &
               input_layer_ids = parent_ids, &
               input_rank = this%model(layer_id)%layer%input_rank &
          )
       case(2) ! add
          t_merge_layer = add_layer_type( &
               input_layer_ids = parent_ids, &
               input_rank = this%model(layer_id)%layer%input_rank &
          )
          ! case(3) ! multiply
          !    t_merge_layer = multiply_layer_type( &
          !         input_layer_ids = parent_vertices &
          !    )
       case default
          write(0,*) "invalid merge operator: ", operator
          call stop_program("invalid merge operator")
          return
       end select
       t_merge_layer%use_graph_input = this%model(layer_id)%layer%use_graph_input
       t_merge_layer%use_graph_output = t_merge_layer%use_graph_input
       call this%add( &
            t_merge_layer, &
            input_list = parent_ids, &
            output_list = [layer_id] &
       )
       deallocate(t_merge_layer)
    end do merge_loop
    call this%build_vertex_order()


    ! Update number of layers
    !---------------------------------------------------------------------------
    this%num_layers = size(this%model,dim=1)



    !---------------------------------------------------------------------------
    ! Initialise layers
    !---------------------------------------------------------------------------
    do i = 1, size(this%vertex_order, dim = 1)
       vertex_idx = this%vertex_order(i)
       layer_id = this%auto_graph%vertex(vertex_idx)%id
       if(allocated(this%model(layer_id)%layer%input_shape))then
          l_set_input_shape = .false.
       else
          l_set_input_shape = .true.
       end if
       if(l_set_input_shape)then
          layer_rank = this%model(layer_id)%layer%input_rank
          parent_rank = 0

          select type( layer => this%model(layer_id)%layer )
          class is(merge_layer_type)
             ! loop over all parent layers
             allocate( &
                  merge_shape( &
                       this%model(layer_id)%layer%input_rank, &
                       size(layer%input_layer_ids) &
                  ) &
             )
             do k = 1, size(layer%input_layer_ids)
                merge_shape(:,k) = &
                     this%model(layer%input_layer_ids(k))%layer%output_shape
             end do
             input_shape = layer%calc_input_shape(merge_shape)
             deallocate(merge_shape)
          class default

             allocate( &
                  input_shape(this%model(layer_id)%layer%input_rank), &
                  source = 0 &
             )
             do j = 1, this%auto_graph%num_vertices
                if(this%auto_graph%adjacency(j,vertex_idx).eq.0) cycle
                parent_id = this%auto_graph%vertex(j)%id
                parent_rank = this%model(parent_id)%layer%output_rank

                if(layer_rank .eq. parent_rank)then
                   input_shape(:) = input_shape(:) + &
                        this%model(parent_id)%layer%output_shape
                elseif(layer_rank .eq. 1)then
                   input_shape(1) = input_shape(1) + &
                        product( this%model(parent_id)%layer%output_shape )
                end if
             end do
          end select
          call this%model(layer_id)%layer%init( &
               input_shape = input_shape, &
               verbose = verbose_ &
          )
          deallocate(input_shape)
       end if
       if(verbose_.gt.0)then
          write(*,*) "layer: ", layer_id, this%model(layer_id)%layer%type
          write(*,*) this%model(layer_id)%layer%input_shape
          write(*,*) this%model(layer_id)%layer%output_shape
       end if
    end do


    ! Set number of outputs
    !---------------------------------------------------------------------------
    this%num_outputs = 0
    call this%build_leaf_vertices()
    do i = 1, size(this%leaf_vertices,1)
       this%num_outputs = this%num_outputs + &
            product( &
                 this%model( &
                      this%auto_graph%vertex(this%leaf_vertices(i))%id &
                 )%layer%output_shape &
            )
    end do
    if( &
         this%model( &
              this%auto_graph%vertex(this%leaf_vertices(1))%id &
         )%layer%use_graph_output &
    )then
       this%use_graph_output = .true.
    else
       this%use_graph_output = .false.
    end if


    !---------------------------------------------------------------------------
    ! Confirm input_shape of each layer matches data going into it
    !---------------------------------------------------------------------------
    do i = 1, size(this%vertex_order, dim = 1)
       vertex_idx = this%vertex_order(i)
       layer_id = this%auto_graph%vertex(vertex_idx)%id
       if(this%model(layer_id)%layer%type.eq."inpt") cycle

       ! Get all parent vertices that feed into this layer
       parent_vertices = pack( &
            [(j, j=1,size(this%auto_graph%adjacency(:,vertex_idx)))], &
            this%auto_graph%adjacency(:,vertex_idx) .ne. 0 &
       )
       if(size(parent_vertices).eq.0) cycle
       select type( layer => this%model(layer_id)%layer )
       class is(merge_layer_type)
          operator = layer%merge_mode
       class default
          if(size(parent_vertices).gt.1)then
             call stop_program( &
                  "layer "//trim(layer%name)// &
                  " is not a merge layer but has multiple inputs" &
             )
             return
          end if
       end select

       ! Calculate expected input size from parent layers
       num_inputs = 0
       do j = 1, size(parent_vertices)
          parent_vertex = parent_vertices(j)

          select case(operator)
          case(1) ! pointwise - all inputs should have same size
             if(num_inputs.eq.0)then
                if(this%model(layer_id)%layer%use_graph_input)then
                   num_inputs = this%model(parent_vertex)%layer%output_shape(1)
                else
                   num_inputs = product(this%model(parent_vertex)%layer%output_shape)
                end if
             end if
          case(2) ! concatenate
             if(this%model(layer_id)%layer%use_graph_input)then
                num_inputs = num_inputs + &
                     this%model(parent_vertex)%layer%output_shape(1)
             else
                num_inputs = num_inputs + &
                     product(this%model(parent_vertex)%layer%output_shape)
             end if
          end select
       end do

       ! Verify calculated input size matches layer's expected input size
       if(this%model(layer_id)%layer%use_graph_input)then
          if(num_inputs.ne.this%model(layer_id)%layer%input_shape(1) .and. &
               num_inputs.ne.0)then
             write(*,*) "Expected:", num_inputs, "Got:", &
                  this%model(layer_id)%layer%input_shape(1)
             call stop_program( &
                  "input_shape of layer "//&
                  trim(this%model(layer_id)%layer%name)// &
                  " does not match data going into it" &
             )
          end if
       else
          if(num_inputs.ne.product(this%model(layer_id)%layer%input_shape) .and. &
               num_inputs.ne.0)then
             write(*,*) "Expected:", num_inputs, "Got:", &
                  product(this%model(layer_id)%layer%input_shape)
             call stop_program( &
                  "input_shape of layer "//&
                  trim(this%model(layer_id)%layer%name)// &
                  " does not match data going into it" &
             )
          end if
       end if

    end do

    !---------------------------------------------------------------------------
    ! Initialise optimiser
    !---------------------------------------------------------------------------
    this%num_params = this%get_num_params()
    if(present(optimiser))then
       this%optimiser = optimiser
    end if
    if(.not.allocated(this%optimiser))then
       call stop_program("No optimiser is defined for the network")
       return
    else
       call this%optimiser%init(num_params=this%num_params)
    end if


    !---------------------------------------------------------------------------
    ! Pre-compute forward pass navigation
    !---------------------------------------------------------------------------
    block
      integer :: nv, l_idx, v_idx, lid, parent_v
      nv = size(this%vertex_order, 1)
      if(allocated(this%fwd_layer_id))   deallocate(this%fwd_layer_id)
      if(allocated(this%fwd_num_inputs)) deallocate(this%fwd_num_inputs)
      if(allocated(this%fwd_parent_id))  deallocate(this%fwd_parent_id)
      if(allocated(this%fwd_layer_type)) deallocate(this%fwd_layer_type)
      allocate(this%fwd_layer_id(nv))
      allocate(this%fwd_num_inputs(nv))
      allocate(this%fwd_parent_id(nv))
      allocate(this%fwd_layer_type(nv))
      this%fwd_parent_id = 0
      do l_idx = 1, nv
         v_idx = this%vertex_order(l_idx)
         lid = this%auto_graph%vertex(v_idx)%id
         this%fwd_layer_id(l_idx) = lid
         this%fwd_num_inputs(l_idx) = &
              count(this%auto_graph%adjacency(:,v_idx).gt.0)
         if(this%fwd_num_inputs(l_idx).eq.1)then
            parent_v = maxloc( &
                 this%auto_graph%adjacency(:,v_idx), dim=1)
            this%fwd_parent_id(l_idx) = &
                 this%auto_graph%vertex(parent_v)%id
         end if
         ! Determine layer type: 0=input, 1=merge, 2=default
         select type(layer => this%model(lid)%layer)
         class is(input_layer_type)
            this%fwd_layer_type(l_idx) = 0
         class is(merge_layer_type)
            this%fwd_layer_type(l_idx) = 1
         class default
            this%fwd_layer_type(l_idx) = 2
         end select
      end do
    end block


    !---------------------------------------------------------------------------
    ! Pre-compute parameter segment layout
    !---------------------------------------------------------------------------
    block
      integer :: l_idx, p_idx, seg_count, s_idx, e_idx
      ! First pass: count segments
      seg_count = 0
      do l_idx = 1, this%num_layers
         select type(current => this%model(l_idx)%layer)
         class is(learnable_layer_type)
            seg_count = seg_count + size(current%params)
         end select
      end do
      this%param_num_segments = seg_count
      if(allocated(this%param_seg_layer)) deallocate(this%param_seg_layer)
      if(allocated(this%param_seg_pidx))  deallocate(this%param_seg_pidx)
      if(allocated(this%param_seg_start)) deallocate(this%param_seg_start)
      if(allocated(this%param_seg_end))   deallocate(this%param_seg_end)
      allocate(this%param_seg_layer(seg_count))
      allocate(this%param_seg_pidx(seg_count))
      allocate(this%param_seg_start(seg_count))
      allocate(this%param_seg_end(seg_count))
      ! Second pass: fill layout
      seg_count = 0
      e_idx = 0
      do l_idx = 1, this%num_layers
         select type(current => this%model(l_idx)%layer)
         class is(learnable_layer_type)
            do p_idx = 1, size(current%params)
               seg_count = seg_count + 1
               s_idx = e_idx + 1
               e_idx = e_idx + size(current%params(p_idx)%val, 1)
               this%param_seg_layer(seg_count) = l_idx
               this%param_seg_pidx(seg_count) = p_idx
               this%param_seg_start(seg_count) = s_idx
               this%param_seg_end(seg_count) = e_idx
            end do
         end select
      end do
    end block


    !---------------------------------------------------------------------------
    ! Set batch size, if provided
    !---------------------------------------------------------------------------
    if(present(batch_size)) this%batch_size = batch_size

  end subroutine compile