!-----------------------------------------------------------------------
! DESCRIPTION
!	program to compute the traveltimes of primaries
!	and specularly-reflected multiples from a two flat-layer model.
!------------------------------------------------------------------
!	Input:
!		zwb:     depth of the flat water bottom (m)
!		zref:    depth of the flat reflector (m)
!		vw:      water velocity (m/s)
!		vr:      sediment velocity (m/s)
!		nh:      number of offsets
!		oh:      minimum offset (m)
!		dh:      offset sampling interval (m)
!		ncmp:    number of cmps
!		ocmp:    minimum cmp (m)
!		dcmp:    cmp sampling interval (m)
!		nt:      number of samples in output traces
!		dt:      time sampling interval (s)
!		fpeak:   peak frequency of ricker wavelet (Hz)
!		rho1:	density of water (g/cc)
!		rho2:	density of sediment layer (g/cc)
!		rho3:	density of half-space below sediments (g/cc)
!		amps:	=1 to compue amplitudes. =0 otherwise 
!
!	Output:
!		SEPFILE cube (nt, nhD, ncmp) of primary and multiple
!		reflections from the flat two-layer model.
!
!	Examples:
!		two_flat_layers.x zwb=500 zref=800 vw=1500 vr=2500 nh=100 \
!		oh=100 dh=20 ncmp=1 ocmp=2000 dcmp=20 nt=800 dt=0.004 \
!		fpeak=20 rho1=1. rho2=2. rho3=2.5 amps=y > two_flat_layer_cmps.H
!---------------------------------------------------------------------
program two_flat_layers

	use sep
	use auxsu
	!use solve_for

	implicit none

	! input variables
	real	:: zwb 
	real	:: zref
	real	:: vw
	real	:: vr
	integer	:: nh
	real	:: oh
	real	:: dh
	integer	:: ncmp
	real	:: ocmp
	real	:: dcmp
	integer	:: nt
	real	:: dt
	real	:: fpeak
	real	:: rho1
	real	:: rho2
	real	:: rho3
	character(len=1) :: prim1
	character(len=1) :: prim2
	character(len=1) :: wb1
	character(len=1) :: wb2
	character(len=1) :: wb3
	character(len=1) :: sm1
	character(len=1) :: im1
	character(len=1) :: plg1
	character(len=1) :: amps

	! internal variables
	integer	:: icmp, ih, lw, iw
	integer	:: it,itt
	real	:: x1, x121, x101, x1010101, x1210121, x12121, x12101
	real	:: y1, y121, y101, y1010101, y1210121, y12121, y12101
	real	:: cmp, hh, rho, zwb2, hh2
	real	:: thref, thref2
	real	:: solve_for_x1
	real	:: rc_comp, tc_comp
	real	:: rc1,rc2,rc3,rc4,rc5,rc6,rc7,tc1,tc2,tc3,tc4,alpha,vhs
	!real,external					:: solve_for_x1
	real,dimension(8)				:: tt=0.,rc=1.		!! reflection coefficients
	real,allocatable,dimension(:)	:: temp,trace
	real,dimension(:),pointer 		:: wavelet

	! initialize SEP stuff
	call initpar()

	!! get processing parameters and set their defaults
	call from_param("zwb"  ,zwb  , 500.)
	call from_param("zref" ,zref , 800.)
	call from_param("vw"   ,vw   ,1500.)
	call from_param("vr"   ,vr   ,2500.)
	call from_param("vhs"  ,vhs  ,3000.)
	call from_param("rhow" ,rho1 ,   1.)
	call from_param("rhor" ,rho2 ,   2.)
	call from_param("rhohs",rho3 ,  2.5)
	call from_param("nh"   ,nh   ,    1)
	call from_param("oh"   ,oh   , 100.)
	call from_param("dh"   ,dh   ,  20.)
	call from_param("ncmp" ,ncmp ,    1)
	call from_param("ocmp" ,ocmp , 100.)
	call from_param("dcmp" ,dcmp ,  20.)
	call from_param("nt"   ,nt   ,  800)
	call from_param("dt"   ,dt   ,0.004)
	call from_param("fpeak",fpeak,  20.)
	call from_param("prim1",prim1,  'y')
	call from_param("prim2",prim2,  'y')
	call from_param("wb1",wb1,  'y')
	call from_param("wb2",wb2,  'y')
	call from_param("wb3",wb3,  'y')
	call from_param("sm1",sm1,  'y')
	call from_param("im1",im1,  'y')
	call from_param("plg1",plg1,  'y')
	call from_param("amps",amps,  'y')

	!! write SEP parameters to history file
	call to_history("n1"        ,      nt)
	call to_history("n2"        ,      nh)
	call to_history("n3"        ,    ncmp)
	call to_history("o1"        ,      0.)
	call to_history("o2"        ,      oh)
	call to_history("o3"        ,    ocmp)
	call to_history("d1"        ,      dt)
	call to_history("d2"        ,      dh)
	call to_history("d3"        ,    dcmp)
	call to_history("label1"    ,  "time")
	call to_history("label2"    ,"offset")
	call to_history("label3"    ,   "cmp")
	call to_history("esize"     ,       4)

	!! compute auxiliary variables
	rho=vr/vw
	zwb2=zwb**2
	thref=zref-zwb
	thref2=thref**2
	
	!! generate ricker wavelet
	call rickersu(fpeak,dt,lw,iw,wavelet)
	write(0,*)'DESPUES DE RICKER'

	!! allocate working arrays
    allocate(trace(nt),temp(nt))

	!! loop over cmps
	do icmp=1,ncmp
	
		! set coordinate of this cmp
		cmp=ocmp+(icmp-1)*dcmp
