!! ------------------------------------------------------
!! The routines in this file set up geometry for
!! advection and diffusion terms.  These do not need
!! to be modified for specific geometries.  All such
!! modifications can be made in the <geo>_setup.f90 files,
!! where <geo> is 'channel', 'annulus', 'irr', or 'user'
!! (for user defined geometries).
!! ------------------------------------------------------


!! -------------------------------------------------------
!! Call initialization routines
!! -------------------------------------------------------
SUBROUTINE initialize_bdry(mx,my,dx,dy,x,y,ax,ay,bx,by, iprob)
  IMPLICIT REAL (a-h,o-z)
  INCLUDE "parms.h"

  SELECT CASE (iprob)
  CASE(channel)
     CALL channel_initialize(mx,my,dx,dy,x,y,ax,ay,bx,by)
  CASE(annulus)
     CALL annulus_initialize(mx,my,dx,dy,x,y,ax,ay,bx,by)
  CASE(irr_objs)
     CALL irr_initialize(mx,my,dx,dy,x,y,ax,ay,bx,by)
  CASE(user_domain)
     CALL user_initialize(mx,my,dx,dy,x,y,ax,ay,bx,by)
  END SELECT

END SUBROUTINE initialize_bdry

!! -------------------------------------------------------
!! Level set defining flow domain;  Call individual
!! level set routines
!! -------------------------------------------------------
REAL FUNCTION fbody(x,y)
  IMPLICIT REAL (a-h,o-z)
  INCLUDE "parms.h"

  REAL irr_fbody
  COMMON /problem_type/ iprob

  SELECT CASE(iprob)
  CASE(channel)
     fbody = channel_fbody(x,y)

  CASE(annulus)
     fbody = annulus_fbody(x,y)

  CASE(irr_objs)
     fbody = irr_fbody(x,y)

  CASE(user_domain)
     fbody = user_fbody(x,y)

  END SELECT
  RETURN

END FUNCTION fbody


!! -------------------------------------------------------
!! Set up cell and edge fractions, centers of mass, midpoints
!! (for approximating diffusive fluxes at partial edges)
!! and boundary segments.
!! -------------------------------------------------------
SUBROUTINE get_geometry(mx,my,dx,dy,x,y,frside,kappa, midpts, &
     & centers, bdry_seg)
  IMPLICIT REAL(a-h,o-z)
  INCLUDE "parms.h"

  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc) :: kappa, bdry_seg
  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc,2) :: centers, frside
  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc,2,2) :: midpts

  REAL, DIMENSION(-1:maxgrid) :: x,y

  DO i = 1-mbc,mx+mbc
     DO j = 1-mbc,my + mbc
        xi = x(i) - dx/2
        yj = y(j) - dy/2
        CALL cartfr(xi,yj,dx,dy,frarea,frl,frb, xc, yc, xm, ym,blen)
        frside(i,j,1) = frl
        frside(i,j,2) = frb
        kappa(i,j) = frarea
        centers(i,j,1) = xc
        centers(i,j,2) = yc
        midpts(i,j,1,1) = xi
        midpts(i,j,1,2) = ym
        midpts(i,j,2,1) = xm
        midpts(i,j,2,2) = yj
        bdry_seg(i,j) = blen
     END DO
  END DO
END SUBROUTINE get_geometry


!! -------------------------------------------------------
!! Set up stream function. This assumes that individual
!! routines set up values in ghost cells as well.
!! -------------------------------------------------------
SUBROUTINE set_psi(mx,my,dx,dy,x,y,psi,iprob)
  IMPLICIT REAL(a-h,o-z)
  INCLUDE "parms.h"

  SELECT CASE(iprob)
  CASE(channel)
     CALL channel_stream(mx,my,dx,dy,x,y,psi)
  CASE(annulus)
     CALL annulus_stream(mx,my,dx,dy,x,y,psi)
  CASE(irr_objs)
     CALL irr_stream(mx,my,dx,dy,x,y,psi)
  CASE(user_domain)
     CALL user_stream(mx,my,dx,dy,x,y,psi)
  END SELECT
END SUBROUTINE set_psi


