! Cross Gradient Module
! Author:  Mohammad Maysami
!          SEP, Stanford University
!          Feb 09

module xgrad_mod
  use adj_mod
  use sep3d_struct_mod

  implicit none
  integer, private :: n1,n2,n3
  real, private    :: d1,d2,d3
  real, private,      pointer :: aux(:)
  real, private,save, pointer :: temp(:)


contains
! ----------------------------------------------
subroutine xgrad_init(aux_s,aux_in)
  ! Initialize XGrad function
  ! aux_s types 3D argument passes the dimensions 
  ! aux_in passes the soft data
  type(sep3d)                  :: aux_s
  real, dimension(:),target    :: aux_in

  ! Assumes 1D array as input
  aux => aux_in
  
  ! Set size of aux
  n1=aux_s%n(1)
  n2=aux_s%n(2)
  d1=aux_s%d(1)
  d2=aux_s%d(2)

  if (sep3d_ndims(aux_s) >2) then
     n3=aux_s%n(3) ! Default to zero
     d3=aux_s%n(3)
     allocate(temp(n1*n2*n3*3))
  else
     n3=1
     d3=1.
     allocate(temp(n1*n2*n3*2))
  end if
  temp=0.

end subroutine


! ==============================================
function xgrad_lop (adj,add,mm,dd)  result(stat)
  ! General X-Gradient function
  ! Seems to work better with central difference
  logical, intent(in) :: adj,add
  real,dimension(:)   :: mm,dd 
  integer             :: stat,stat1,stat2 

  call adjnull (adj,add,mm,dd )

  if (n3==1) then 
     if (adj) then
        stat2=cross_lop2(adj, .false., temp,dd)
        stat1=grad_lop2(adj, add, mm,temp)

     else
        stat1=grad_lop2(adj, .false., mm, temp)
        stat2=cross_lop2(adj, add, temp, dd)

     end if
  else
     if (adj) then
        stat2=cross_lop3(adj, .false., temp,dd)
        stat1=grad_lop3(adj, add, mm,temp)

     else
       stat1=grad_lop3(adj, .false., mm, temp)
       stat2=cross_lop3(adj, add, temp, dd)
    end if
 end if
 
 stat=0
end function 


! =============================================
function grad_lop2 (adj,add,mm,dd) result(stat)
  ! Low-order 2D gradient function
  logical, intent(in)         :: adj,add
  real,dimension(n1,n2,n3)    :: mm
  real,dimension(n1,n2,n3,2)  :: dd
  integer                     :: stat,i

  call gradnull(adj,add,mm,dd)

  if (adj) then
     mm(1:n1-2,:,:)  =mm(1:n1-2,:,:) -(0.5/d1)*dd(2:n1-1,:,:,1)
     mm(3:n1  ,:,:)  =mm(3:n1  ,:,:) +(0.5/d1)*dd(2:n1-1,:,:,1)
  else
     dd(2:n1-1,:,:,1)=dd(2:n1-1,:,:,1) +(0.5/d1)*(mm(3:n1,:,:)-mm(1:n1-2,:,:))
  end if


  if (adj) then
     mm(:,1:n2-2,:)  =mm(:,1:n2-2,:) -(0.5/d2)*dd(:,2:n2-1,:,2)
     mm(:,3:n2  ,:)  =mm(:,3:n2  ,:) +(0.5/d2)*dd(:,2:n2-1,:,2)
  else
     dd(:,2:n2-1,:,2)= dd(:,2:n2-1,:,2) +(0.5/d2)*(mm(:,3:n2,:)-mm(:,1:n2-2,:))
  end if

  call xgrad_close()
  stat=0
end function

