module traceback_mod
  use sep3d_struct_mod
  use basic_util_mod
  use sep
  use put_3d_mod 
  use tau_z_conv_3d_mod
  implicit none
  integer,private,save :: ns(3),ntrs,npts
  real,private,save    :: os(3),ds(3),d_t
  real, pointer,private:: slow(:,:,:),sigmax(:,:,:),sigmay(:,:,:)
  logical, private,save:: tau
  
  contains

  logical function init_traceback_func(slow_tag,sigmax_tag,sigmay_tag,dtt)
    character(len=*) ::  slow_tag,sigmax_tag,sigmay_tag
    type(sep3d)      ::  slow_s,sigma_x,sigmay_s
    optional         ::  sigmax_tag,sigmay_tag 
    integer,external :: sreed
    integer          :: ntau,nt(3)
    real             :: otau,dtau,ot(3),dt(3),dtt
    real,pointer     :: map(:,:,:)
    
    init_traceback_func=.false.

    d_t=dtt
    ntrs=0
    npts=0

    call init_sep3d(slow_tag,slow_s,"INPUT")
    call set_n_o_d_3d(slow_s,ns,os,ds)
    call init_coord(ns,os,ds)
    nt=ns;dt=ds;ot=os


    call from_param("tau",tau,.false.)

    allocate(slow(ns(1),ns(2),ns(3)))
    if(4*size(slow)/=sreed(slow_tag,slow,size(slow)*4)) then
      write(0,*) "trouble reading in slowness"
      return
    end if

    if(tau) then
      slow=1./slow
      call from_param("ntau",ntau);nt(1)=ntau
      call from_param("dtau",dtau);dt(1)=dtau
      call from_param("otau",otau);ot(1)=otau
      allocate(map(ns(1),ns(2),ns(3)))

      if(.not.init_tau_z_conv_3d(ns,os,ds,nt,ot,dt)) then
        write(0,*) "trouble initializing tau conversion"
        return
      end if
      if(.not. create_tau_map_3d(slow,map)) then
        write(0,*) "trouble creating tau map"
        return
      end if
      slow=1./slow
    end if
     
    
    if(tau .and. present(sigmax_tag)) then
      tau=.true.
      allocate(sigmax(ns(1),ns(2),ns(3)))
      if(4*size(sigmax)/=sreed(sigmax_tag,sigmax,size(sigmax)*4)) then
        write(0,*) "trouble reading in sigmax"
        return
      end if
    end if
    if(present(sigmay_tag) .and. ns(3)>1) then
      allocate(sigmay(ns(1),ns(2),ns(3)))
      if(4*size(sigmay)/=sreed(sigmay_tag,sigmay,size(sigmay)*4)) then
        write(0,*) "trouble reading in sigmay"
        return
      end if
    end if
    init_traceback_func=.true.
  end function

  integer function tb_length(ray)
    real :: ray(:,:)
    real :: last
    
    if(any(ray(:,size(ray,2)-1) /= ray(:,size(ray,2))))  then
      tb_length=size(ray,2)
    else
      tb_length=1; 
      do while(.not.(all(ray(:,tb_length)==ray(:,tb_length+1) .and. ray(3,tb_length)< .01)))
        tb_length=tb_length+1
      end do
    end if
  end function

  logical function calc_center(ray_in,center)
    real ray_in(:,:),center(:,:)
    calc_center=.false.
    
    center(:,1:size(ray_in,2)-1)=0.5*(ray_in(:,2:size(ray_in,2)-1)&
       +ray_in(:,1:size(ray_in,2)-1))
    calc_center=.true.
  end function

  logical function calc_time(center,delta,time)
    real :: center(:,:)
    real :: delta(:,:),time(:)
    real    :: s,sx,sy
    integer ::  i
    calc_time=.false.

    i=1
    if(.not. tau) then
      time(i)=delta(4,i)*get_val_lin_pos(slow,center(3,i),&
         center(1,i),center(2,i))
      i=size(time)-1
      time(i)=delta(4,1)*get_val_lin_pos(slow,center(3,i),&
        center(1,i),center(2,i))
    else if(ns(3)==1) then
     s=get_val_lin_pos(slow,center(3,i),center(1,i),center(2,i))
      sx=get_val_lin_pos(sigmax,center(3,i),center(1,i),center(2,i))
      time(i)=sqrt((s*delta(1,i))**2+(s*delta(2,i))**2+&
        .25*(delta(3,i)-sx*delta(1,i))**2)
      i=size(center,2)
      s=get_val_lin_pos(slow,center(3,i),center(1,i),center(2,i))
      sx=get_val_lin_pos(sigmax,center(3,i),center(1,i),center(2,i))
      time(i)=sqrt((s*delta(1,i))**2+(s*delta(2,i))**2+&
        .25*(delta(3,i)-sx*delta(1,i))**2)
    else
      s=get_val_lin_pos(slow,center(3,i),center(1,i),center(2,i))
      sx=get_val_lin_pos(sigmax,center(3,i),center(1,i),center(2,i))
      sy=get_val_lin_pos(sigmay,center(3,i),center(1,i),center(2,i))
      time(i)=sqrt((s*delta(1,i))**2+(s*delta(2,i))**2+&
          .25*(delta(3,i)-sx*delta(1,i)-sy*delta(2,i))**2)
      i=size(center,2)
      s=get_val_lin_pos(slow,center(3,i),center(1,i),center(2,i))
      sx=get_val_lin_pos(sigmax,center(3,i),center(1,i),center(2,i))
      sy=get_val_lin_pos(sigmay,center(3,i),center(1,i),center(2,i))
      time(i)=sqrt((s*delta(1,i))**2+(s*delta(2,i))**2+&
        .25*(delta(3,i)-sx*delta(1,i)-sy*delta(2,i))**2)
    end if 
    time(2:size(center,2)-1)=d_t
    calc_time=.true.
  end function    

  logical function get_ray_prop(center,out,field)
     real :: center(:,:),field(:,:,:),out(:)
     integer :: i
     get_ray_prop=.false.

     do i=1,size(center,2)
       out(i)=get_val_lin_pos(field,center(3,i),center(1,i),center(2,i))
     end do
     get_ray_prop=.true.

  end function
  logical function calc_delta(ray_in,delta)
    real ray_in(:,:),delta(:,:)
    calc_delta=.false.
    
    delta(1:3,1:size(ray_in,2)-1)=(ray_in(:,2:size(ray_in,2)-1)-&
      ray_in(:,1:size(ray_in,2)-1))
    delta(1:3,size(ray_in,2))=delta(1:3,size(ray_in,2)-1)
    delta(4,:)=sqrt(sum(delta(1:3,:)**2,1))
    calc_delta=.true.
  end function

  logical function trace_back_depth(ray,ray_s,scale,iloc,fract,n1)
    integer :: n1,j,ray_type,iloc(:,:),i
    real :: ray(:,:),ray_s
    real :: scale(:,:),fract(:,:)
    real  :: center(3,size(ray,2)),delta(4,size(ray,2))!,fract(3,n1)
    real  :: sl(size(ray,2)),sx(size(ray,2)),sy(size(ray,2))
    real  :: time(size(ray,2)),dsig(size(ray,2))

    trace_back_depth=.false.

    n1=tb_length(ray)

    if(.not. calc_center(ray(:,:n1),center(:,:n1))) then
      write(0,*) "        ERROR: calculating the center"
      return
    end if

    if(.not. calc_delta(ray(:,:n1),delta)) then
      write(0,*) "        ERROR: calculating the delta"
      return
    end if
    n1=n1-1

    if(.not. calc_iloc_fract(center(:,:n1),ray_s,iloc(:,:n1),fract(:,:n1),&
      scale(:,:n1))) then
      write(0,*) "        ERROR: calculating iloc and fract"
      return
    end if
    scale(1,:n1)=scale(1,:n1)*delta(4,1:n1)
    if(n1<3)  scale(:,:n1)=0.
    trace_back_depth=.true.
  end function



  logical function trace_back_tau(ray,ray_s,scale,iloc,fract,n1)
    integer :: n1,j,iloc(:,:),i
    real :: ray(:,:),ray_s
    real :: scale(:,:),fract(:,:)
    real  :: center(3,size(ray,2)),delta(4,size(ray,2))!,fract(3,n1)
    real  :: sl(size(ray,2)),sx(size(ray,2)),sy(size(ray,2))
    real  :: time(size(ray,2)),dsig(size(ray,2))

    trace_back_tau=.false.

    n1=tb_length(ray)

    call convert_ray_traj_z_tau_3d(ray(:,1:n1))

    if(.not. calc_center(ray(:,:n1),center(:,:n1))) then
      write(0,*) "        ERROR: calculating the center"
      return
    end if

    if(.not. calc_delta(ray(:,:n1),delta)) then
      write(0,*) "        ERROR: calculating the delta"
      return
    end if
    n1=n1-1

    if(.not. calc_iloc_fract(center(:,:n1),ray_s,iloc(:,:n1),fract(:,:n1),&
      scale(:,:n1))) then
      write(0,*) "        ERROR: calculating iloc and fract"
      return
    end if

    if(.not. get_ray_prop(center(:,1:n1),sl(1:n1),slow))then
       write(0,*) "        ERROR: grabbing slowness along ray";return
    end if
    if(.not. calc_time(center(:,1:n1),delta(:,1:n1),time)) then
       write(0,*) "        ERROR: grabbing dt along ray";return
    end if
    if(.not. get_ray_prop(center(:,1:n1),sx(1:n1),sigmax)) then
       write(0,*) "        ERROR: grabbing sigmax along ray";return
    end if
    if(ns(3)==1) then
      scale(1,1:n1)=delta(1,:n1)**2*sl(:n1)/ time(:n1)*scale(1,:n1)
      scale(2,1:n1)=-.25*delta(1,:n1)*(delta(3,:n1)-delta(1,:n1)*sx(:n1))/&
        time(:n1)*scale(2,:n1)
      do i=1,n1
        if(time(i) < .0001) scale(:,i)=0.
      end do
      scale(3,1:n1)=0.
    else
      if(.not. get_ray_prop(center(:,1:n1),sy(1:n1),sigmay)) then
         write(0,*) "        ERROR: grabbing slowness along ray";return
      end if
      scale(1,1:n1)=&
        (delta(1,:n1)**2+delta(2,:n1)**2)*sl(:n1)/time(:n1)*scale(1,1:n1)
      dsig(1:n1) =-.25*(delta(3,:n1)-delta(1,:n1)&
       *sx(:n1)-delta(2,:n1)*sy(:n1))/time(:n1)
      scale(2,1:n1)=delta(1,:)*dsig(1:n1)*scale(2,1:n1)
      scale(3,1:n1)=delta(2,:)*dsig(1:n1)*scale(3,1:n1)
      do i=1,n1
        if(time(i) < .0001) scale(:,i)=0.
      end do
    end if
    if(n1<3)  scale(:,:n1)=0.
    trace_back_tau=.true.
  end function



  logical function calc_iloc_fract(center,ray_s,iloc,fract,scale) 
    real :: center(:,:),ray_s
    integer :: iloc(:,:)
    real    :: fract(:,:)
    real    :: scale(:,:)
    integer :: i

    calc_iloc_fract=.false.
    
    fract(1,:)=(center(3,:)-os(1))/ds(1)+1.
    fract(2,:)=(center(1,:)-os(2))/ds(2)+1.
    fract(3,:)=(center(2,:)-os(3))/ds(3)+1.
    scale(1,:)=ray_s
    scale(2,:)=ray_s
    scale(3,:)=ray_s
    iloc=fract
    fract=fract-iloc
    if(ns(3)/=1) then
      do i=1,size(fract,2)
        if(iloc(1,i) < 1 .or. iloc(1,i) >= ns(1) .or. &
          iloc(2,i) < 1 .or. iloc(2,i) >= ns(2) .or. &
          iloc(3,i) < 1 .or. iloc(3,i) >= ns(3) )then
            fract(:,i)=.0
            scale(:,i)=0. !OUTSIDE MODEL SPACE SCALE TO 0
            iloc(:,i)=1
         end if
      end do
    else
      do i=1,size(fract,2)
        if(iloc(1,i) < 1 .or. iloc(1,i) >= ns(1) .or. &
          iloc(2,i) < 1 .or. iloc(2,i) >= ns(2) ) then
            fract(1:2,i)=.0
            scale(1:2,i)=0. !OUTSIDE MODEL SPACE SCALE TO 0
            iloc(1:2,i)=1
         end if
      end do
    end if
    calc_iloc_fract=.true.
  end function
  subroutine add_traceback(scale,iloc,fract,nh,n1,tr_scale,ipos,spos,var,glob)
    integer :: nh,n1(:)
    real    :: scale(:,:,:),fract(:,:,:),tr_scale(:),spos(:),var(:)
    real :: aa
    integer :: iloc(:,:,:),iray,ipos(:),i,ntot
    logical :: glob


   if(.not. glob) then
     call srite("var",var,nh*4)
     call srite("err",tr_scale,nh*4)
   else
     call srite("var",var,4)
     call srite("err",tr_scale,4)
   end if
    aa=sum(tr_scale(1:nh))
    i=0
    if(.not. tau) then
       if(ns(3)==1) then
          call add_2d(nh,n1,i,iloc,fract,1,scale,ipos,spos)
       else 
          call add_3d(nh,n1,i,iloc,fract,1,scale,ipos,spos)
       endif
       call srite("iloc",ipos,i*4)
       call srite("scale",spos,i*4)
    else
       if(ns(3)==1) then
          call add_2d(nh,n1,i,iloc,fract,2,scale,ipos,spos)
          call srite("scale_sigx",spos,i*4)
          call add_2d(nh,n1,i,iloc,fract,1,scale,ipos,spos)
       else 
          call add_3d(nh,n1,i,iloc,fract,3,scale,ipos,spos)
          call srite("scale_sigy",spos,i*4)
          call add_3d(nh,n1,i,iloc,fract,2,scale,ipos,spos)
          call srite("scale_sigx",spos,i*4)
          call add_3d(nh,n1,i,iloc,fract,1,scale,ipos,spos)
       endif
       call srite("iloc",ipos,i*4)
       call srite("scale",spos,i*4)
    end if
    if(ns(3)==1) then 
      n1=n1*4
    else
      n1=n1*8
    end if
    ntot=sum(n1(1:nh))
   if(.not. glob) then
     ntrs=ntrs+nh
     call srite("nlen",n1,nh*4)
   else
     call srite("nlen",ntot,4)
     ntrs=ntrs+1
   end if
    npts=npts+i

  end subroutine



  subroutine add_3d(nh,n1,i,iloc,fract,iscale,scale,ipos,spos)
    integer :: ih,i1,nh,n1(:),iloc(:,:,:),ipos(:),i,iscale
    real    :: fract(:,:,:),scale(:,:,:),spos(:)
    i=0
    do ih=1,nh
      do i1=1,n1(ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih)-1)+ns(1)*ns(2)*&
             (iloc(3,i1,ih)-1)
         spos(i)=(1.-fract(1,i1,ih))*(1.-fract(2,i1,ih))*(1.-fract(3,i1,ih))&
           *scale(iscale,i1,ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih)-1)+1+ns(1)*ns(2)*&
           (iloc(3,i1,ih)-1)
         spos(i)=(fract(1,i1,ih))*(1.-fract(2,i1,ih))*(1.-fract(3,i1,ih))&
           *scale(iscale,i1,ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih))+ns(1)*ns(2)*&
            (iloc(3,i1,ih)-1)
         spos(i)=(1.-fract(1,i1,ih))*(fract(2,i1,ih))*(1.-fract(3,i1,ih))&
           *scale(iscale,i1,ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih))+1+ns(1)*ns(2)*&
           (iloc(3,i1,ih)-1)
         spos(i)=(fract(1,i1,ih))*(fract(2,i1,ih))*(1.-fract(3,i1,ih))&
           *scale(iscale,i1,ih)
         i=i+1

         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih)-1)+ns(1)*ns(2)*&
           (iloc(3,i1,ih))
         spos(i)=(1.-fract(1,i1,ih))*(1.-fract(2,i1,ih))*(fract(3,i1,ih))&
           *scale(iscale,i1,ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih)-1)+1+ns(1)*ns(2)*&
           (iloc(3,i1,ih))
         spos(i)=(fract(1,i1,ih))*(1.-fract(2,i1,ih))*(fract(3,i1,ih))&
           *scale(iscale,i1,ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih))+ns(1)*ns(2)*(iloc(3,i1,ih))
         spos(i)=(1.-fract(1,i1,ih))*(fract(2,i1,ih))*(fract(3,i1,ih))&
           *scale(iscale,i1,ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih))+1+ns(1)*ns(2)*&
           (iloc(3,i1,ih))
         spos(i)=(fract(1,i1,ih))*(fract(2,i1,ih))*(fract(3,i1,ih))&
           *scale(iscale,i1,ih)
      end do
    end do
  end subroutine

  subroutine add_2d(nh,n1,i,iloc,fract,iscale,scale,ipos,spos)
    integer :: ih,i1,nh,n1(:),i,iloc(:,:,:),ipos(:),iscale
    real    :: fract(:,:,:),scale(:,:,:),spos(:)
    do ih=1,nh
      do i1=1,n1(ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih)-1)
         spos(i)=(1.-fract(1,i1,ih))*(1.-fract(2,i1,ih))*scale(iscale,i1,ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih)-1)+1
         spos(i)=(fract(1,i1,ih))*(1.-fract(2,i1,ih))*scale(iscale,i1,ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih))
         spos(i)=(1.-fract(1,i1,ih))*(fract(2,i1,ih))*scale(iscale,i1,ih)
         i=i+1
         ipos(i)=iloc(1,i1,ih)+ns(1)*(iloc(2,i1,ih))+1
         spos(i)=(fract(1,i1,ih))*(fract(2,i1,ih))*scale(iscale,i1,ih)
      end do
    end do
  end subroutine


  subroutine finish_tracebacks()
    call to_history("n1",ntrs,"var")
    call to_history("n1",ntrs,"err")
    call to_history("n1",ntrs,"nlen")
    call to_history("n1",npts,"iloc")
    call to_history("n1",npts,"scale")
    if(tau) then
      call to_history("n1",npts,"scale_sigx")
      if(ns(3)/=1) call to_history("n1",npts,"scale_sigy")
    end if
  end subroutine
end module
