module athena__fixed_lno_layer
  !! Module containing implementation of a Laplace Neural Operator layer
  !!
  !! This module implements a Laplace Neural Operator (LNO) layer that
  !! approximates an integral kernel operator in the Laplace-transform domain.
  !! It combines a spectral pathway (encode → spectral mixing → decode)
  !! with a local affine bypass:
  !!
  !! \[ \mathbf{v} = \sigma\!\bigl(
  !!    \underbrace{\mathbf{D}\,\mathbf{R}\,\mathbf{E}\,\mathbf{u}}_{\text{spectral}}
  !!  + \underbrace{\mathbf{W}\,\mathbf{u}}_{\text{local}}
  !!  + \mathbf{b}\bigr) \]
  !!
  !! where:
  !!   - \(\mathbf{u} \in \mathbb{R}^{n_{in}}\) is the discretised input
  !!   - \(\mathbf{E} \in \mathbb{R}^{M \times n_{in}}\) is the Laplace
  !!     encoder basis:  \(E_{k,j}=\exp(-s_k\,t_j)\),
  !!     \(s_k = k\pi\), \(t_j = (j{-}1)/(n_{in}{-}1)\)
  !!   - \(\mathbf{R} \in \mathbb{R}^{M \times M}\) are learnable spectral
  !!     mixing weights
  !!   - \(\mathbf{D} \in \mathbb{R}^{n_{out} \times M}\) is the Laplace
  !!     decoder basis:  \(D_{i,k}=\exp(-s_k\,\tau_i)\),
  !!     \(\tau_i = (i{-}1)/(n_{out}{-}1)\)
  !!   - \(\mathbf{W} \in \mathbb{R}^{n_{out} \times n_{in}}\) are the
  !!     local (bypass) weights
  !!   - \(\mathbf{b} \in \mathbb{R}^{n_{out}}\) is the bias
  !!   - \(\sigma\) is the activation function
  !!   - \(M\) = num_modes, the number of Laplace spectral modes
  !!
  !! Number of parameters (learnable):
  !!   \(M^2 + n_{out}\,n_{in}\) without bias,
  !!   \(M^2 + n_{out}\,n_{in} + n_{out}\) with bias.
  use coreutils, only: real32, stop_program, pi
  use athena__base_layer, only: learnable_layer_type, base_layer_type
  use athena__misc_types, only: base_actv_type, base_init_type, &
       onnx_attribute_type, &
       onnx_node_type, onnx_initialiser_type, onnx_tensor_type
  use athena__onnx_nop_utils, only: emit_nop_input_transpose, &
       emit_nop_output_tail, emit_float_initialiser, emit_matrix_initialiser
  use diffstruc, only: array_type, matmul, operator(+)
  implicit none


  private

  public :: fixed_lno_layer_type
  public :: read_fixed_lno_layer


  type, extends(learnable_layer_type) :: fixed_lno_layer_type
     !! Type for a fixed-basis Laplace Neural Operator layer
     integer :: num_inputs
     !! Number of inputs (discretisation points)
     integer :: num_outputs
     !! Number of outputs (discretisation points)
     integer :: num_modes
     !! Number of Laplace spectral modes
     type(array_type) :: encoder_basis
     !! Fixed Laplace encoder basis E [num_modes x num_inputs]
     type(array_type) :: decoder_basis
     !! Fixed Laplace decoder basis D [num_outputs x num_modes]
     type(array_type), dimension(1) :: z
     !! Temporary array for pre-activation values
   contains
     procedure, pass(this) :: get_num_params => get_num_params_fixed_lno
     procedure, pass(this) :: set_hyperparams => set_hyperparams_fixed_lno
     procedure, pass(this) :: init => init_fixed_lno
     procedure, pass(this) :: print_to_unit => print_to_unit_fixed_lno
     procedure, pass(this) :: read => read_fixed_lno

     procedure, pass(this) :: forward => forward_fixed_lno
     procedure, pass(this) :: get_attributes => get_attributes_fixed_lno
     procedure, pass(this) :: emit_onnx_nodes => emit_onnx_nodes_fixed_lno

     final :: finalise_fixed_lno
  end type fixed_lno_layer_type

  interface fixed_lno_layer_type
     module function layer_setup( &
          num_outputs, num_modes, num_inputs, use_bias, &
          activation, &
          kernel_initialiser, bias_initialiser, verbose &
     ) result(layer)
       integer, intent(in) :: num_outputs
       integer, intent(in) :: num_modes
       integer, optional, intent(in) :: num_inputs
       logical, optional, intent(in) :: use_bias
       class(*), optional, intent(in) :: activation
       class(*), optional, intent(in) :: kernel_initialiser, bias_initialiser
       integer, optional, intent(in) :: verbose
       type(fixed_lno_layer_type) :: layer
     end function layer_setup
  end interface fixed_lno_layer_type