! =============================================
function grad_lop2h (adj,add,mm,dd) result(stat)
  ! High-order 2D gradient function
  logical, intent(in)         :: adj,add
  real,dimension(n1,n2,n3)    :: mm
  real,dimension(n1,n2,n3,2)  :: dd
  integer                     :: stat,i

  call gradnull(adj,add,mm,dd)

  if (adj) then
     mm(1:n1-2,:,:)  =mm(1:n1-2,:,:) +(0.5/d1)  *dd(1:n1-2,:,:,1)
     mm(2:n1-1,:,:)  =mm(2:n1-1,:,:) -(0.5/d1)*4*dd(1:n1-2,:,:,1)
     mm(3:n1  ,:,:)  =mm(3:n1  ,:,:) +(0.5/d1)*3*dd(1:n1-2,:,:,1)
  else
     dd(1:n1-2,:,:,1)=dd(1:n1-2,:,:,1) +(0.5/d1)*(mm(1:n1-2,:,:)-4*mm(2:n1-1,:,:)+3*mm(3:n1,:,:))
  end if


  if (adj) then
     mm(:,1:n2-2,:)  =mm(:,1:n2-2,:) +(0.5/d2)  *dd(:,1:n2-2,:,2)
     mm(:,2:n2-1,:)  =mm(:,2:n2-1,:) -(0.5/d2)*4*dd(:,1:n2-2,:,2)
     mm(:,3:n2  ,:)  =mm(:,3:n2  ,:) +(0.5/d2)*3*dd(:,1:n2-2,:,2)
  else
     dd(:,1:n2-2,:,2)= dd(:,1:n2-2,:,2) +(0.5/d2)*(mm(:,1:n2-2,:)-4*mm(:,2:n2-1,:)+3*mm(:,3:n2,:))
  end if

  call xgrad_close()
  stat=0
end function



! =============================================
function cross_lop2 (adj,add,mm,dd) result(stat)
  ! 2D cross mulitplier function
  logical, intent(in)         :: adj,add
  real,dimension(n1,n2,n3,2)  :: mm,gaux
  real,dimension(n1,n2,n3,1)    :: dd
  integer                     :: stat,stat1

  call crossnull(adj,add,mm,dd)
  stat1=grad_lop2 (.false.,.false.,aux,gaux)

  if (adj) then
     mm(:,:,:,1)=mm(:,:,:,1) -gaux(:,:,:,2)*dd(:,:,:,1)
     mm(:,:,:,2)=mm(:,:,:,2) +gaux(:,:,:,1)*dd(:,:,:,1)
  else
     ! It is actually dd(:,:,:,3)
     dd(:,:,:,1)=dd(:,:,:,1) -gaux(:,:,:,2)*mm(:,:,:,1)  +gaux(:,:,:,1)*mm(:,:,:,2)
  end if

  stat=0
end function


! =============================================
function grad_lop3 (adj,add,mm,dd) result(stat)
  ! Low-order 3D gradient function
  logical, intent(in)         :: adj,add
  real,dimension(n1,n2,n3)    :: mm
  real,dimension(n1,n2,n3,3)  :: dd
  integer                     :: stat,i

  call gradnull(adj,add,mm,dd)

  if (adj) then
     mm(1:n1-2,:,:)  =mm(1:n1-2,:,:) -(0.5/d1)*dd(2:n1-1,:,:,1)
     mm(3:n1  ,:,:)  =mm(3:n1  ,:,:) +(0.5/d1)*dd(2:n1-1,:,:,1)
  else
     dd(2:n1-1,:,:,1)=dd(2:n1-1,:,:,1) +(0.5/d1)*(mm(3:n1,:,:)-mm(1:n1-2,:,:))
  end if


  if (adj) then
     mm(:,1:n2-2,:)  =mm(:,1:n2-2,:) -(0.5/d2)*dd(:,2:n2-1,:,2)
     mm(:,3:n2  ,:)  =mm(:,3:n2  ,:) +(0.5/d2)*dd(:,2:n2-1,:,2)
  else
     dd(:,2:n2-1,:,2)= dd(:,2:n2-1,:,2) +(0.5/d2)*(mm(:,3:n2,:)-mm(:,1:n2-2,:))
  end if


  if (adj) then
     mm(:,:,1:n3-2)  =mm(:,:,1:n3-2) -(0.5/d3)*dd(:,:,1:n3-2,3)
     mm(:,:,3:n3  )  =mm(:,:,3:n3  ) +(0.5/d3)*dd(:,:,3:n3,3)
  else
     dd(:,:,2:n3-1,3)=dd(:,:,2:n3-1,3) +(0.5/d3)*(mm(:,:,3:n3)-mm(:,:,1:n3-2))
  end if

  stat=0
  call xgrad_close()
end function