write(0,*)'Processing cmp',cmp

		! loop over offsets
		do ih=1,nh

			! set half offset of current trace
			hh=(oh+(ih-1)*dh)/2.
			hh2=hh**2

			! compute traveltime of water-bottom primary
			if (prim1=='y') then
				tt(1)=2/vw*sqrt(hh2+zwb2)
				if (amps=='y') then 
					alpha=atan(hh/zwb)
					rc(1)=rc_comp(vw,rho1,vr,rho2,alpha)
				else
					rc(1)=1.
				end if
			end if
!write(0,*)'cmp=',cmp,'hh=',hh,'t1=',tt(1)

			!
			! compute the traveltime of the primary from the sediment reflector
			! the 121 refers to bounces at reflectors 1, 2 and 1
			!
			if (prim2=='y') then
				! first reflection point at the water-bottom
				if (hh==0) then
					x121=0
				else
					x121=solve_for_x1(rho,hh,zwb,thref,1)	
				end if
				! second reflection point at the water-bottom
				y121=hh-x121			
				! compute the traveltime
				tt(2)=2/vw*sqrt(x121**2+zwb2)+2/vr*sqrt(y121**2+thref2)
!write(0,*)'x121=',x121,'y121=',y121,'t121=',tt(2)
				if (amps=='y') then
					alpha=atan(x121/zwb)
					tc1=tc_comp(vw,rho1,vr,rho2,alpha)
					alpha=atan((hh-x121)/thref)
					rc1=rc_comp(vr,rho2,vhs,rho3,alpha)
					tc2=tc_comp(vr,rho2,vw,rho1,alpha)
					rc(2)=tc1*rc1*tc2
				else
					rc(2)=1.
				end if
			end if

			!
			! compute the traveltime of the first, second and third order WB multiples
			!
			if (wb1=='y') then
				tt(3)=4/vw*sqrt((hh/2.)**2+zwb2)
				if (amps=='y') then
					alpha=atan(hh/(2*zwb))
					rc1=rc_comp(vw,rho1,vr,rho2,alpha)
					rc2=-1
					rc3=rc1
					rc(3)=rc1*rc2*rc3
				else 
					rc(3)=-1.
				end if
			end if
			if (wb2=='y') then 
				tt(4)=6/vw*sqrt((hh/3.)**2+zwb2)
				if (amps=='y') then
					alpha=atan(hh/(3*zwb))
					rc1=rc_comp(vw,rho1,vr,rho2,alpha)
					rc2=-1
					rc3=rc1
					rc4=-1
					rc5=rc1
					rc(4)=rc1*rc2*rc3*rc4*rc5
				else
					rc(4)=1.
				end if
			end if
			if (wb3=='y') then
				tt(5)=8/vw*sqrt((hh/4.)**2+zwb2)
				if (amps=='y') then
					alpha=atan(hh/(4*zwb))
					rc1=rc_comp(vw,rho1,vr,rho2,alpha)
					rc2=-1
					rc3=rc1
					rc4=-1
					rc5=rc1
					rc6=-1
					rc7=rc1
					rc(5)=rc1*rc2*rc3*rc4*rc5*rc6*rc7
				else
					rc(5)=-1.
				end if
			end if
!write(0,*)'t101=',tt(3),'t1010101=',tt(4)

			!
			! compute the traveltime of the first order surface multiple of the 
			! second layer
			!
			if (sm1=='y') then
				! first reflection point at the water-bottom
				if (hh==0) then
					x1210121=0
				else
					x1210121=solve_for_x1(rho,hh/2,zwb,thref,1)	
				end if
				y1210121=hh/2-x1210121
				tt(6)=4/vw*sqrt(x1210121**2+zwb2)+4/vr*sqrt(y1210121**2+thref2)
