! -------------------------------
! A multi-layer neural network supporting multiple hidden layers
! (c) 2011 Bui Hoang Hai - Modified for multi hidden layers and weight saving/loading
! -------------------------------

! It is assumed that the following global variables are declared elsewhere:
!   INTEGER :: Nin, Nhl, Nout
!   INTEGER, ALLOCATABLE :: hidden(:)       ! hidden(1) ... hidden(Nhl) : neurons in each hidden layer
!   REAL, ALLOCATABLE :: weight_ih(0:Nin, :)   ! For input -> first hidden layer
!   REAL, ALLOCATABLE :: moment_ih(0:Nin, :)
!   REAL, ALLOCATABLE :: weight_hl(:,:)        ! For connections between hidden layers (an array of matrices)
!   REAL, ALLOCATABLE :: moment_hl(:,:)        ! Array of same shape as weight_hl for momentum
!   REAL, ALLOCATABLE :: weight_ho(0: hidden(Nhl), Nout) ! For last hidden -> output
!   REAL, ALLOCATABLE :: moment_ho(0: hidden(Nhl), Nout)
!   REAL, ALLOCATABLE :: Inputs(:,:), Outputs(:,:), Yhq(:,:)
!   REAL, ALLOCATABLE :: fin(0: Nin), fhid(0: maxval(hidden)), fout(Nout)
!   REAL, ALLOCATABLE :: ftarget(Nout), delta_o(Nout)
!   REAL, ALLOCATABLE :: delta_h(:)  ! temporary for hidden layers
!   REAL :: learning_rate, momentum, error
!   INTEGER :: Nepoc, Nsample
!   LOGICAL :: training_mode, debug
!   CHARACTER(LEN=*) :: save_file
!   REAL :: mean_error, mean_error_check
!
! Also, variables used in controller_* routines (min_in, max_in, min_out, max_out, etc.)
! are assumed to be declared globally.
! ---------------------------------------------

Subroutine ANN_2024(Xinp, Youtp, Nspl, MXvar, MYvar, Miss)
  ! Modified main subroutine supporting multi hidden layers.
  integer :: Nspl, MXvar, MYvar, i, ic, Y1d(Nspl)
  Real :: Xinp(MXvar, Nspl), Youtp(MYvar, Nspl), X1d(MXvar), Miss
  Real :: Z1d(MYvar)
  
  ! Process valid samples as in your original code.
  Nin = MXvar
  Nout = MYvar
  ic = 0
  Y1d(:) = 0
  if (training_mode) then
     do i = 1, Nspl
        X1d(1:MXvar) = Xinp(1:MXvar, i)
        Z1d(1:MYvar) = Youtp(1:MYvar, i)
        if (.not. Empty_Ar1(Z1d, MYvar, Miss) .AND. .not. Empty_Ar1(X1d, MXvar, Miss)) then
           ic = ic + 1
           Y1d(i) = 1
        endif
     enddo
  else
     do i = 1, Nspl
        X1d(1:MXvar) = Xinp(1:MXvar, i)
        if (.not. Empty_Ar1(X1d, MXvar, Miss)) then
           ic = ic + 1
           Y1d(i) = 1
        endif
     enddo
  endif
  Nsample = ic

  if (Allocated(Yest)) deallocate(Yest)
  Allocate(Yest(Nspl, Nout))
  Yest(:,:) = Miss

  if (Nsample < Nin) return

  if (Allocated(Inputs)) deallocate(Inputs)
  if (Allocated(Outputs)) deallocate(Outputs)
  if (Allocated(Yhq)) deallocate(Yhq)
  Allocate(Inputs(Nin, Nsample), Outputs(Nout, Nsample), Yhq(Nsample, Nout))
  
  ic = 0
  do i = 1, Nspl
     if (Y1d(i) == 1) then
        ic = ic + 1
        Inputs(1:Nin, ic) = Xinp(1:Nin, i)
        Outputs(1:Nout, ic) = Youtp(1:Nout, i)
     endif
  enddo

  debug = .false.
  Call controller_init   ! Initializes ANN (including multi hidden layers) and scaling factors.
  if (training_mode) then
     Call controller_train_data
  endif
  Call controller_output_data
  ic = 0
  do i = 1, Nspl
     if (Y1d(i) == 1) then
        ic = ic + 1
        Yest(i, 1:Nout) = Yhq(ic, 1:Nout)
     endif
  enddo
  Call controller_deinit