! =============================================
function grad_lop3h (adj,add,mm,dd) result(stat)
  ! High-order 3D gradient function
  logical, intent(in)         :: adj,add
  real,dimension(n1,n2,n3)    :: mm
  real,dimension(n1,n2,n3,3)  :: dd
  integer                     :: stat,i

  call gradnull(adj,add,mm,dd)


  if (adj) then
     mm(1:n1-2,:,:)  =mm(1:n1-2,:,:) +(0.5/d1)  *dd(1:n1-2,:,:,1)
     mm(2:n1-1,:,:)  =mm(2:n1-1,:,:) -(0.5/d1)*4*dd(1:n1-2,:,:,1)
     mm(3:n1  ,:,:)  =mm(3:n1  ,:,:) +(0.5/d1)*3*dd(1:n1-2,:,:,1)
  else
     dd(1:n1-2,:,:,1)=dd(1:n1-2,:,:,1) +(0.5/d1)*(mm(1:n1-2,:,:)-4*mm(2:n1-1,:,:)+3*mm(3:n1,:,:))
  end if


  if (adj) then
     mm(:,1:n2-2,:)  =mm(:,1:n2-2,:) +(0.5/d2)  *dd(:,1:n2-2,:,2)
     mm(:,2:n2-1,:)  =mm(:,2:n2-1,:) -(0.5/d2)*4*dd(:,1:n2-2,:,2)
     mm(:,3:n2  ,:)  =mm(:,3:n2  ,:) +(0.5/d2)*3*dd(:,1:n2-2,:,2)
  else
     dd(:,1:n2-2,:,2)= dd(:,1:n2-2,:,2) +(0.5/d2)*(mm(:,1:n2-2,:)-4*mm(:,2:n2-1,:)+3*mm(:,3:n2,:))
  end if


  if (adj) then
     mm(:,:,1:n3-2)  =mm(:,:,1:n3-2) +(0.5/d3)  *dd(:,:,1:n3-2,3)
     mm(:,:,2:n3-1)  =mm(:,:,2:n3-1) -(0.5/d3)*4*dd(:,:,1:n3-2,3)
     mm(:,:,3:n3  )  =mm(:,:,3:n3  ) +(0.5/d3)*3*dd(:,:,1:n3-2,3)
  else
     dd(:,:,1:n3-2,3)= dd(:,:,1:n3-2,3) +(0.5/d3)*(mm(:,:,1:n3-2)-4*mm(:,:,2:n3-1)+3*mm(:,:,3:n3))
  end if

  stat=0
  call xgrad_close()
end function




! =============================================
function cross_lop3 (adj,add,mm,dd) result(stat)
  ! 3D cross multiplier function
  logical, intent(in)        :: adj,add
  real,dimension(n1,n2,n3,3) :: mm, dd, gaux
  integer                    :: stat

  call crossnull(adj,add,mm,dd)
  stat=grad_lop3 (.false.,.false.,aux,gaux)

  if (adj) then
     mm(:,:,:,1)=mm(:,:,:,1) +gaux(:,:,:,3)*dd(:,:,:,2) -gaux(:,:,:,2)*dd(:,:,:,3)
     mm(:,:,:,2)=mm(:,:,:,2) -gaux(:,:,:,3)*dd(:,:,:,1) +gaux(:,:,:,1)*dd(:,:,:,3)
     mm(:,:,:,3)=mm(:,:,:,3) +gaux(:,:,:,2)*dd(:,:,:,1) -gaux(:,:,:,1)*dd(:,:,:,2)
  else
     dd(:,:,:,1)=dd(:,:,:,1) -gaux(:,:,:,3)*mm(:,:,:,2) +gaux(:,:,:,2)*mm(:,:,:,3)
     dd(:,:,:,2)=dd(:,:,:,2) +gaux(:,:,:,3)*mm(:,:,:,1) -gaux(:,:,:,1)*mm(:,:,:,3)
     dd(:,:,:,3)=dd(:,:,:,3) -gaux(:,:,:,2)*mm(:,:,:,1) +gaux(:,:,:,1)*mm(:,:,:,2)
  end if

  stat=0
end function


! ----------------------------------------------
subroutine xgrad_close()
!   deallocate(temp)
end subroutine 

! ----------------------------------------------
subroutine gradnull(adj, add, x, y)
  logical, intent (in) :: adj, add
  real  :: x (:,:,:), y (:,:,:,:)
  if (.not. add) then 
     if (adj) then
        x = 0.
     else 		    
        y = 0.
     end if
  end if
end subroutine

! ----------------------------------------------
subroutine crossnull(adj, add, x, y)
  logical, intent (in) :: adj, add
  real  :: x (:,:,:,:), y (:,:,:,:)
  if (.not. add) then 
     if (adj) then
        x = 0.
     else
        y = 0.
     end if
  end if
end subroutine


end module xgrad_mod
