    module kdtree
!
!  SYNOPSIS
!    use kdtree
!
!  DESCRIPTION
!    These routines implement a KD tree for rapid lookup of observations
!    contained in a given patch (cube) for the data assimilation routines.
!
!  REQUIRED PACKAGES
!    ferror, resolution - in support directory
!
!  AUTHOR
!    Eric Kostelich, kostelich@asu.edu, 480-965-5006
!
!  MODIFICATION HISTORY
!    06/03/05 - searches on all 3 coordinates of any observation.
!    10/12/04 - first implementation.
!
    use grid_types
    use ferror
    use resolution
    implicit none
!--------------------
!  Public module data
!
    type obsinfo
       real::coord(3)  ! 1:x (lon), 2:y (lat), 3:sigma (z) coord of observation
       integer::index  ! into obs array
       integer::typ    ! so we can grab pressure or other selected obs types
    end type obsinfo
!
    type kdnode
       type(obsinfo)::value
       integer::left  ! index (in nodearray) of left child
       integer::right  ! index (in nodearray) of right child
       integer::depth  ! 1=root node
    end type kdnode
!------------------
!  Private module data
!
    type(kdnode),allocatable,save,private::nodearray(:)  ! the KD tree
    integer,save,private::nalloc=0  ! # of entries filled in nodearray
    private::nodealloc,insert,shuffle,lookfor,append
    contains
!-----------------------------------------------------------------------------
    subroutine search(sw,ne,found,nfound,obstype)
!  SEARCH - look for observations whose coordinates sit in the closure of
!  the box whose lower left coords are given by SW and upper right by NE.
!  The box may not straddle the poles but is allowed to straddle the
!  0-degree longitude line.
!  Place the list of nodes in FOUND and return the number of valid nodes
!  in NFOUND.  Public subroutine.
!  If OBSTYPE is present, then return only those observations whose type
!  is one of those in OBSTYPE.
!
    type(patch_index),intent(in)::sw,ne
    type(obsinfo),allocatable,intent(inout)::found(:)
    integer,intent(out)::nfound
    integer,intent(in),optional::obstype(:)
!
    integer,parameter::BUFSIZ=100
    type(obsinfo),allocatable::wfound(:),efound(:)  ! left or right of LONF
    integer::wf,ef  ! number found left or right of LONF
    type(patch_index)::bdy
!
!  First check whether the box straddles the wrap-around longitude LONF.
!  If it does (sw(x) > ne(x)) then do two searches in the KD tree:
!  one from sw(x) to LONF and another from 1 to ne(x).  The KD tree
!  does not understand modulo distances.
!
    nfound=0
    if(sw%x.gt.ne%x) then  ! wrap at LONF
       allocate(wfound(BUFSIZ),efound(BUFSIZ))
       bdy=patch_index(LONF,ne%y,ne%z)   ! x to LONF
       wf=0
       call lookfor(sw,bdy,1,1,wfound,wf,obstype)
       bdy=patch_index(1,sw%y,sw%z)  ! 1=LONF to x
       ef=0
       call lookfor(bdy,ne,1,1,efound,ef,obstype)
       nfound=wf+ef
       if(allocated(found)) then
          if(size(found).lt.nfound) then
             deallocate(found)
             allocate(found(nfound))
          endif
       else
          allocate(found(nfound))
       endif
       found(1:wf)=wfound(1:wf)
       found(wf+1:nfound)=efound(1:ef)
    else  ! ordinary lookup
       if(.not.allocated(found)) allocate(found(BUFSIZ))
       call lookfor(sw,ne,1,1,found,nfound,obstype)
    endif
    return
    end subroutine search
!-----------------------------------------------------------------------------
    recursive subroutine lookfor(low,high,t,lev,found,nfound,obstype)
