program deeponet
    use torch_ftn
    use iso_fortran_env

    implicit none
    integer :: m,n,l,i, j, k

    type(torch_module) :: torch_mod
    type(torch_tensor_wrap) :: input_tensors
    type(torch_tensor) :: out_tensor

    integer :: arglen, stat
    integer :: NX, NY, NZ, NT, NX1, NY1 
    integer :: NXY, NXY1, NZT
    real(real32), allocatable :: data_in(:)    ! NX*NY
    real(real32), allocatable :: data_trunk(:) ! NT

    real(real32), allocatable :: encin(:,:,:,:)
    real(real32), pointer     :: encout(:,:) 
    character(:), allocatable :: filename

   !real(real32), allocatable :: dnt_dat(:) ! 248 x 36
    real(real32), allocatable :: dnt_t(:, :)
    real(real32), allocatable :: dntin(:, :, :, :)
    real(real32), allocatable :: dntin_f(:, :, :, :)
    real(real32), pointer     :: dntout(:, :, :, :)

    real(real32), allocatable :: decin(:, :)
    real(real32), pointer :: decout(:, :)

    NX   = 70  ! No of Longitude
    NY   = 70  ! No of Latitude 
    NZ   = 72  ! No of Levels  (46??)   
    NT   = 248 ! No of Times
    NX1  = 6   ! ML NX 
    NY1  = 6   ! ML NY
    NXY  = NX*NY
    NXY1 = NX1*NY1
    NZT  = NZ*NT

    allocate (data_in(NXY))
    open(12, file="V_bf_revised_checked.out")
    read(12,*) data_in   
    close(12)

    !allocate (data_trunk(NT))
    !open(51, file="x_trunk.out")
    !read(51, *) data_trunk
    !close(51)
    !dnt_t = reshape(data_trunk, (/1, NT/))
    !print *, "Output time: ", dnt_t
    !stop
 
    allocate (encin(NX,NY,1,NZ)) ! (FD(NX) , FD(NY), NCCHANNEL, BATCHSIZE)
    allocate (dnt_t(1,NT))
    allocate (dntin(NXY1,NT,1,NZ))
    allocate (decin(NXY1,NZ))
    do l = 1,NT
      dnt_t(1,l) = 2.0* (l -1)/(NT - 1) - 1.0
    end do
   ! print *, "Output time: ", dnt_t
   ! stop
    do l = 1,NZ 
      do i = 1,NX
        do j = 1, NY 
          n = (i-1)*NY + j 
          encin(i,j,1,l) = data_in(n) 
        end do
      end do 
    end do 
    !call encoder to convert E3SM to deepONet space 
    call input_tensors%create
    call input_tensors%add_array(encin)
    call torch_mod%load("V_Encoder.pt")
    call torch_mod%forward(input_tensors, out_tensor)
    call out_tensor%to_array(encout)
    print *, "Output after encoder is: ", shape(encout)
    print *, minval(encout), maxval(encout)

    ! call DeepONet model 
    do l = 1,NT
      do k = 1,NZ
        do i = 1,NX1
          do j = 1, NY1
            n = (i-1)*NY1 + j
            m = (l-1)*NZ + k
            dntin(n,l,1,k) = encout(n,k)
          end do
        end do 
      end do 
    end do
    print *, "Input before deepONet is: ", shape(dntin)
    print *, minval(dntin), maxval(dntin)

    call input_tensors%create
    call input_tensors%add_array(dntin)
    call input_tensors%add_array(dnt_t)
    call torch_mod%load("V_DeepONet.pt")
    call torch_mod%forward(input_tensors, out_tensor, 1)
    call out_tensor%to_array(dntout)
    print *, "Output after deepONet is", shape(dntout)
    print *, minval(dntout), maxval(dntout)

    ! call DeepONet model 
    do i = 1,NX1
      do j = 1, NY1
        do k = 1,NZ
          n = (i-1)*NY1 + j
          decin(n,k) = dntout(i,j,1,k)
        end do 
      end do
    end do
    call input_tensors%create
    call input_tensors%add_array(decin)  ! allocate (decin(NX,NY,1,NT))
    call torch_mod%load("V_Decoder.pt")
    call torch_mod%forward(input_tensors, out_tensor)
    call out_tensor%to_array(decout)
    print *, "Output after decoder is", shape(decout)
    print *, minval(decout), maxval(decout)

    !open(18, file = "dnet.dat", action="write", status = "replace")
    !write(18, *) output
    !close(18)

end program deeponet