End Subroutine ANN_2024

!---------------------------------------------------------------------------
Logical Function Empty_Ar1(X, N, Miss)
  integer :: N, i
  real :: X(N), Miss
  Logical :: Empty
  Empty = .False.
  do i = 1, N
     if (X(i) == Miss) then
        Empty = .True.
        Exit
     endif
  enddo
  Empty_Ar1 = Empty
End Function Empty_Ar1
!---------------------------------------------------------------------------

Subroutine controller_init
  ! Initialize controller, ANN, and scaling factors.
  Integer :: i
  Real :: tmp
  Call ann_init  ! Now initializes multi hidden layers.
  Allocate(min_in(Nin))
  Allocate(max_in(Nin))
  Allocate(min_out(Nout))
  Allocate(max_out(Nout))
  if (training_mode) then
     do i = 1, Nin
        min_in(i) = MINVAL(Inputs(i, :))
        max_in(i) = MAXVAL(Inputs(i, :))
     enddo
     do i = 1, Nout
        min_out(i) = MINVAL(Outputs(i, :))
        max_out(i) = MAXVAL(Outputs(i, :))
     enddo
  else
     call controller_load_info
  endif
End Subroutine controller_init

!---------------------------------------------    
Subroutine controller_train_data
  Integer :: epoc, i, j, k, temp, idx(Nsample)
  Real :: tmp, eps, r
  do k = 1, Nsample; idx(k) = k; enddo
  ! Shuffle idx (Fisher-Yates)
  do i = Nsample, 2, -1
     call random_number(r)
     j = 1 + int(i * r)
     temp = idx(i)
     idx(i) = idx(j)
     idx(j) = temp
  end do
  mean_error_check = 9999999.0
  eps = 0.001 
  do epoc = 1, Nepoc
     mean_error = 0.0
     do k = 1, Nsample
        j = idx(k)
        do i = 1, Nin
           Fin(i) = rescale(Inputs(i, j), min_in(i), max_in(i), -1.0, 1.0)
        enddo
        Call ann_forward
        do i = 1, Nout
           Ftarget(i) = rescale(Outputs(i, j), min_out(i), max_out(i), 0.1, 0.9)
        enddo
        Call ann_backward
        mean_error = mean_error + error / Nsample
     enddo
     if (epoc == 1) then
        mean_error_check = mean_error
     else
        tmp = ABS((mean_error_check - mean_error) / mean_error_check)
        mean_error_check = mean_error
        if (epoc > 5) then
           if (tmp < eps) then
              ! Optionally adjust learning_rate here.
           endif
        endif
     endif
  enddo
  call controller_save_info
End Subroutine controller_train_data

!---------------------------------------------    
Subroutine controller_output_data
  Integer :: i, j, NN
  do j = 1, Nsample
     do i = 1, Nin
        Fin(i) = rescale(Inputs(i, j), min_in(i), max_in(i), -1.0, 1.0)
     enddo
     Call ann_forward
     do i = 1, Nout
        Ftarget(i) = rescale(Fout(i), 0.1, 0.9, min_out(i), max_out(i))
     enddo
     Yhq(j, 1:Nout) = Ftarget(1:Nout)
  enddo
End Subroutine controller_output_data