!! -------------------------------------------------------
!! Set up velocities (u = aux(i,j,1); v = aux(i,j,2)),
!! modified kappa (aux(i,j,3)) and beta for diffusive
!! fluxes.
!! ------------------------------------------------------
SUBROUTINE set_aux(mx,my,dx,dy,x,y,frside,kappa,psi,D,aux,beta)
  IMPLICIT REAL(a-h,o-z)

  INCLUDE "parms.h"

  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc) :: psi, kappa
  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc,*) :: aux,frside, beta
  REAL, DIMENSION(-1:maxgrid) :: x,y

  !! Initialize to zero
  aux(1-mbc:mx+mbc,1-mbc:my+mbc,1:3) = 0.0

  DO i = 1-mbc,mx+mbc
     DO j = 1-mbc,my+mbc
        !! This requires that psi has valid values in
        !! ghost cells.
        aux(i,j,1) = (psi(i,j+1) - psi(i,j)) / dy
        aux(i,j,2) = (psi(i,j) - psi(i+1,j)) / dx
        aux(i,j,3) = kappa(i,j)

        DO iside = 1,2
           beta(i,j,iside) = frside(i,j,iside)*D
        END DO
        beta(i,j,3) = D  !! Don't need a fraction of D here; we multiply
                         !! by a length when determining RHS in src2_irr.

        !! Just to be sure that velocities are zero in no-flow domain
        !! This may not be conservative!
        IF (ABS(frside(i,j,1)) <= 1e-13) THEN
           aux(i,j,1) = 0.0
        ENDIF

        IF (ABS(frside(i,j,2)) <= 1e-13) THEN
           aux(i,j,2) = 0.0
        ENDIF

        !! Triangular cell fix up
        hmax = MAX(MAX(frside(i,j,1),frside(i,j,2)), &
             & MAX(frside(i+1,j,1),frside(i,j+1,2)))
        IF (hmax > 0.d0) THEN
           aux(i,j,3) = aux(i,j,3) / hmax
        ENDIF

        aux(i,j,3) = MAX(aux(i,j,3),1e-3)
     END DO
  END DO
END SUBROUTINE set_aux


!! -------------------------------------------------------
!! Determine coefficients needed to approximate diffusive
!! fluxes at full and partial edges. All routines that
!! follow are used by this routine.
!! -------------------------------------------------------
SUBROUTINE set_diffusion(mx,my,dx,dy,x,y,frside, kappa, &
     & centers, midpts, neighbors, coeffs)
  IMPLICIT REAL(a-h,o-z)
  INCLUDE "parms.h"

  REAL, DIMENSION(-1:maxgrid) :: x,y
  INTEGER, DIMENSION(1-mbc:maxmx+mbc, 1-mbc:maxmy+mbc,2,4,2) :: neighbors
  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc,2,4) :: coeffs
  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc) :: kappa
  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc,2) :: centers, frside
  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc,2,2) :: midpts

  REAL c(4)
  integer find_neighbors
  INTEGER, DIMENSION(2) :: i_idx, j_idx
  INTEGER, DIMENSION(4) :: n_idx, n_idy

  DATA i_idx /-1, 0/
  DATA j_idx /0, -1/

  !! Set neighbors for diffusion approximations
  DO i = 1,mx+1
     DO j = 1,my+1
        DO iside = 1,2
           ix = i + i_idx(iside) !! Left or bottom neighbor
           iy = j + j_idx(iside)
           ignum = find_neighbors(iside,i,j,kappa)
           CALL set_box_idx(ignum,n_idx,n_idy)
           DO k = 1,4
              ixn = i + n_idx(k)
              iyn = j + n_idy(k)
              neighbors(i,j,iside,k,1) = ixn
              neighbors(i,j,iside,k,2) = iyn
           END DO
           !! If Cij or neighbor is partial, we need to approximate
           !! flux on partial side.
           IF ( (ABS(kappa(i,j)-1.0) > 1e-15 .OR. &
                & ABS(kappa(ix,iy)-1.0) > 1e-15) .AND. &
                & (ABS(kappa(i,j)) > 1e-15 .AND. &
                & ABS(kappa(ix,iy)) > 1e-15)) THEN
              xm = midpts(i,j,iside,1)
              ym = midpts(i,j,iside,2)
              CALL compute_coeffs(mx,my,dx,dy,x,y,iside,i,j, &
                   & centers,xm,ym, neighbors,coeffs)
           ELSE
              c = 0.0
              SELECT CASE(iside)
              CASE(1)  !! Left edge
                 SELECT CASE(ignum)
                 CASE(1)
                    c(1) = 1.0/dx
                    c(4) = -1.0/dx
                 CASE(2)
                    c(1) = 1.0/dx
                    c(2) = -1.0/dx
                 END SELECT
              CASE(2)  !! Bottom edge
                 SELECT CASE(ignum)
                 CASE(1)
                    c(1) = 1.0/dy
                    c(2) = -1.0/dy
                 CASE(4)
                    c(1) = 1.0/dy
                    c(4) = -1.0/dy
                 END SELECT
              END SELECT

              DO k = 1,4
                 coeffs(i,j,iside,k) = c(k)
              END DO
           END IF
        END DO
     END DO
  END DO
END SUBROUTINE set_diffusion