!  LOOKFOR - look for nodes in the KD tree rooted at T for nodes whose coords
!  sit in the closure of the box whose lower left coords are in LOW and
!  whose upper right coords are in HIGH.  Append the coordinates of all
!  NFOUND enclosed points in the allocatable array FOUND (which may be
!  longer than NFOUND entries).
!  If OBSTYPE is present, then we gather only those nodes corresponding
!  to observations whose type is one of those in OBSTYPE.
!
    type(patch_index),intent(in)::low,high
    integer,intent(in)::t,lev
    type(obsinfo),allocatable,intent(inout)::found(:)
    integer,intent(inout)::nfound
    integer,intent(in),optional::obstype(:)
!
!  Local variables
!
    integer::nextlev
    real::nodeval(3),lowcoord(3),highcoord(3)
!
    if(t.eq.0) return  ! leaf node
    if(t.lt.0) call fabort('lookfor: botched t')
    if(lev.lt.1.or.lev.gt.3) call fabort('lookfor: botched level')
!
!  If we're inside the box, then append the node at T to the list in FOUND.
!  Append only OBS of the specified types if requested.
!
    nodeval=nodearray(t)%value%coord
    lowcoord=(/low%x,low%y,low%z/)
    highcoord=(/high%x,high%y,high%z/)
    if(all(lowcoord.le.nodeval.and.nodeval.le.highcoord)) then
       if(present(obstype)) then
          if(any(nodearray(t)%value%typ.eq.obstype)) &
             call append(nodearray(t)%value,found,nfound)
       else
          call append(nodearray(t)%value,found,nfound)
       endif
    endif
!
!  Recurse left or right as needed.
!
    nextlev=mod(lev,3)+1
    if(lowcoord(lev).le.nodeval(lev)) &
       call lookfor(low,high,nodearray(t)%left,nextlev,found,nfound,obstype)
    if(nodeval(lev).le.highcoord(lev)) &
       call lookfor(low,high,nodearray(t)%right,nextlev,found,nfound,obstype)
    return
    end subroutine lookfor
!-----------------------------------------------------------------------------
    subroutine build_kdtree(list,n)
!  BUILD_KDTREE - build a KD tree for the observations given a list of
!  N coordinates.  Public subroutine.
!
    integer,intent(in)::n
    type(obsinfo),intent(inout)::list(n)
!
    integer::j,t,depth
!
!  If there are no observations, then just allocate a dummy root node.
!  Otherwise, build the tree.
!
    if(n.eq.0) then
       allocate(nodearray(1))
       nodearray(1)%left=0
       nodearray(1)%right=0
       nodearray(1)%value%coord=0  ! so it matches no coords
       nodearray(1)%value%index=0
       nodearray(1)%depth=1
    else
       call shuffle(list,n)
       allocate(nodearray(n))
       t=0
       do j=1,n
          if(t.eq.0) then
             depth=1
          else
             depth=nodearray(t)%depth
          endif
          call insert(list(j),t,1,depth)
       enddo
    endif
    return
    end subroutine build_kdtree
!-----------------------------------------------------------------------------
    recursive subroutine insert(val,t,ic,depth)
!  INSERT - insert a node containing VAL into the KD tree rooted at T.
!  We construct the tree with ordinary comparisons on all coordinates.
!  Internal subroutine.
!  VAL  : value to be inserted
!  T :=: on entry index of node where we start searching (0=leaf node).
!      On return, the index of the node allocated for VAL.
!  IC : coordinate on which the comparison is based (1 <= IC <= 3)
!  DEPTH : current depth of the tree.
!
    type(obsinfo),intent(in)::val
    integer,intent(inout)::t
    integer,intent(in)::ic,depth
!
    if(t.eq.0) then
       t=nodealloc()
       nodearray(t)%left=0
       nodearray(t)%right=0
       nodearray(t)%value=val
       nodearray(t)%depth=depth
    else if(val%coord(ic).le.nodearray(t)%value%coord(ic)) then
       call insert(val,nodearray(t)%left,mod(ic,3)+1,depth+1)
    else
       call insert(val,nodearray(t)%right,mod(ic,3)+1,depth+1)
    endif
    return
    end subroutine insert