!---------------------------------------------    
Subroutine controller_print_info
  Integer :: i, j
  write(*,*) "Input, Hidden, Output:"
  write(*,*) Nin, Nhl, Nout
  write(*,*) "Weight_input-hidden:"
  do i = 0, Nin
     write(*,*) (weight_ih(i, j), j = 1, hidden(1))
  enddo
  if (Nhl > 1) then
     write(*,*) "Weight_hidden-hidden (for each connection):"
     do i = 1, Nhl - 1
        write(*,*) "Layer", i, "to", i+1, weight_hl(i)
     enddo
  endif
  write(*,*) "Weight_hidden-output:"
  do i = 0, hidden(Nhl)
     write(*,*) (weight_ho(i, j), j = 1, Nout)
  enddo
  write(*,*) "min_in:"; write(*,*) min_in
  write(*,*) "Max_in:"; write(*,*) max_in
  write(*,*) "min_out:"; write(*,*) min_out
  write(*,*) "Max_out:"; write(*,*) max_out
End Subroutine controller_print_info

!---------------------------------------------
Subroutine controller_save_info
  ! Save weights (for all layers) and scaling factors.
  Integer :: i, j
  open(1, file=save_file, form="formatted", status="replace")
  ! Save dimensions: Nin, Nhl, Nout and hidden array.
  write(1,*) Nin, Nhl, Nout
  write(1,*) hidden(1:Nhl)
  ! Save weight_ih:
  do i = 0, Nin
     write(1,*) (weight_ih(i, j), j = 1, hidden(1))
  enddo
  ! Save each hidden-hidden weight matrix if Nhl > 1:
  if (Nhl > 1) then
     do i = 1, Nhl-1
        write(1,*) "Layer", i, "to", i+1, size(weight_hl(i),1), size(weight_hl(i),2)
        do j = 0, hidden(i)
           write(1,*) (weight_hl(i)(j, k), k = 1, hidden(i+1))
        enddo
     enddo
  endif
  ! Save weight_ho:
  do i = 0, hidden(Nhl)
     write(1,*) (weight_ho(i, j), j = 1, Nout)
  enddo
  ! Save scaling factors:
  write(1,*) min_in
  write(1,*) max_in
  write(1,*) min_out
  write(1,*) max_out
  close(1)
End Subroutine controller_save_info

!---------------------------------------------
Subroutine controller_load_info
  Integer :: Nin1, Nhl1, Nout1, i, j, dummy
  open(1, file=save_file, form="formatted", status="old")
  read(1,*) Nin1, Nhl1, Nout1
  read(1,*) hidden(1:Nhl)  ! Assume dimensions match.
  if ((Nin /= Nin1) .or. (Nout /= Nout1)) then
     write(*,*) "Mismatch in dimensions between saved file and input file!"
     stop
  endif
  do i = 0, Nin
     read(1,*) (weight_ih(i, j), j = 1, hidden(1))
  enddo
  if (Nhl > 1) then
     do i = 1, Nhl-1
        read(1,*) dummy, dummy,  ! For label and dimensions (can be ignored)
        do j = 0, hidden(i)
           read(1,*) (weight_hl(i)(j, k), k = 1, hidden(i+1))
        enddo
     enddo
  endif
  do i = 0, hidden(Nhl)
     read(1,*) (weight_ho(i, j), j = 1, Nout)
  enddo
  read(1,*) min_in
  read(1,*) max_in
  read(1,*) min_out
  read(1,*) max_out
  close(1)
End Subroutine controller_load_info

!---------------------------------------------
Subroutine controller_deinit
  if (allocated(min_in))    deallocate(min_in)
  if (allocated(max_in))    deallocate(max_in)
  if (allocated(min_out))   deallocate(min_out)
  if (allocated(max_out))   deallocate(max_out)
  if (allocated(weight_ih)) deallocate(weight_ih)
  if (allocated(weight_ho)) deallocate(weight_ho)
  if (Nhl > 1) then
     if (allocated(weight_hl)) deallocate(weight_hl)
     if (allocated(moment_hl)) deallocate(moment_hl)
  endif
  if (allocated(fin))       deallocate(fin)
  if (allocated(fhid))      deallocate(fhid)
  if (allocated(fout))      deallocate(fout)
  if (allocated(ftarget))   deallocate(ftarget)
  if (allocated(delta_o))   deallocate(delta_o)
  if (allocated(delta_h))   deallocate(delta_h)
  if (allocated(moment_ih)) deallocate(moment_ih)
  if (allocated(moment_ho)) deallocate(moment_ho)
  if (debug) close(111)