INTEGER FUNCTION  find_neighbors(iside,ix,iy, kappa)
  IMPLICIT REAL (a-h,o-z)

  INCLUDE "parms.h"

  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc) :: kappa

   SELECT CASE (iside)
   CASE(1)
      IF (kappa(ix-1,iy-1) >= kappa(ix-1,iy+1)) THEN
         ignum = 1
      ELSE
         ignum = 2
      END IF
   CASE(2)
      IF (kappa(ix-1,iy-1) >= kappa(ix+1,iy-1)) THEN
         ignum = 1
      ELSE
         ignum = 4
      END IF
   CASE default
      WRITE(6,*) 'find_neighbors : Can not find four neighbors '
      WRITE(6,*) iside
      STOP
   END SELECT

  find_neighbors = ignum
  RETURN
END FUNCTION find_neighbors


SUBROUTINE set_box_idx(ignum,n_idx,n_idy)
  IMPLICIT REAL (a-h,o-z)
  INCLUDE "parms.h"

  INTEGER, DIMENSION(1:4) :: n_idx,n_idy,ig_idx1,ig_idx2,ig_idx3,ig_idx4,&
       & ig_idy1,ig_idy2, ig_idy3,ig_idy4

  DATA ig_idx1 /0,  0, -1, -1/
  DATA ig_idx2 /0, -1, -1,  0/
  DATA ig_idx3 /0,  0,  1,  1/
  DATA ig_idx4 /0,  1,  1,  0/

  DATA ig_idy1 /0, -1, -1,  0/
  DATA ig_idy2 /0,  0,  1,  1/
  DATA ig_idy3 /0,  1,  1,  0/
  DATA ig_idy4 /0,  0, -1, -1/

  SELECT CASE(ignum)
  CASE(1)
     n_idx = ig_idx1
     n_idy = ig_idy1
  CASE(2)
     n_idx = ig_idx2
     n_idy = ig_idy2
  CASE(3)
     n_idx = ig_idx3
     n_idy = ig_idy3
  CASE(4)
     n_idx = ig_idx4
     n_idy = ig_idy4
  CASE default
     WRITE(6,*) 'set_box_idx : ignum = ',ignum
     STOP
  END SELECT
END SUBROUTINE set_box_idx

SUBROUTINE compute_coeffs(mx,my,dx,dy,x,y,iside,ix,iy, &
     & centers,xm,ym,neighbors,coeffs)
  IMPLICIT REAL (a-h,o-z)

  INCLUDE "parms.h"

  INTEGER, DIMENSION(1-mbc:maxmx+mbc, 1-mbc:maxmy+mbc,2,4,2) :: neighbors
  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc,2,4) :: coeffs
  REAL, DIMENSION(1-mbc:maxmx+mbc,1-mbc:maxmy+mbc,2) :: centers

  REAL, DIMENSION(4,4) :: A,B,C, id4, Atmp
  REAL, DIMENSION(-1:maxgrid) :: x, y

  DATA id4 /1,0,0,0, &
         &  0,1,0,0, &
         &  0,0,1,0, &
         &  0,0,0,1/  !! 4x4 identity matrix


  !!First form matrix for computing coefficients
  A_loop : DO i = 1,4
     ixn = neighbors(ix,iy,iside,i,1)
     iyn = neighbors(ix,iy,iside,i,2)
     xc = centers(ixn,iyn,1)
     yc = centers(ixn,iyn,2)

     !! Solve system A'x to get rows of inv(A)
     A(1,i) = 1.0
     A(2,i) = xc
     A(3,i) = yc
     A(4,i) = xc*yc
  END DO A_loop

  !! now invert linear system to get rows 2,3,4 of inv(A').
  B = id4  !! RHS Set up above
  Atmp = A
  CALL dgesv(4,4,A,4,IPIV,B,4,info)

  IF (info /= 0) THEN
     WRITE(6,*) 'setup : dgesv failed ',info
     DO j = 1,4
        WRITE(6,'(1P4E28.16)') (Atmp(i,j),i=1,4)
     END DO
     WRITE(6,*) ' '
     WRITE(6,'(4I5)') (neighbors(ix,iy,iside,i,1),i=1,4)
     WRITE(6,*) ' '
     WRITE(6,'(4I5)') (neighbors(ix,iy,iside,i,2),i=1,4)
     WRITE(6,*) ' '
     WRITE(6,*) ix,iy
     STOP
  END IF

130 FORMAT(I5,3F8.3)

  !! Now compute coefficients
  DO i = 1,4
     SELECT CASE(iside)
     CASE(1)
        coeffs(ix,iy,iside,i) = B(i,2) + B(i,4)*ym   !! p_x(xi,ym)
     CASE(2)
        coeffs(ix,iy,iside,i) = B(i,3) + B(i,4)*xm   !! p_y(xm,yj)
     END SELECT
  END DO

END SUBROUTINE compute_coeffs
