module ktae

  ! Module for Eigenvalue Bessel 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)

  character(*), parameter, private :: version = &
       & '$Id: ktae.f90,v 1.6 2003/12/18 21:25:49 pletzer Exp $'
  integer, parameter  :: Nd = 1 ! space dimension
  integer, parameter  :: Nc = 1 ! no of components

  type ktae_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
     
     ! equation parameters
     integer :: mpol, ntor ! poloidal/toroidal mode number
     real(r8) :: q0, q1    ! axis/edgesafety factor values
     real(r8) :: rhobar, rho_delta ! gyroradius

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

  end type ktae_obj

contains

  subroutine ktae_init(this, Nz, Nj, sigma_dx2, sampling)
    implicit none
    type(ktae_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 = 1
    this % tol = 1.e-10_r8

    this % mpol = 1
    this % ntor = 1
    this % q0 = 1.1_r8
    this % q1 = 2.1_r8
    this % rhobar = 0.01_r8
    this % rho_delta = 0

    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 ktae_init

  subroutine ktae_free(this)
    implicit none
    type(ktae_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 ktae_free

  subroutine ktae_set_domain(this, xmin, xmax)
    implicit none
    type(ktae_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 ktae_set_domain

  subroutine ktae_set_tol(this, tol)
    implicit none
    type(ktae_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 ktae_set_tol

  subroutine ktae_assemble(this)
    implicit none
    type(ktae_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, x, xb
    complex(r8) :: val, gabor_fct, d1, d2, d3, d4, x4_del_square
    logical newcol_flag
    integer rank, ier, i, nnz, base, i2, i1, m

    ! 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
    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 *,'Ktae::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)

          ! Boundary conditions

          ! Dirichlet are applied on the boundary collocation row
          if( abs(xpos(1)-xmin(1)) < this%eps ) then
             val = gabor_fct
             if(abs(val) > this % tol ) then
                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, val, ier)
                call sparseC16_error(this % bmat, 6, ier)
                newcol_flag = .FALSE.
                cycle
             endif
          end if
          if( abs(xpos(1)-xmax(1)) < this%eps  ) then
             val = gabor_fct
             if(abs(val) > this % tol ) then
                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, val, ier)
                call sparseC16_error(this % bmat, 6, ier)
                newcol_flag = .FALSE.
                cycle
             endif
          end if
          ! Higher order derivative BCs (modify row that is one node inside the domain)
          xb = xmin(1)+xg%delta_x(1)
          if( abs(xpos(1)-xb) < this%eps ) then
             gabor_fct = GaborN(xmin, zpos, kwav, sigma)
             if(this % mpol /= 1) then
                ! want y'(0) = 0 for m/=1
                call derivative(1, xmin(1), kwav(1), zpos(1), sigma(1), val, ier)
             else
                ! want y''(0) = 0 for m=1
                call derivative(2, xmin(1), kwav(1), zpos(1), sigma(1), val, ier)
             endif
             val = val * gabor_fct

             call sparseC16_set_next(this % bmat, i1-1, newcol_flag, val, ier)
             call sparseC16_error(this % bmat, 6, ier)

             call sparseC16_set_next(this % amat, i1-1, newcol_flag, val, ier)
             call sparseC16_error(this % amat, 6, ier)             
             newcol_flag = .FALSE.
             cycle
          end if
          xb = xmax(1)-xg%delta_x(1)
          if( abs(xpos(1)-xb) < this%eps  ) then
             gabor_fct = GaborN(xmax, zpos, kwav, sigma)
             call derivative(1, xmax(1), kwav(1), zpos(1), sigma(1), val, ier)

             val = val * gabor_fct

             call sparseC16_set_next(this % bmat, i1-1, newcol_flag, val, ier)
             call sparseC16_error(this % bmat, 6, ier)

             call sparseC16_set_next(this % amat, i1-1, newcol_flag, val, ier)
             call sparseC16_error(this % amat, 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(4, xpos(1), kwav(1), zpos(1), sigma(1), D4, ier)
          call derivative(3, xpos(1), kwav(1), zpos(1), sigma(1), D3, ier)
          call derivative(2, xpos(1), kwav(1), zpos(1), sigma(1), D2, ier)
          call derivative(1, xpos(1), kwav(1), zpos(1), sigma(1), D1, ier)

          x = xpos(1)
          m = this % mpol

          ! Left hand side

          ! x^4 rho^2 del^4
          val = m**2*(-4 + m**2) + x*((1 + 2*m**2)*D1 + x*(-((1 + 2*m**2)*D2) + x*(2*D3 + x*D4)))
          val = (this%rhobar*(1+(0,1)*this%rho_delta))**2 * val * gabor_fct

          ! x^4 del^2
          x4_del_square = x**2 *(x**2 *D2 + x*D1 - m**2) * gabor_fct

          val = val -  &
               & K_square(this, x,m)*x4_del_square - &
               & x**4 * K_square_prime(this, x,m)*gabor_fct*(D1 - 1/x) 

          
          call sparseC16_set_next(this % amat, i1-1, newcol_flag, val, ier)
          call sparseC16_error(this % amat, 6, ier)
          

          ! Right hand side

          ! - x^4 del^2
          val = -  x4_del_square
          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 ktae_assemble

  subroutine ktae_solve(this, xvec, lambda, nmax)
    implicit none
    type(ktae_obj) :: this
    complex(r8), intent(inout) :: xvec(:) ! initial eigenvector guess on input/eigenvector on output
    complex(r8), intent(inout) :: lambda ! initial eigenvalue 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

    complex(r8) :: res

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

  end subroutine ktae_solve

  subroutine ktae_get_solution(this, xvec, x, y)
    implicit none
    type(ktae_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 ktae_get_solution

  subroutine ktae_plot(this, x, y, title)
    implicit none
    type(ktae_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 = 'Ktae '//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 ktae_plot

  subroutine ktae_test(this)

    ! Test module 

    use supralu_mod, only : sparseC16_get_size
    implicit none
    type(ktae_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 = 10 !! 20
    Nj = 2 !! 5
    sigma_dx2 = 2
    sampling = 1.5_r8    
    call ktae_init(this, Nz, Nj, sigma_dx2, sampling)

    xmin = 0
    xmax = 1
    call ktae_set_domain(this, xmin, xmax)

    tol = 1.e-6_r8 !! 1.e-10_r8
    call ktae_set_tol(this, tol)

    this % mpol = 1
    this % ntor = 1

    this % q0 = 2
    this % q1 = 2.1
    this % rhobar = 0.1_r8 ! 0.01_r8
    this % rho_delta = 0

    call ktae_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.262984_r8 !0.271574_r8 !0.493198_r8 !0.410244_r8 ! 0.638599_r8 !0.346723_r8 !0.29076_r8  ! initial eigenvalue guess
    lambda = 0.5236810816371932_r8 ! 0.55_r8 !0.51_r8
    nmax = 20 ! max no of iterations
    call ktae_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 ktae_get_solution(this, xvec, x, y)
    ! normalize
    y = y / maxval(abs(y))
    write(title, '(a,2f10.4)') 'Lambda=', lambda
    call ktae_plot(this, x, y, title)
    deallocate(x, y)
    deallocate(xvec)

    call ktae_free(this)


  end subroutine ktae_test

!...................................................................................
! extra methods

  function q(this, x)

    ! safety factor profile

    implicit none
    type(ktae_obj) :: this
    real(r8), intent(in) :: x
    real(r8) :: q
    q = this%q0 * (1 + (this%q1/this%q0 - 1) * x**2)

  end function q

  function q_prime(this, x)

    ! derivative of safety factor profile

    implicit none
    type(ktae_obj) :: this
    real(r8), intent(in) :: x
    real(r8) :: q_prime
    q_prime = 2 * (this%q1 - this%q0) * x

  end function q_prime

  function K_square(this, x, m)

    implicit none
    type(ktae_obj) :: this
    real(r8), intent(in) :: x
    integer, intent(in) :: m
    real(r8) :: K_square
    
    K_square = (this%ntor - m/q(this,x))**2

  end function K_square

  function K_square_prime(this, x, m)

    implicit none
    type(ktae_obj) :: this
    real(r8), intent(in) :: x
    integer, intent(in) :: m
    real(r8) :: K_square_prime

    integer n
    real(r8) :: q0, q1

    q0 = this % q0
    q1 = this % q1
    n = this % ntor
    
    K_square_prime = (-4*m*(q0 - q1)*x*(m - n*q0 + n*(q0 - q1)*x**2))/ &
         & (-q0 + (q0 - q1)*x**2)**3

  end function K_square_prime

end module ktae