End Subroutine controller_deinit

!---------------------------------------------
Real function rescale(x, x1, x2, y1, y2)
  Real :: x, x1, x2, y1, y2
  rescale = y1 + (y2 - y1) * (x - x1) / (x2 - x1)
End Function rescale
!---------------------------------------------

Subroutine ann_init
  ! Initialize ANN weights for multi hidden layers.
  Integer :: iin, ihid, iout, k, j
  Real :: tmp
  ! Allocate input->first hidden connection:
  allocate(weight_ih(0:Nin, hidden(1)))
  allocate(moment_ih(0:Nin, hidden(1)))
  do iin = 0, Nin
     do ihid = 1, hidden(1)
        call random_number(tmp)
        weight_ih(iin, ihid) = 2.0 - tmp
     enddo
  enddo
  ! Allocate hidden-hidden connections if more than one hidden layer:
  if (Nhl > 1) then
     allocate(weight_hl(Nhl-1))
     allocate(moment_hl(Nhl-1))
     do k = 1, Nhl-1
        allocate(weight_hl(k)(0:hidden(k), hidden(k+1)))
        allocate(moment_hl(k)(0:hidden(k), hidden(k+1)))
        do iin = 0, hidden(k)
           do ihid = 1, hidden(k+1)
              call random_number(tmp)
              weight_hl(k)(iin, ihid) = 2.0 - tmp
           enddo
        enddo
     enddo
  endif
  ! Allocate last hidden -> output connection:
  allocate(weight_ho(0:hidden(Nhl), Nout))
  allocate(moment_ho(0:hidden(Nhl), Nout))
  do iin = 0, hidden(Nhl)
     do iout = 1, Nout
        call random_number(tmp)
        weight_ho(iin, iout) = 2.0 - tmp
     enddo
  enddo

  moment_ih = 0.0
  if (Nhl > 1) then
     do k = 1, Nhl-1
        moment_hl(k) = 0.0
     enddo
  endif
  moment_ho = 0.0

  allocate(fin(0:Nin))
  allocate(fhid(0:maxval(hidden)))  ! Use maximum hidden layer size for temporary storage.
  allocate(fout(Nout))
  allocate(ftarget(Nout))
  allocate(delta_o(Nout))
  if (Nhl > 0) then
     allocate(delta_h(maxval(hidden)))
  endif

  fin(0) = 1.0
  if (debug) open(111, file='debug.txt')
End Subroutine ann_init

!---------------------------------------------
Subroutine ann_forward
  ! Forward propagation for multi hidden layers.
  Real :: asum
  Integer :: i, j, k, layer
  ! Compute first hidden layer activation:
  do j = 1, hidden(1)
     asum = 0.0
     do i = 0, Nin
        asum = asum + fin(i) * weight_ih(i, j)
     enddo
     fhid(j) = sigma(asum)
  enddo
  ! For additional hidden layers:
  if (Nhl > 1) then
     do layer = 2, Nhl
        do j = 1, hidden(layer)
           asum = 0.0
           do i = 0, hidden(layer-1)
              asum = asum + fhid(i) * weight_hl(layer-1)(i, j)
           enddo
           fhid(j) = sigma(asum)
        enddo
     enddo
  endif
  ! Compute output layer:
  do j = 1, Nout
     asum = 0.0
     do i = 0, hidden(Nhl)
        asum = asum + fhid(i) * weight_ho(i, j)
     enddo
     fout(j) = sigma(asum)
  enddo
  if (debug) then
     write(111,*) 'FORWARD'
     write(111,*) 'fin:', fin
     write(111,*) 'weight_ih:', weight_ih
     if (Nhl > 1) then
        do layer = 1, Nhl-1
           write(111,*) 'weight_hl(', layer, '):', weight_hl(layer)
        enddo
     endif
     write(111,*) 'weight_ho:', weight_ho
     write(111,*) 'fout:', fout
  endif