!write(0,*)'x1210121=',x1210121,'y1210121=',y1210121,'t1210121=',tt(5)
				if (amps=='y') then
					alpha=atan(x1210121/zwb)
					tc1=tc_comp(vw,rho1,vr,rho2,alpha)
					alpha=atan(hh/(2*thref))
					rc1=rc_comp(vr,rho2,vhs,rho3,alpha)
					tc2=tc_comp(vr,rho2,vw,rho1,alpha)
					rc2=-1
					tc3=tc1
					rc3=rc1
					tc4=tc2
					rc(6)=tc1*rc1*tc2*rc2*tc3*rc3*tc4
				else
					rc(6)=1.
				end if
			end if

			!
			! compute the traveltime of the first order internal multiple 
			! of the second layer
			!
			if (im1=='y') then
				! first reflection point at the water-bottom
				if (hh==0) then
					x12121=0
				else
					x12121=solve_for_x1(rho,hh,zwb,2*thref,1)
				end if
				y12121=(hh-x12121)/2.
				tt(7)=2/vw*sqrt(x12121**2+zwb2)+4/vr*sqrt(y12121**2+thref2);
!write(0,*)'x12121=',x12121,'y12121=',y12121,'t12121=',tt(6)
				if (amps=='y') then
					alpha=atan(x12121/zwb)
					tc1=tc_comp(vw,rho1,vr,rho2,alpha)
					alpha=atan((hh-x12121)/(2*thref))
					rc1=rc_comp(vr,rho2,vhs,rho3,alpha)
					rc2=rc_comp(vr,rho2,vw,rho1,alpha)
					rc3=rc1
					tc2=tc_comp(vr,rho2,vw,rho1,alpha)
					rc(7)=tc1*rc1*rc2*rc3*tc2
				else
					rc(7)=1.
				end if
			end if

			!
			! compute the traveltime of the peg-leg multiple of the second layer
			!
			if (plg1=='y') then
				! first reflection point at the water-bottom
				if (hh==0) then
					x12101=0
				else
					x12101=solve_for_x1(rho,hh,zwb,thref,2)
				end if
				y12101=(hh-2*x12101)
				tt(8)=4/vw*sqrt(x12101**2+zwb2)+2/vr*sqrt(y12101**2+thref2)
!write(0,*)'x121=',x121,'x12121=',x12121,'x1210121=',x1210121,'x12101=',x12101,'hh=',hh
!write(0,*)'rho=',rho,'zwb=',zwb,'thref=',thref
				if (amps=='y') then
					alpha=atan(x12101/zwb)
					tc1=tc_comp(vw,rho1,vr,rho2,alpha)
					alpha=asin(vr*sin(alpha)/vw)
					rc1=rc_comp(vr,rho2,vhs,rho3,alpha)
					tc2=rc_comp(vr,rho2,vw,rho1,alpha)
					alpha=asin(vw*sin(alpha)/vr)
					rc2=-1
					rc3=rc_comp(vw,rho1,vr,rho2,alpha)
					rc(8)=tc1*rc1*tc2*rc2*rc3
				else
					rc(8)=1.
				end if
			end if

			! convolve spikes with a Ricker wavelet
			temp=0.
			do it=1,size(tt)
				if (tt(it)==0.) cycle
				itt=NINT(tt(it)/dt)
!write(0,*)'it=',it,'itt=',itt
				if (itt>0.AND.itt<=nt) then
            		temp(itt)=temp(itt)+rc(it)
            		!temp(itt)=temp(itt)+1
				else
					write(0,*)'WARNING, it>nt for cmp=',cmp,'hh=',hh
				end if
			end do
			call convsu(lw,iw+1,wavelet,nt,1,temp,nt,1,trace)

			! output this trace
            call srite("out",trace,4*nt)
		end do
	end do

	! deallocate working space
    deallocate(wavelet,trace,temp)
	
	call exit (0)
end program two_flat_layers

real function rc_comp(v1,rho1,v2,rho2,alpha)
	!---------------------------------------------------------
	! Function to compute angle-dependent reflection coefficient
	!----------------------------------------------------------

	implicit none

	! argument variables
	real, intent(in)	:: v1 !! velocity in incident medium
	real, intent(in)	:: v2 !! velocity in transmitted medium
	real, intent(in)	:: rho1 !!density in incident medium
	real, intent(in)	:: rho2 !!density in transmitted medium
	real, intent(in)	:: alpha !! incidence angle

	!! internal variables
	real				:: s_beta,c_beta
	s_beta=(v2*sin(alpha)/v1)
	c_beta=sqrt(1-s_beta**2)

	rc_comp=(rho2*v2/c_beta-rho1*v1/cos(alpha))/(rho2*v2/c_beta+rho1*v1/cos(alpha))

	return