contains

!###############################################################################
  subroutine finalise_fixed_lno(this)
    !! Finalise the fixed-basis Laplace neural operator layer
    implicit none

    ! Arguments
    type(fixed_lno_layer_type), intent(inout) :: this
    !! Layer instance to release

    if(allocated(this%input_shape)) deallocate(this%input_shape)
    if(allocated(this%output)) deallocate(this%output)
    if(this%z(1)%allocated) call this%z(1)%deallocate()
    if(this%encoder_basis%allocated) call this%encoder_basis%deallocate()
    if(this%decoder_basis%allocated) call this%decoder_basis%deallocate()

  end subroutine finalise_fixed_lno
!###############################################################################


!###############################################################################
  pure function get_num_params_fixed_lno(this) result(num_params)
    !! Return the number of learnable parameters for the layer
    implicit none

    ! Arguments
    class(fixed_lno_layer_type), intent(in) :: this
    !! Layer instance
    integer :: num_params
    !! Total number of learnable parameters

    ! R: num_modes^2, W: n_out * n_in, b: n_out (optional)
    num_params = this%num_modes * this%num_modes + &
         this%num_outputs * this%num_inputs
    if(this%use_bias) num_params = num_params + this%num_outputs

  end function get_num_params_fixed_lno
!###############################################################################


!###############################################################################
  module function layer_setup( &
       num_outputs, num_modes, num_inputs, &
       use_bias, &
       activation, &
       kernel_initialiser, bias_initialiser, verbose &
  ) result(layer)
    use athena__activation, only: activation_setup
    use athena__initialiser, only: initialiser_setup
    implicit none

    ! Arguments
    integer, intent(in) :: num_outputs
    !! Number of output features
    integer, intent(in) :: num_modes
    !! Number of Laplace spectral modes
    integer, optional, intent(in) :: num_inputs
    !! Number of input features when known at construction time
    logical, optional, intent(in) :: use_bias
    !! Whether to allocate a bias term
    class(*), optional, intent(in) :: activation
    !! Activation function specification
    class(*), optional, intent(in) :: kernel_initialiser, bias_initialiser
    !! Kernel and bias initialiser specifications
    integer, optional, intent(in) :: verbose
    !! Verbosity level

    type(fixed_lno_layer_type) :: layer
    !! Constructed fixed LNO layer

    ! Local variables
    integer :: verbose_ = 0
    !! Effective verbosity level
    logical :: use_bias_ = .true.
    !! Effective bias flag
    class(base_actv_type), allocatable :: activation_
    !! Materialised activation object
    class(base_init_type), allocatable :: kernel_initialiser_, bias_initialiser_
    !! Materialised kernel and bias initialisers

    if(present(verbose)) verbose_ = verbose
    if(present(use_bias)) use_bias_ = use_bias

    if(present(activation))then
       activation_ = activation_setup(activation)
    else
       activation_ = activation_setup("none")
    end if

    if(present(kernel_initialiser))then
       kernel_initialiser_ = initialiser_setup(kernel_initialiser)
    end if
    if(present(bias_initialiser))then
       bias_initialiser_ = initialiser_setup(bias_initialiser)
    end if

    call layer%set_hyperparams( &
         num_outputs = num_outputs, &
         num_modes = num_modes, &
         use_bias = use_bias_, &
         activation = activation_, &
         kernel_initialiser = kernel_initialiser_, &
         bias_initialiser = bias_initialiser_, &
         verbose = verbose_ &
    )

    if(present(num_inputs)) call layer%init(input_shape=[num_inputs])

  end function layer_setup
!###############################################################################