End Subroutine ann_forward

!---------------------------------------------
Subroutine ann_compute_error
  error = 0.0
  Integer :: iout
  do iout = 1, Nout
     delta_o(iout) = (ftarget(iout) - fout(iout))
     error = error + (delta_o(iout)**2) / Nout
  enddo
End Subroutine ann_compute_error

!---------------------------------------------
Subroutine ann_backward
  ! Backward propagation for multi hidden layers.
  Real :: delta_ih, delta_hl, delta_ho
  Integer :: iout, ihid, iin, layer
  error = 0.0
  ! Update weights for output layer:
  do iout = 1, Nout
     delta_o(iout) = (ftarget(iout) - fout(iout))
     error = error + (delta_o(iout)**2) / Nout
     delta_o(iout) = delta_o(iout) * fout(iout) * (1.0 - fout(iout))
     do ihid = 0, hidden(Nhl)
        delta_ho = learning_rate * delta_o(iout) * fhid(ihid)
        weight_ho(ihid, iout) = weight_ho(ihid, iout) + delta_ho + moment_ho(ihid, iout)
        moment_ho(ihid, iout) = momentum * delta_ho
     enddo
  enddo

  ! Backpropagate error for last hidden layer:
  allocate(delta_h(maxval(hidden)))
  do ihid = 1, hidden(Nhl)
     delta_h(ihid) = 0.0
     do iout = 1, Nout
        delta_h(ihid) = delta_h(ihid) + weight_ho(ihid, iout) * delta_o(iout)
     enddo
     delta_h(ihid) = delta_h(ihid) * fhid(ihid) * (1.0 - fhid(ihid))
  enddo

  if (Nhl == 1) then
     ! Only one hidden layer; update weight_ih.
     do iin = 0, Nin
        do ihid = 1, hidden(1)
           delta_ih = learning_rate * delta_h(ihid) * fin(iin)
           weight_ih(iin, ihid) = weight_ih(iin, ihid) + delta_ih + moment_ih(iin, ihid)
           moment_ih(iin, ihid) = momentum * delta_ih
        enddo
     enddo
  else
     ! For multiple hidden layers, update the connection between the last two hidden layers.
     do iin = 0, hidden(Nhl-1)
        do ihid = 1, hidden(Nhl)
           delta_hl = learning_rate * delta_h(ihid) * fhid(iin)
           weight_hl(Nhl-1)(iin, ihid) = weight_hl(Nhl-1)(iin, ihid) + delta_hl + moment_hl(Nhl-1)(iin, ihid)
           moment_hl(Nhl-1)(iin, ihid) = momentum * delta_hl
        enddo
     enddo
     ! Further backpropagation through earlier hidden layers would follow here.
  endif

  if (debug) then
     write(111,*) 'BACK'
     write(111,*) '--fout', fout
     write(111,*) '--ftarget', ftarget
     write(111,*) 'delta_o', delta_o
     write(111,*) 'weight_ih', weight_ih
     if (Nhl > 1) then
        do layer = 1, Nhl-1
           write(111,*) 'weight_hl(', layer, '):', weight_hl(layer)
        enddo
     endif
     write(111,*) 'weight_ho', weight_ho
  endif
  deallocate(delta_h)
End Subroutine ann_backward

!---------------------------------------------
Real function sigma(x)
  Real :: x, tmp
  tmp = 1.0 / (1.0 + exp(-x))
  sigma = tmp
End Function sigma
!---------------------------------------------