end function rc_comp

real function tc_comp(v1,rho1,v2,rho2,alpha)
	!---------------------------------------------------------
	! Function to compute angle-dependent transmission coefficient
	!----------------------------------------------------------

	implicit none

	! argument variables
	real, intent(in)	:: v1 !! velocity in incident medium
	real, intent(in)	:: v2 !! velocity in transmitted medium
	real, intent(in)	:: rho1 !!density in incident medium
	real, intent(in)	:: rho2 !!density in transmitted medium
	real, intent(in)	:: alpha !! incidence angle

	!! internal variables
	real				:: s_beta,c_beta
	s_beta=(v2*sin(alpha)/v1)
	c_beta=sqrt(1-s_beta**2)

	tc_comp=(2*rho2*v2/c_beta)/(rho2*v2/c_beta+rho1*v1/cos(alpha))

	return
end function tc_comp

real function solve_for_x1(rho,hh,z1,z2,fn)
	
	!----------------------------------------------------------
	! Function to solve a quartic polynomial equation by the 
	! method of bisection. The solution is the distance between
	! the shot and the first bounce of the multiple at the WB.
	! The equation is:
   	! (rho^2-1)x1^4-2(rho^2-1)hh*x1^3+[(rho^2-1)hh^2+rho^2*
	!	z2^2-z1^2]x1^2+2hh*z1^2*x1-z1^2*hh^2=0
	! Ax1^4+Bx1^3+Cx1^2+Dx1+E=0
	!----------------------------------------------------------

	implicit none

	! argument variables
	real, intent(in)	:: rho ! v2/v1
	real, intent(in)	:: hh  ! half-offset
	real, intent(in)	:: z1	! depth of water bottom
	real, intent(in)	:: z2	! depth of sediement reflector
	integer,intent(in)	:: fn	! =1 for multiples symmetric wrt the cmp
								! /=1 for other multiples
                                                                                     
	! internal variables
	integer	:: k
	real	:: A, B, C, D, E	! function coefficients
	real	:: x1, x2, x3		! trial values
	real	:: f1, f2, f3		! trial function values
	real	:: hh_scl,z1_scl,z2_scl

	! scale the input values by 1000
	hh_scl=hh/1000.
	z1_scl=z1/1000.
	z2_scl=z2/1000.

	! compute function coefficients
	if (fn==1) then
		A=(rho**2-1)
		B=-2*(rho**2-1)*hh_scl
		C=(rho**2-1)*hh_scl**2+rho**2*z2_scl**2-z1_scl**2
		D=2*hh_scl*z1_scl**2
		E=-z1_scl**2*hh_scl**2
	else if (fn==2) then
		A=4*(rho**2-1)
		B=-4*(rho**2-1)*hh_scl
		C=(rho**2-1)*hh_scl**2+rho**2*z2_scl**2-4*z1_scl**2
		D=4*hh_scl*z1_scl**2
		E=-z1_scl**2*hh_scl**2
	else
		write(0,*)'ERROR: fn must be 1 or 2'
		call exit(9)
	end if

!write(0,*)'A=',A,'B=',B,'C=',C,'D=',D,'E=',E

	! set initial guesses
	x1=0;
	x2=hh_scl;
	x3=(x1+x2)/2.;
	
	f1=A*x1**4+B*x1**3+C*x1**2+D*x1+E;
	f2=A*x2**4+B*x2**3+C*x2**2+D*x2+E;
	f3=A*x3**4+B*x3**3+C*x3**2+D*x3+E;
!write(0,*)'f1=',f1,'f2=',f2,'f3=',f3
                                                                                     
	! loop until solution found
	k=1;
	do
		f1=A*x1**4+B*x1**3+C*x1**2+D*x1+E;
		f2=A*x2**4+B*x2**3+C*x2**2+D*x2+E;
		f3=A*x3**4+B*x3**3+C*x3**2+D*x3+E;
		if (abs(f3)<0.00001) then
			solve_for_x1=x3*1000.
!write(0,*)'x1=',x1*1000.,'x2=',x2*1000.,'x3=',x3*1000.
!write(0,*)'f1=',f1,'f2=',f2,'f3=',f3
			return
       	end if
       	if (f1*f3<0) then
			x2=x3;
		else
			x1=x3;
		end if
		x3=(x1+x2)/2;

		if (k>100) then
			write(0,*)'ERROR: could not solve for x1 in 1000 iters'
			return;
		end if
		k=k+1;
	end do
	return
end function solve_for_x1