!###############################################################################
  subroutine set_hyperparams_fixed_lno( &
       this, num_outputs, num_modes, &
       use_bias, &
       activation, &
       kernel_initialiser, bias_initialiser, &
       verbose &
  )
    use athena__activation, only: activation_setup
    use athena__initialiser, only: get_default_initialiser, initialiser_setup
    implicit none

    ! Arguments
    class(fixed_lno_layer_type), intent(inout) :: this
    !! Layer instance to configure
    integer, intent(in) :: num_outputs
    !! Number of output features
    integer, intent(in) :: num_modes
    !! Number of Laplace spectral modes
    logical, intent(in) :: use_bias
    !! Whether to use a bias term
    class(base_actv_type), allocatable, intent(in) :: activation
    !! Activation function object
    class(base_init_type), allocatable, intent(in) :: &
         kernel_initialiser, bias_initialiser
    !! Kernel and bias initialiser objects
    integer, optional, intent(in) :: verbose
    !! Verbosity level

    ! Local variables
    character(len=256) :: buffer
    !! Buffer for default initialiser lookup

    this%name = "fixed_lno"
    this%type = "nop"
    this%input_rank = 1
    this%output_rank = 1
    this%use_bias = use_bias
    this%num_outputs = num_outputs
    this%num_modes = num_modes

    if(allocated(this%activation)) deallocate(this%activation)
    if(.not.allocated(activation))then
       this%activation = activation_setup("none")
    else
       allocate(this%activation, source=activation)
    end if
    if(allocated(this%kernel_init)) deallocate(this%kernel_init)
    if(.not.allocated(kernel_initialiser))then
       buffer = get_default_initialiser(this%activation%name)
       this%kernel_init = initialiser_setup(buffer)
    else
       allocate(this%kernel_init, source=kernel_initialiser)
    end if
    if(allocated(this%bias_init)) deallocate(this%bias_init)
    if(.not.allocated(bias_initialiser))then
       buffer = get_default_initialiser( &
            this%activation%name, &
            is_bias=.true. &
       )
       this%bias_init = initialiser_setup(buffer)
    else
       if(allocated(this%bias_init)) deallocate(this%bias_init)
       allocate(this%bias_init, source=bias_initialiser)
    end if

    if(present(verbose))then
       if(abs(verbose).gt.0)then
          write(*,'("fixed_lno activation: ",A)') &
               trim(this%activation%name)
       end if
    end if

  end subroutine set_hyperparams_fixed_lno
!###############################################################################


!###############################################################################
  subroutine init_fixed_lno(this, input_shape, verbose)
    !! Initialise parameter storage, fixed bases and output buffers
    implicit none

    ! Arguments
    class(fixed_lno_layer_type), intent(inout) :: this
    !! Layer instance to initialise
    integer, dimension(:), intent(in) :: input_shape
    !! Input shape used to infer num_inputs
    integer, optional, intent(in) :: verbose
    !! Verbosity level

    ! Local variables
    integer :: num_inputs, j, k, i, idx
    !! Effective fan-in size and basis-construction indices
    integer :: verbose_ = 0
    !! Effective verbosity level
    real(real32) :: s, t
    !! Spectral pole value and normalised coordinate

    if(present(verbose)) verbose_ = verbose

    !---------------------------------------------------------------------------
    ! Set shapes
    !---------------------------------------------------------------------------
    if(.not.allocated(this%input_shape)) call this%set_shape(input_shape)
    this%num_inputs = this%input_shape(1)
    this%output_shape = [this%num_outputs]
    this%num_params = this%get_num_params()


    !---------------------------------------------------------------------------
    ! Allocate learnable parameters
    !
    ! params(1): R  spectral mixing weights [num_modes x num_modes]
    ! params(2): W  local bypass weights    [num_outputs x num_inputs]
    ! params(3): b  bias                    [num_outputs]  (optional)
    !---------------------------------------------------------------------------
    allocate(this%weight_shape(2,2))
    this%weight_shape(:,1) = [ this%num_modes, this%num_modes ]
    this%weight_shape(:,2) = [ this%num_outputs, this%num_inputs ]

    if(this%use_bias)then
       this%bias_shape = [ this%num_outputs ]
       allocate(this%params(3))
    else
       allocate(this%params(2))
    end if

    ! R: spectral mixing weights
    call this%params(1)%allocate([this%num_modes, this%num_modes, 1])
    call this%params(1)%set_requires_grad(.true.)
    this%params(1)%fix_pointer = .true.
    this%params(1)%is_sample_dependent = .false.
    this%params(1)%is_temporary = .false.

    ! W: local bypass weights
    call this%params(2)%allocate([this%num_outputs, this%num_inputs, 1])
    call this%params(2)%set_requires_grad(.true.)
    this%params(2)%fix_pointer = .true.
    this%params(2)%is_sample_dependent = .false.
    this%params(2)%is_temporary = .false.

    num_inputs = this%num_inputs
    if(this%use_bias)then
       num_inputs = this%num_inputs + 1
       call this%params(3)%allocate([this%bias_shape, 1])
       call this%params(3)%set_requires_grad(.true.)
       this%params(3)%fix_pointer = .true.
       this%params(3)%is_sample_dependent = .false.
       this%params(3)%is_temporary = .false.
    end if


    !---------------------------------------------------------------------------
    ! Initialise learnable parameters
    !---------------------------------------------------------------------------
    call this%kernel_init%initialise( &
         this%params(1)%val(:,1), &
         fan_in = this%num_modes, fan_out = this%num_modes, &
         spacing = [ this%num_modes ] &
    )
    call this%kernel_init%initialise( &
         this%params(2)%val(:,1), &
         fan_in = num_inputs, fan_out = this%num_outputs, &
         spacing = [ this%num_outputs ] &
    )
    if(this%use_bias)then
       call this%bias_init%initialise( &
            this%params(3)%val(:,1), &
            fan_in = num_inputs, fan_out = this%num_outputs &
       )
    end if


    !---------------------------------------------------------------------------
    ! Build fixed encoder basis E [num_modes x num_inputs]
    !   E(k,j) = exp(-s_k * t_j)
    !   s_k = k * pi,  t_j = (j-1)/(n_in-1)
    !---------------------------------------------------------------------------
    if(this%encoder_basis%allocated) call this%encoder_basis%deallocate()
    call this%encoder_basis%allocate( &
         [this%num_modes, this%num_inputs, 1])
    this%encoder_basis%is_sample_dependent = .false.
    this%encoder_basis%requires_grad = .false.
    this%encoder_basis%fix_pointer = .true.
    this%encoder_basis%is_temporary = .false.

    do j = 1, this%num_inputs
       if(this%num_inputs .gt. 1)then
          t = real(j-1, real32) / real(this%num_inputs-1, real32)
       else
          t = 0.0_real32
       end if
       do k = 1, this%num_modes
          s = real(k, real32) * pi
          idx = k + (j-1) * this%num_modes
          this%encoder_basis%val(idx, 1) = exp(-s * t)
       end do
    end do


    !---------------------------------------------------------------------------
    ! Build fixed decoder basis D [num_outputs x num_modes]
    !   D(i,k) = exp(-s_k * tau_i)
    !   s_k = k * pi,  tau_i = (i-1)/(n_out-1)
    !---------------------------------------------------------------------------
    if(this%decoder_basis%allocated) call this%decoder_basis%deallocate()
    call this%decoder_basis%allocate( &
         [this%num_outputs, this%num_modes, 1])
    this%decoder_basis%is_sample_dependent = .false.
    this%decoder_basis%requires_grad = .false.
    this%decoder_basis%fix_pointer = .true.
    this%decoder_basis%is_temporary = .false.

    do k = 1, this%num_modes
       s = real(k, real32) * pi
       do i = 1, this%num_outputs
          if(this%num_outputs .gt. 1)then
             t = real(i-1, real32) / real(this%num_outputs-1, real32)
          else
             t = 0.0_real32
          end if
          idx = i + (k-1) * this%num_outputs
          this%decoder_basis%val(idx, 1) = exp(-s * t)
       end do
    end do


    !---------------------------------------------------------------------------
    ! Allocate output arrays
    !---------------------------------------------------------------------------
    if(allocated(this%output)) deallocate(this%output)
    allocate(this%output(1,1))
    if(this%z(1)%allocated) call this%z(1)%deallocate()

  end subroutine init_fixed_lno