!-----------------------------------------------------------------------------
    integer function nodealloc()
!  NODEALLOC - return the index of a free node in NODEARRAY.  Internal function.
!
    nalloc=nalloc+1
    if(nalloc.gt.size(nodearray)) call fabort('nodealloc: out of space')
    nodealloc=nalloc
    return
    end function nodealloc
!-----------------------------------------------------------------------------
    subroutine shuffle(list,n)
!  SHUFFLE - randomize the contents of LIST, but let LIST(1) be the element
!  that is closest to the average value of each coordinate.  Use the maximum
!  norm for this purpose.  (We assume that the average is close to the median
!  for the data.)
!
    integer,intent(in)::n
    type(obsinfo),intent(inout)::list(n)
!
    type(obsinfo)::temp
    integer::j,k,minidx
    real::x,avg(size(list(1)%coord)),dmin,dv
!
    do j=1,n
       call random_number(x)
       k=int(x*n)+1
       temp=list(k)
       list(k)=list(1)
       list(1)=temp
    enddo
!
!  Find the average value of each coordinate and the index MINIDX of
!  the LIST element that is closest to the average.  Swap LIST(1) with
!  LIST(MINIDX).
!
    avg=list(1)%coord
    do j=2,n
       avg=avg+list(j)%coord
    enddo
    avg=avg/n
    minidx=1
    dmin=sum(abs(list(1)%coord-avg))
    do j=2,n
       dv=sum(abs(list(j)%coord-avg))
       if(dv.lt.dmin) then
          dmin=dv
          minidx=j
       endif
    enddo
    temp=list(minidx)
    list(minidx)=list(1)
    list(1)=temp
    return
    end subroutine shuffle
!-----------------------------------------------------------------------------
    subroutine append(val,found,nfound)
!  Append VAL to the list of FOUND nodes, and enlarge FOUND as needed.
!
    type(obsinfo),intent(in)::val
    type(obsinfo),allocatable,intent(inout)::found(:)
    integer,intent(inout)::nfound
!
    if(nfound.ge.size(found)) call enlarge(found)
    nfound=nfound+1
    found(nfound)=val
    return
    end subroutine append
!-----------------------------------------------------------------------------
    subroutine enlarge(f,n)
!  ENLARGE - If N is present, then enlarge F to size N.  Otherwise,
!  enlarge F by 50% or 1 element, whichever is greater.
!  Preserve the initial elements in F.
!
    type(obsinfo),allocatable,intent(inout)::f(:)
    integer,intent(in),optional::n
!
    type(obsinfo)::temp(size(f))
    integer::sz,newsize
!
    sz=size(f)
    if(present(n)) then
       newsize=max(n,sz)
    else
       newsize=max((3*sz)/2,sz+1)
    endif
    if(newsize.gt.sz) then
       temp=f
       deallocate(f)
       allocate(f(newsize))
       f(1:sz)=temp
    endif
    return
    end subroutine enlarge
!-----------------------------------------------------------------------------
    subroutine kdstats(nodes,maxdepth,nodecount)
!  KDSTATS - report statistics on the KD tree.
!  Arguments:
!  NODES := number of allocated nodes in the tree
!  MAXDEPTH := maximum depth of the tree
!  NODECOUNT := Kth element is number of nodes in the tree at depth K.
!
    integer,intent(out)::nodes,maxdepth
    integer,allocatable,intent(out)::nodecount(:)
!
    integer::j,idepth(nalloc)
!
    nodes=nalloc
    idepth=nodearray(1:nalloc)%depth
    maxdepth=maxval(idepth)
    allocate(nodecount(maxdepth))
    nodecount=0
    do j=1,nalloc
       nodecount(idepth(j))=nodecount(idepth(j))+1
    enddo
    return
    end subroutine kdstats
!-----------------------------------------------------------------------------
    end module kdtree
