module eigenwave

  ! Module for Eigenvalue wave test problem

  use supralu_mod
  use indexer
  use Gabor
  use kzgrid
  use xgrid
  use operators
  use plotmtv


  implicit none
  integer, parameter, private :: r8 = selected_real_kind(12,100)
  real(r8), parameter, private :: twopi = 6.2831853071795865_r8

  character(*), parameter, private :: version = &
       & '$Id: eigenwave.f90,v 1.2 2003/12/15 16:52:21 pletzer Exp $'
  integer, parameter  :: Nd = 1 ! space dimension
  integer, parameter  :: Nc = 1 ! no of components

  type eigenwave_obj
     integer :: Nz ! No of Gabor intervals
     integer :: Nj ! Fourier modes -Nj .. + Nj
     real(r8) :: sampling  ! sampling rate > 1
     real(r8) :: tol ! sparsity tolerance
     real(r8) :: xmin, xmax ! domain boundary

     real(r8) :: eps
     real(r8) :: sigma
     integer :: Nx
     type(sparseC16_obj) :: amat, bmat
     type(indexer_obj) :: iter2
     type(kzgrid_obj) :: kzg

  end type eigenwave_obj

contains

  subroutine eigenwave_init(this, Nz, Nj, sigma_dx2, sampling)
    implicit none
    type(eigenwave_obj) :: this
    integer, intent(in) :: Nz ! No of Gabor intervals
    integer, intent(in) :: Nj ! Fourier modes -Nj .. + Nj
    real(r8), intent(in)  :: sigma_dx2 ! inverse Gabor extent normalized to grid size
    real(r8), intent(in)  :: sampling  ! sampling rate > 1

    real(r8) :: delta_z

    this % xmin = 0
    this % xmax = twopi
    this % tol = 1.e-10_r8

    this % Nz = Nz
    this % Nj = Nj
    delta_z = (this%xmax - this%xmin)/this%Nz
    this % sigma = sigma_dx2 / delta_z**2
    this % sampling = sampling

    this % eps = 10 * epsilon(1._r8)
    this % Nx = (2*Nj+1)*(Nz+1) - 1

  end subroutine eigenwave_init

  subroutine eigenwave_free(this)
    implicit none
    type(eigenwave_obj) :: this

    integer iok
    call sparseC16_free(this % amat, iok)
    call sparseC16_free(this % bmat, iok)
    call indexer_free(this % iter2, iok)
    call kzgrid_free(this % kzg, iok)
  end subroutine eigenwave_free

  subroutine eigenwave_set_domain(this, xmin, xmax)
    implicit none
    type(eigenwave_obj) :: this
    real(r8), intent(in) :: xmin, xmax ! domain boundary

    this % xmin = xmin
    this % xmax = xmax
    this % sigma = this % sigma/(xmax - xmin)**2
  end subroutine eigenwave_set_domain

  subroutine eigenwave_set_tol(this, tol)
    implicit none
    type(eigenwave_obj) :: this
    real(r8), intent(in) :: tol ! sparsity tolerance

    ! A low value of tol will make the matrices less sparse
    this % tol = max(this%eps, tol)
  end subroutine eigenwave_set_tol

  subroutine eigenwave_assemble(this)
    implicit none
    type(eigenwave_obj) :: this

    type(indexer_obj) :: iter1
    type(xgrid_obj) :: xg
    real(r8) :: minus_log_tol

    integer all_ind1(1+Nd), all_ind2(1+Nd*2), Nx(Nd), Nz(Nd), Nj(Nd)
    integer sizes1(1+Nd), starts1(1+Nd), sizes2(1+2*Nd), starts2(1+2*Nd)
    integer kz_map(2*Nd), x_map(Nd)
    real(r8) :: xmin(Nd), xmax(Nd), sigma(Nd)
    real(r8) :: sampling(Nd), xpos(Nd), kwav(Nd), zpos(Nd), dx(Nd), LARGE
    complex(r8) :: val, gabor_fct
    logical newcol_flag
    integer rank, ier, i, nnz, base, i2, i1

    ! In order to exercize iterators, the matrix 
    ! is assembled using the general, multi-dimensional
    ! iterator based method, although for a 1-d 
    ! equation it might be more efficient to use do-loops.

    xmin = this % xmin
    xmax = this % xmax
    dx = xmax - xmin
    LARGE = 1._r8/this % tol
    sigma = this % sigma
    Nz = this % Nz
    Nj = this % Nj
    Nx = this % Nx
    sampling = this % sampling


    minus_log_tol = -log(this % tol) 

    ! iterators
    sizes1(1) = Nc; starts1(1) = 1;
    do i = 1, Nd
       sizes1(1+   i) = Nx(i)+1; starts1(1+   i) = 0
    end do
    call indexer_init(iter1, sizes1, starts1, ier)

    sizes2(1) = Nc; starts2(1) = 1;
    do i = 1, Nd
       sizes2(1+   i) = 2*Nj(i)+1; starts2(1+   i) = -Nj(i)
       sizes2(1+Nd+i) = Nz(i)  +1; starts2(1+Nd+i) = 0
    end do
    call indexer_init(this % iter2, sizes2, starts2, ier)

    ! grid data extractors
    do i = 1, Nd
       kz_map(i   ) = i   +1
       kz_map(i+Nd) = i+Nd+1
    end do
    call kzgrid_init(this % kzg, xmin, xmax, Nz, Nj, sampling, kz_map, ier)

    do i = 1, Nd
       x_map(i) = i + 1
    end do
    call xgrid_init(xg, xmin, xmax, Nx, x_map, ier)

    rank = iter1%ntot
    if(this % iter2%ntot /= rank) then
       print *,'Eigenwave::ERROR rank mismatch'
    endif
    nnz = rank**2 ! max no of elements
    base = 0
    call sparseC16_init(this % amat, rank, nnz, base, ier)
    call sparseC16_error(this % amat, 6, ier)
    call sparseC16_init(this % bmat, rank, nnz, base, ier)
    call sparseC16_error(this % bmat, 6, ier)

    ! start assembling

    do i2 = 1, this % iter2%ntot
       ! Iterate over Gabors
       newcol_flag = .TRUE.
       call indexer_next(this % iter2, ier)
       call indexer_get(this % iter2, all_ind2, ier)
       !ic2 = all_ind2(1) ! component
       ! get kwav & zpos
       call kzgrid_get(this % kzg, all_ind2, kwav, zpos, ier) 

       do i1 = 1, iter1%ntot
          ! Iterate over collocation points
          call indexer_next(iter1, ier)
          call indexer_get(iter1, all_ind1, ier)
          !ic1 = all_ind1(1) ! component
          ! get collocation point
          call xgrid_get(xg, all_ind1, xpos, ier)

          gabor_fct = GaborN(xpos, zpos, kwav, sigma)

          if( all(xpos <= (xmin+dx*this%eps)) ) then
             val = gabor_fct * LARGE ! force y(xmin) -> 0
             call sparseC16_set_next(this % amat, i1-1, newcol_flag, val, ier)
             call sparseC16_error(this % amat, 6, ier)
             call sparseC16_set_next(this % bmat, i1-1, newcol_flag, gabor_fct, ier)
             call sparseC16_error(this % bmat, 6, ier)
             newcol_flag = .FALSE.
             cycle
          end if
          if( all(xpos >= (xmax-dx*this%eps)) ) then
             val = gabor_fct * LARGE ! force y(xmax) -> 0
             call sparseC16_set_next(this % amat, i1-1, newcol_flag, val, ier)
             call sparseC16_error(this % amat, 6, ier)
             call sparseC16_set_next(this % bmat, i1-1, newcol_flag, gabor_fct, ier)
             call sparseC16_error(this % bmat, 6, ier)
             newcol_flag = .FALSE.
             cycle
          end if

          if (sum(0.5_r8*sigma*(xpos-zpos)**2) > minus_log_tol) cycle

          ! PDE operator

          call derivative(2, xpos(1), kwav(1), zpos(1), sigma(1), val, ier)
          val = - val * gabor_fct ! - y''
          call sparseC16_set_next(this % amat, i1-1, newcol_flag, val, ier)
          call sparseC16_error(this % amat, 6, ier)
          
          val = 1 * gabor_fct ! y
          call sparseC16_set_next(this % bmat, i1-1, newcol_flag, val, ier)
          call sparseC16_error(this % bmat, 6, ier)

          newcol_flag = .FALSE.
       end do
    end do

    call sparseC16_compare_shapes(this%amat, this%bmat, ier)
    call sparseC16_error(this % amat, 6, ier)
    if(ier/=0) stop

    call indexer_free(iter1, ier)
    call xgrid_free(xg, ier)

  end subroutine eigenwave_assemble

  subroutine eigenwave_solve(this, xvec, lambda, nmax)
    implicit none
    type(eigenwave_obj) :: this
    complex(r8), intent(inout) :: xvec(:) ! initial egenvector guess on input/eigenvector on output
    complex(r8), intent(inout) :: lambda ! initial egenvalue guess on input/eigenvalue on output
    integer, intent(inout) :: nmax ! max number of iterations on input/no of iterations on output

    real(r8) :: tol

    integer ier

    tol = 1.e-10_r8
    call sparseC16_eigen(this%amat, this%bmat, xvec, lambda, tol, nmax, ier)

  end subroutine eigenwave_solve

  subroutine eigenwave_get_solution(this, xvec, x, y)
    implicit none
    type(eigenwave_obj) :: this
    complex(r8), intent(in) :: xvec(:) ! eigenvector returned by solve method
    real(r8), intent(in) :: x(:) ! abscissae 
    complex(r8), intent(out) :: y(:) ! ordinates

    integer n, i1, i2, ier, all_ind2(1+Nd*2)
    real(r8) :: kwav(Nd), zpos(Nd)

    y = 0
    n = size(x)
    if(n==0) return
    do i1 = 1, n
       y(i1) = 0
       this % iter2 % counter = 0 ! reset
       do i2 = 1, this % iter2 % ntot
          call indexer_next(this % iter2, ier)
          call indexer_get(this % iter2, all_ind2, ier)
          !ic2 = all_ind2(1) ! component
          ! get kwav & zpos
          call kzgrid_get(this % kzg, all_ind2, kwav, zpos, ier)
          y(i1) = y(i1) + xvec(this % iter2%counter)* &
               & GaborN( (/x(i1)/), zpos, kwav, (/this % sigma/))
       end do
    end do
  end subroutine eigenwave_get_solution

  subroutine eigenwave_plot(this, x, y, title)
    implicit none
    type(eigenwave_obj) :: this
    type(plotmtv_obj) :: plt_obj
    real(r8), intent(in) :: x(:) ! abscissae
    complex(r8), intent(in) :: y(:) ! ordinates
    character(*), intent(in) :: title

    call plotmtv_init(plt_obj)
    plt_obj % toplabel = 'Eigenwave '//trim(adjustl(title))
    plt_obj % xlabel = 'x'
    plt_obj % ylabel = 'y'
    plt_obj % linelabel = 'Re'
    plt_obj % display = .FALSE.
    plt_obj % hold = .TRUE.
    call plotmtv_plot(plt_obj, x, real(y), 'b|', '  ') 
    plt_obj % linelabel = 'Im'
    plt_obj % display = .TRUE.
    call plotmtv_plot(plt_obj, x, aimag(y), 'rl', '  ')
    call plotmtv_free(plt_obj)

  end subroutine eigenwave_plot

  subroutine eigenwave_test(this)

    ! Test module 

    use supralu_mod, only : sparseC16_get_size
    implicit none
    type(eigenwave_obj) :: this
    integer :: Nz, Nj, nmax
    real(r8) sigma_dx2, sampling, xmin, xmax, tol
    real(r8), allocatable :: x(:)
    complex(r8), allocatable :: y(:), xvec(:)
    complex(r8) :: lambda
    character(len=80) :: title
    

    integer n, i, ier, nrows, ncols

    Nz = 20
    Nj = 1
    sigma_dx2 = 4
    sampling = 1.5_r8    
    call eigenwave_init(this, Nz, Nj, sigma_dx2, sampling)

    xmin = 0
    xmax = twopi
    call eigenwave_set_domain(this, xmin, xmax)

    tol = 1.e-8_r8
    call eigenwave_set_tol(this, tol)

    call eigenwave_assemble(this)
    call sparseC16_save(this%amat, 'amat.nc', ier)
    call sparseC16_save(this%bmat, 'bmat.nc', ier)
    call sparseC16_get_size(this%amat, nrows, ncols, ier)
    allocate(xvec(ncols))
    xvec = 1 ! initial eigenvector guess
    lambda = 0.25_r8 ! initial eigenvalue guess
    nmax = 10 ! max no of iterations
    call eigenwave_solve(this, xvec, lambda, nmax)
    print *,'Eigenvalue = ', lambda, ' obtained after ', nmax, ' iterations'

    n = this % iter2 % ntot * 10
    allocate(x(n), y(n))
    x = this%xmin + (this%xmax-this%xmin)*(/ (real(i-1,r8)/real(n-1, r8), i=1,n) /)
    call eigenwave_get_solution(this, xvec, x, y)
    ! normalize
    y = y / maxval(abs(y))
    write(title, '(a,2f6.3)') 'Lambda=', lambda
    call eigenwave_plot(this, x, y, title)
    deallocate(x, y)
    deallocate(xvec)

    call eigenwave_free(this)


  end subroutine eigenwave_test


end module eigenwave