!###############################################################################


!###############################################################################
  subroutine print_to_unit_fixed_lno(this, unit)
    !! Print fixed LNO settings and parameters to a unit
    use coreutils, only: to_upper
    implicit none

    ! Arguments
    class(fixed_lno_layer_type), intent(in) :: this
    !! Layer instance to print
    integer, intent(in) :: unit
    !! Output unit number

    write(unit,'(3X,"NUM_INPUTS = ",I0)') this%num_inputs
    write(unit,'(3X,"NUM_OUTPUTS = ",I0)') this%num_outputs
    write(unit,'(3X,"NUM_MODES = ",I0)') this%num_modes
    write(unit,'(3X,"USE_BIAS = ",L1)') this%use_bias
    if(this%activation%name .ne. 'none')then
       call this%activation%print_to_unit(unit)
    end if

    write(unit,'("WEIGHTS")')
    write(unit,'(5(E16.8E2))') this%params(1)%val(:,1)   ! R
    write(unit,'(5(E16.8E2))') this%params(2)%val(:,1)   ! W
    if(this%use_bias)then
       write(unit,'(5(E16.8E2))') this%params(3)%val(:,1) ! b
    end if
    write(unit,'("END WEIGHTS")')

  end subroutine print_to_unit_fixed_lno
!###############################################################################


!###############################################################################
  subroutine read_fixed_lno(this, unit, verbose)
    use athena__tools_infile, only: assign_val, assign_vec, move
    use coreutils, only: to_lower, to_upper, icount
    use athena__activation, only: read_activation
    use athena__initialiser, only: initialiser_setup
    implicit none

    ! Arguments
    class(fixed_lno_layer_type), intent(inout) :: this
    !! Layer instance to populate from file data
    integer, intent(in) :: unit
    !! Input unit number
    integer, optional, intent(in) :: verbose
    !! Verbosity level

    ! Local variables
    integer :: stat, verbose_ = 0
    !! I/O status and effective verbosity level
    integer :: j, k, c, itmp1, iline
    !! Loop counters and parser scratch integers
    integer :: num_inputs, num_outputs, num_modes
    !! Parsed layer dimensions
    logical :: use_bias = .true.
    !! Parsed bias flag
    character(14) :: kernel_initialiser_name='', bias_initialiser_name=''
    !! Parsed initialiser names
    class(base_actv_type), allocatable :: activation
    !! Parsed activation object
    class(base_init_type), allocatable :: kernel_initialiser, bias_initialiser
    !! Parsed initialiser objects
    character(256) :: buffer, tag, err_msg
    !! Input buffer, parsed tag and formatted error message
    real(real32), allocatable, dimension(:) :: data_list
    !! Temporary storage for flattened parameter blocks
    integer :: param_line, final_line, num_vals
    !! Weights-section line markers and current block size

    if(present(verbose)) verbose_ = verbose

    iline = 0
    param_line = 0
    final_line = 0
    tag_loop: do
       read(unit,'(A)',iostat=stat) buffer
       if(stat.ne.0)then
          write(err_msg,'("file encountered error (EoF?) before END ",A)') &
               to_upper(this%name)
          call stop_program(err_msg)
          return
       end if
       if(trim(adjustl(buffer)).eq."") cycle tag_loop

       if(trim(adjustl(buffer)).eq."END "//to_upper(trim(this%name)))then
          final_line = iline
          backspace(unit)
          exit tag_loop
       end if
       iline = iline + 1

       tag=trim(adjustl(buffer))
       if(scan(buffer,"=").ne.0) tag=trim(tag(:scan(tag,"=")-1))

       select case(trim(tag))
       case("NUM_INPUTS")
          call assign_val(buffer, num_inputs, itmp1)
       case("NUM_OUTPUTS")
          call assign_val(buffer, num_outputs, itmp1)
       case("NUM_MODES")
          call assign_val(buffer, num_modes, itmp1)
       case("USE_BIAS")
          call assign_val(buffer, use_bias, itmp1)
       case("ACTIVATION")
          iline = iline - 1
          backspace(unit)
          activation = read_activation(unit, iline)
       case("KERNEL_INITIALISER", "KERNEL_INIT", "KERNEL_INITIALIZER")
          call assign_val(buffer, kernel_initialiser_name, itmp1)
       case("BIAS_INITIALISER", "BIAS_INIT", "BIAS_INITIALIZER")
          call assign_val(buffer, bias_initialiser_name, itmp1)
       case("WEIGHTS")
          kernel_initialiser_name = 'zeros'
          bias_initialiser_name   = 'zeros'
          param_line = iline
       case default
          if(scan(to_lower(trim(adjustl(buffer))),&
               'abcdfghijklmnopqrstuvwxyz').eq.0)then
             cycle tag_loop
          elseif(tag(:3).eq.'END')then
             cycle tag_loop
          end if
          write(err_msg,'("Unrecognised line in input file: ",A)') &
               trim(adjustl(buffer))
          call stop_program(err_msg)
          return
       end select
    end do tag_loop
    kernel_initialiser = initialiser_setup(kernel_initialiser_name)
    bias_initialiser = initialiser_setup(bias_initialiser_name)

    call this%set_hyperparams( &
         num_outputs = num_outputs, &
         num_modes = num_modes, &
         use_bias = use_bias, &
         activation = activation, &
         kernel_initialiser = kernel_initialiser, &
         bias_initialiser = bias_initialiser, &
         verbose = verbose_ &
    )
    call this%init(input_shape=[num_inputs])

    if(param_line.eq.0)then
       write(0,*) "WARNING: WEIGHTS card in " // trim(this%name) // " not found"
    else
       call move(unit, param_line - iline, iostat=stat)

       ! Read R (num_modes^2 values)
       num_vals = num_modes * num_modes
       allocate(data_list(num_vals), source=0._real32)
       c = 1
       k = 1
       do while(c.le.num_vals)
          read(unit,'(A)',iostat=stat) buffer
          if(stat.ne.0) exit
          k = icount(buffer)
          read(buffer,*,iostat=stat) (data_list(j),j=c,c+k-1)
          c = c + k
       end do
       this%params(1)%val(:,1) = data_list
       deallocate(data_list)

       ! Read W (num_outputs * num_inputs values)
       num_vals = num_outputs * num_inputs
       allocate(data_list(num_vals), source=0._real32)
       c = 1
       k = 1
       do while(c.le.num_vals)
          read(unit,'(A)',iostat=stat) buffer
          if(stat.ne.0) exit
          k = icount(buffer)
          read(buffer,*,iostat=stat) (data_list(j),j=c,c+k-1)
          c = c + k
       end do
       this%params(2)%val(:,1) = data_list
       deallocate(data_list)

       ! Read b if use_bias
       if(use_bias)then
          allocate(data_list(num_outputs), source=0._real32)
          c = 1
          k = 1
          do while(c.le.num_outputs)
             read(unit,'(A)',iostat=stat) buffer
             if(stat.ne.0) exit
             k = icount(buffer)
             read(buffer,*,iostat=stat) (data_list(j),j=c,c+k-1)
             c = c + k
          end do
          this%params(3)%val(:,1) = data_list(1:num_outputs)
          deallocate(data_list)
       end if

       read(unit,'(A)') buffer
       if(trim(adjustl(buffer)).ne."END WEIGHTS")then
          call stop_program("END WEIGHTS not where expected")
          return
       end if
    end if

    call move(unit, final_line - iline, iostat=stat)
    read(unit,'(A)') buffer
    if(trim(adjustl(buffer)).ne."END "//to_upper(trim(this%name)))then
       write(err_msg,'("END ",A," not where expected")') to_upper(this%name)
       call stop_program(err_msg)
       return
    end if

  end subroutine read_fixed_lno
!###############################################################################


!###############################################################################
  function read_fixed_lno_layer(unit, verbose) result(layer)
    !! Read a fixed LNO layer from file and return it
    implicit none

    ! Arguments
    integer, intent(in) :: unit
    !! Input unit number
    integer, optional, intent(in) :: verbose
    !! Verbosity level
    class(base_layer_type), allocatable :: layer
    !! Allocated base-layer instance containing the result

    ! Local variables
    integer :: verbose_ = 0
    !! Effective verbosity level

    if(present(verbose)) verbose_ = verbose
    allocate(layer, source=fixed_lno_layer_type( &
         num_outputs=0, num_modes=1))
    call layer%read(unit, verbose=verbose_)

  end function read_fixed_lno_layer
!###############################################################################


!###############################################################################
  subroutine forward_fixed_lno(this, input)
    !! Forward propagation for the Laplace Neural Operator layer
    !!
    !! Computes:
    !!   v = sigma( D @ R @ E @ u  +  W @ u  +  b )
    implicit none

    ! Arguments
    class(fixed_lno_layer_type), intent(inout) :: this
    !! Layer instance to execute
    class(array_type), dimension(:,:), intent(in) :: input
    !! Input batch tensor collection

    ! Local variables
    type(array_type), pointer :: ptr, ptr_spec, ptr_local
    !! Combined output, spectral-path output and local-path output


    ! Spectral pathway:  D @ R @ E @ u
    !---------------------------------------------------------------------------
    ptr_spec => matmul(this%encoder_basis, input(1,1))  ! [M, batch]
    ptr_spec => matmul(this%params(1), ptr_spec)         ! [M, batch]
    ptr_spec => matmul(this%decoder_basis, ptr_spec)     ! [n_out, batch]

    ! Local bypass:  W @ u
    !---------------------------------------------------------------------------
    ptr_local => matmul(this%params(2), input(1,1))      ! [n_out, batch]

    ! Combine
    !---------------------------------------------------------------------------
    ptr => ptr_spec + ptr_local

    ! Add bias
    !---------------------------------------------------------------------------
    if(this%use_bias)then
       ptr => ptr + this%params(3)
    end if

    ! Apply activation
    !---------------------------------------------------------------------------
    call this%output(1,1)%zero_grad()
    if(trim(this%activation%name) .eq. "none")then
       call this%output(1,1)%assign_and_deallocate_source(ptr)
    else
       call this%z(1)%zero_grad()
       call this%z(1)%assign_and_deallocate_source(ptr)
       this%z(1)%is_temporary = .false.
       ptr => this%activation%apply(this%z(1))
       call this%output(1,1)%assign_and_deallocate_source(ptr)
    end if
    this%output(1,1)%is_temporary = .false.

  end subroutine forward_fixed_lno
!###############################################################################


!###############################################################################
  function get_attributes_fixed_lno(this) result(attributes)
    !! Return list of fixed LNO attributes for ONNX export
    implicit none

    ! Arguments
    class(fixed_lno_layer_type), intent(in) :: this
    !! Instance of the fixed LNO layer
    type(onnx_attribute_type), allocatable, dimension(:) :: attributes
    !! List of attributes for ONNX export

    ! Local variables
    character(32) :: buffer
    !! Buffer for integer-to-string conversion

    allocate(attributes(5))

    write(buffer, '(I0)') this%num_inputs
    attributes(1) = onnx_attribute_type( &
         name='num_inputs', type='int', val=trim(buffer))
    write(buffer, '(I0)') this%num_outputs
    attributes(2) = onnx_attribute_type( &
         name='num_outputs', type='int', val=trim(buffer))
    write(buffer, '(I0)') this%num_modes
    attributes(3) = onnx_attribute_type( &
         name='num_modes', type='int', val=trim(buffer))
    if(this%use_bias)then
       buffer = '1'
    else
       buffer = '0'
    end if
    attributes(4) = onnx_attribute_type( &
         name='use_bias', type='int', val=trim(buffer))
    attributes(5) = onnx_attribute_type( &
         name='activation', type='string', val=trim(this%activation%name))

  end function get_attributes_fixed_lno
!###############################################################################


!###############################################################################
  subroutine emit_onnx_nodes_fixed_lno( &
       this, prefix, nodes, num_nodes, max_nodes, inits, num_inits, &
       max_inits, input_name, is_last_layer, format)
    !! Emit decomposed standard ONNX nodes for a Fixed LNO layer.
    !!
    !! Forward: v = sigma(D * R * E * u + W * u + b)
    !! where E and D are fixed Laplace bases, R is a learnable mixing matrix.
    use coreutils, only: pi
    implicit none

    ! Arguments
    class(fixed_lno_layer_type), intent(in) :: this
    !! Fixed LNO layer instance
    character(*), intent(in) :: prefix
    !! Layer name prefix (e.g. "layer1")
    type(onnx_node_type), intent(inout), dimension(:) :: nodes
    !! Node accumulator
    integer, intent(inout) :: num_nodes
    !! Node counter
    integer, intent(in) :: max_nodes
    !! Node limit
    type(onnx_initialiser_type), intent(inout), dimension(:) :: inits
    !! Initialiser accumulator
    integer, intent(inout) :: num_inits
    !! Initialiser counter
    integer, intent(in) :: max_inits
    !! Initialiser limit
    character(*), optional, intent(in) :: input_name
    !! Name of the input tensor
    logical, optional, intent(in) :: is_last_layer
    !! Whether this is the last layer
    integer, optional, intent(in) :: format
    !! Export format selector

    ! Local variables
    integer :: j, k, idx, n
    real(real32) :: s, t
    real(real32), allocatable :: e_data(:), d_data(:)
    character(128) :: e_name, d_name, r_name, w_name, b_name
    character(128) :: trans_in_out, mm_e_out, mm_r_out, mm_d_out
    character(128) :: mm_w_out, add_out, add_b_out, final_output, &
         output_source
    integer :: format_

    format_ = 1
    if(present(format)) format_ = format
    if(format_ .ne. 2) return
    if(.not.present(input_name)) return
    if(.not.present(is_last_layer)) return

    !--------------------------------------------------------------------------
    ! Build names
    !--------------------------------------------------------------------------
    write(e_name, '(A,".E")') trim(prefix)
    write(d_name, '(A,".D")') trim(prefix)
    write(r_name, '(A,".R")') trim(prefix)
    write(w_name, '(A,".W")') trim(prefix)
    write(b_name, '(A,".b")') trim(prefix)

    write(trans_in_out, '("/",A,"/Transpose_output_0")') trim(prefix)
    write(mm_e_out, '("/",A,"/MatMul_output_0")') trim(prefix)
    write(mm_r_out, '("/",A,"/MatMul_1_output_0")') trim(prefix)
    write(mm_d_out, '("/",A,"/MatMul_2_output_0")') trim(prefix)
    write(mm_w_out, '("/",A,"/MatMul_3_output_0")') trim(prefix)
    write(add_out, '("/",A,"/Add_output_0")') trim(prefix)
    write(add_b_out, '("/",A,"/Add_1_output_0")') trim(prefix)

    !--------------------------------------------------------------------------
    ! Emit nodes
    !--------------------------------------------------------------------------
    ! 1. Transpose(input)
    call emit_nop_input_transpose(trim(prefix), trim(input_name), nodes, &
         num_nodes, trim(trans_in_out))

    ! 2. MatMul(E, x_t)
    num_nodes = num_nodes + 1
    write(nodes(num_nodes)%name, '("/",A,"/MatMul")') trim(prefix)
    nodes(num_nodes)%op_type = 'MatMul'
    allocate(nodes(num_nodes)%inputs(2))
    nodes(num_nodes)%inputs(1) = trim(e_name)
    nodes(num_nodes)%inputs(2) = trim(trans_in_out)
    allocate(nodes(num_nodes)%outputs(1))
    nodes(num_nodes)%outputs(1) = trim(mm_e_out)
    nodes(num_nodes)%attributes_json = ''

    ! 3. MatMul(R, encoded)
    num_nodes = num_nodes + 1
    write(nodes(num_nodes)%name, '("/",A,"/MatMul_1")') trim(prefix)
    nodes(num_nodes)%op_type = 'MatMul'
    allocate(nodes(num_nodes)%inputs(2))
    nodes(num_nodes)%inputs(1) = trim(r_name)
    nodes(num_nodes)%inputs(2) = trim(mm_e_out)
    allocate(nodes(num_nodes)%outputs(1))
    nodes(num_nodes)%outputs(1) = trim(mm_r_out)
    nodes(num_nodes)%attributes_json = ''

    ! 4. MatMul(D, mixed)
    num_nodes = num_nodes + 1
    write(nodes(num_nodes)%name, '("/",A,"/MatMul_2")') trim(prefix)
    nodes(num_nodes)%op_type = 'MatMul'
    allocate(nodes(num_nodes)%inputs(2))
    nodes(num_nodes)%inputs(1) = trim(d_name)
    nodes(num_nodes)%inputs(2) = trim(mm_r_out)
    allocate(nodes(num_nodes)%outputs(1))
    nodes(num_nodes)%outputs(1) = trim(mm_d_out)
    nodes(num_nodes)%attributes_json = ''

    ! 5. MatMul(W, x_t)
    num_nodes = num_nodes + 1
    write(nodes(num_nodes)%name, '("/",A,"/MatMul_3")') trim(prefix)
    nodes(num_nodes)%op_type = 'MatMul'
    allocate(nodes(num_nodes)%inputs(2))
    nodes(num_nodes)%inputs(1) = trim(w_name)
    nodes(num_nodes)%inputs(2) = trim(trans_in_out)
    allocate(nodes(num_nodes)%outputs(1))
    nodes(num_nodes)%outputs(1) = trim(mm_w_out)
    nodes(num_nodes)%attributes_json = ''

    ! 6. Add(spectral, local)
    num_nodes = num_nodes + 1
    write(nodes(num_nodes)%name, '("/",A,"/Add")') trim(prefix)
    nodes(num_nodes)%op_type = 'Add'
    allocate(nodes(num_nodes)%inputs(2))
    nodes(num_nodes)%inputs(1) = trim(mm_d_out)
    nodes(num_nodes)%inputs(2) = trim(mm_w_out)
    allocate(nodes(num_nodes)%outputs(1))
    nodes(num_nodes)%outputs(1) = trim(add_out)
    nodes(num_nodes)%attributes_json = ''

    ! 7. Add(combined, bias)
    if(this%use_bias)then
       num_nodes = num_nodes + 1
       write(nodes(num_nodes)%name, '("/",A,"/Add_1")') trim(prefix)
       nodes(num_nodes)%op_type = 'Add'
       allocate(nodes(num_nodes)%inputs(2))
       nodes(num_nodes)%inputs(1) = trim(add_out)
       nodes(num_nodes)%inputs(2) = trim(b_name)
       allocate(nodes(num_nodes)%outputs(1))
       nodes(num_nodes)%outputs(1) = trim(add_b_out)
       nodes(num_nodes)%attributes_json = ''
    end if

    if(this%use_bias)then
       output_source = add_b_out
    else
       output_source = add_out
    end if
    call emit_nop_output_tail(trim(prefix), trim(this%activation%name), &
         is_last_layer, trim(output_source), nodes, num_nodes, final_output)

    !--------------------------------------------------------------------------
    ! Emit initialisers
    !--------------------------------------------------------------------------
    ! E: fixed encoder basis [M, n_in] in row-major
    n = this%num_modes * this%num_inputs
    allocate(e_data(n))
    do j = 1, this%num_inputs
       if(this%num_inputs .gt. 1)then
          t = real(j - 1, real32) / real(this%num_inputs - 1, real32)
       else
          t = 0.0_real32
       end if
       do k = 1, this%num_modes
          s = real(k, real32) * pi
          idx = (k - 1) * this%num_inputs + j
          e_data(idx) = exp(-s * t)
       end do
    end do
    call emit_float_initialiser(trim(e_name), e_data, &
         [this%num_modes, this%num_inputs], inits, num_inits)
    deallocate(e_data)

    ! D: fixed decoder basis [n_out, M] in row-major
    n = this%num_outputs * this%num_modes
    allocate(d_data(n))
    do k = 1, this%num_modes
       s = real(k, real32) * pi
       do j = 1, this%num_outputs
          if(this%num_outputs .gt. 1)then
             t = real(j - 1, real32) / real(this%num_outputs - 1, real32)
          else
             t = 0.0_real32
          end if
          idx = (j - 1) * this%num_modes + k
          d_data(idx) = exp(-s * t)
       end do
    end do
    call emit_float_initialiser(trim(d_name), d_data, &
         [this%num_outputs, this%num_modes], inits, num_inits)
    deallocate(d_data)

    ! R: spectral mixing [M, M] in row-major
    call emit_matrix_initialiser(trim(r_name), this%params(1)%val(:,1), &
         this%num_modes, this%num_modes, inits, num_inits)

    ! W: bypass weights [n_out, n_in] in row-major
    call emit_matrix_initialiser(trim(w_name), this%params(2)%val(:,1), &
         this%num_outputs, this%num_inputs, inits, num_inits)

    ! b: bias [n_out, 1]
    if(this%use_bias)then
       call emit_float_initialiser(trim(b_name), this%params(3)%val(:,1), &
            [this%num_outputs, 1], inits, num_inits)
    end if

  end subroutine emit_onnx_nodes_fixed_lno
!###############################################################################

end module athena__fixed_lno_layer
