!! module to include cost function for Mauro's well log inversion problem

module ga_Mauro_function
	use auxsu
	use sep
	use Savitzky_Golay

	implicit none

	real, parameter :: FT2M=0.3048		! to convert ft to meters
	
	contains

	!-----------------------------------------------------------------------
	subroutine func(npopsiz,nparam,j,funcval,parent,flag)
	!-----------------------------------------------------------------------
	! This function evaluates the match between a seismic trace and its
	! possible reconstructions using a simple sunthetic seismogram
	! modeling program
	!-----------------------------------------------------------------------

		!! subroutine arguments
		integer, intent(in) :: npopsiz
		integer, intent(in) :: nparam
		integer, intent(in) :: j
		real, intent(out)   :: funcval
		integer, intent(in), optional :: flag
		real, dimension(nparam,npopsiz), intent(in) :: parent

		!! local variables
		integer :: nlyr,option,stat,nt,norm,count,k
		real    :: dt,freq,rho0,dz,tmax
		real,dimension(:),allocatable :: trace
		real,dimension(:),allocatable :: trace0
		real,dimension(:),allocatable :: vel,rho,deltaz

		!! save variables between subroutine calls (make them static)
		save count,trace0
		save nlyr,freq,option,dt,norm,rho0,dz

		! read in log data for reference trace. This trace will be obtained by
		! simple normal incidence plane-plane wave modeling. 
		! The input logs have been "blocked" using Backus average.
		! This is done only the first time the subroutine is called.
		if (count==0) then

			! read in processing flag and parameters associated to the computation
			! of the fitness function
			call from_param("option",option,1)
			call from_param("freq",freq,30.)
			call from_param("dt",dt,0.004)
			call from_param("norm",norm,2)
			call from_param("rho0",rho0,2.)
			call from_param("dz",dz,10.)
write(0,*)'option=',option,'freq=',freq,'dt=',dt,'norm=',norm,'rho0=',rho0,'dz=',dz

			! compute number of layers
			if (option==1) then
				nlyr=nparam
			else if (option==2) then
				nlyr=nparam/2
			else if (option==3) then
				nlyr=nparam/3
			else
				write(0,*)'/********WARN: option has to be either 1,2 or 3. &
					& Resetting to 1**********/'
				option=1
				nlyr=nparam
			end if

			! allocate space for well log arrays
			allocate(vel(nlyr),rho(nlyr),deltaz(nlyr))

			! open file containing the well log data
m_param("option",option,1)
            call from_param("freq",freq,30.)
            call from_param("dt",dt,0.004)
            call from_param("norm",norm,2)
            call from_param("rho0",rho0,2.)
            call from_param("dz",dz,10.)

			open(unit=36,file='layer.txt',status='old',action='read',iostat=stat)
			if (stat/=0) then
				write(0,*)'ERROR: could not open file layer.txt'
				stop
			end if

			! read in log data
			do k=1,nlyr
				if (option==1) then
					!read(36,'E9.7E3') vel(k)
					read(36,*) vel(k)
					rho(k)=rho0
					deltaz(k)=dz*FT2M
				else if (option==2) then
					!read(36,'2E9.7E3') vel(k),rho(k)
					read(36,*) vel(k),rho(k)
					deltaz(k)=dz*FT2M
				else
					!read(36,'3E9.7E3') vel(k),rho(k),deltaz(k)
					read(36,*) vel(k),rho(k),deltaz(k)
					deltaz(k)=deltaz(k)*FT2M
				end if
			end do

			! output original sonic log as a history file
			call to_history("n1",nlyr,"orig_vels")
			call to_history("n2",1,"orig_vels")
			call to_history("n3",1,"orig_vels")
			call to_history("o1",0.,"orig_vels")
			call to_history("esize",4,"orig_vels")
			call srite("orig_vels",vel,4*size(vel))

			!! compute number of time samples in reference seismogram
			tmax=2.0*deltaz(1)/vel(1)
			do k=2,nlyr
				tmax=tmax+2.0*deltaz(k)/vel(k)
			end do
			nt=tmax/dt
			write(0,*)'nt=',nt	

			!! allocate space for reference trace
			allocate(trace0(nt))

			! create reference seismic trace
			call create_trace(nt,nlyr,dt,freq,vel,rho,deltaz,trace0)

			! output reference trace as a history file
			call to_history("n1",nt,"ref_trace")
			call to_history("n2",1,"ref_trace")
			call to_history("n3",1,"ref_trace")
			call to_history("o1",0.,"ref_trace")
			call to_history("d1",dt,"ref_trace")
			call to_history("esize",4,"ref_trace")
			call srite("ref_trace",trace0,4*size(trace0))

			deallocate(vel,rho,deltaz)

		end if
					
		! allocate space for current seismic trace (current individual)
		allocate(trace(nt))

		!! if the best solution has already been found, we only need to
		!! output the corresponding trace and clean up
		if (present(flag)) then

			call synthetic_seismogram(j,option,npopsiz,nparam,nlyr,&
        		& freq,nt,dt,dz,parent,trace,flag)

			deallocate(trace,trace0)
			return

		end if

		!! compute sinthetic seismogram for current individual
		call synthetic_seismogram(j,option,npopsiz,nparam,nlyr,&
        	& freq,nt,dt,dz,parent,trace)

		!! compute fitness value via L1 or L2 norm of the sample-to-sample
		!! difference between current and reference trace
		funcval=0.0
		do k=1,nt
			if (norm==1) then
				funcval=funcval+abs(trace(k)-trace0(k))
			else if (norm==2) then
				funcval=funcval+(trace(k)-trace0(k))**2
			end if
		end do
		if (norm==2) funcval=sqrt(funcval)
			
		! free allocated space
		deallocate(trace)

		count=count+1

		return

	end subroutine func
	!-------------------------------------------------------------------------


	subroutine synthetic_seismogram(j,option,npopsiz,nparam,nlyr,&
		& freq,nt,dt,dz,parent,trace,flag)
	!-------------------------------------------------------------------------
	! Compute a synthetic seismogram for a plane wave at normal incidence
	!-------------------------------------------------------------------------

        !! subroutine arguments
		integer, intent(in) :: j
		integer, intent(in) :: option
        integer, intent(in) :: npopsiz
        integer, intent(in) :: nparam
		integer, intent(in)	:: nlyr				! number of layers
		integer, intent(in) :: nt
		real, intent(in)    :: freq
		real, intent(in)    :: dt
		real, intent(in)    :: dz
		integer,intent(in),optional :: flag
		real, dimension(nt),intent(out)             :: trace 
        real, dimension(nparam,npopsiz), intent(in) :: parent

		!! local variables
		integer :: k,kk,kkk,stat
		real,dimension(nlyr)      :: vel		! velocity
		real,dimension(nlyr)      :: rho		! density
		real,dimension(nlyr)      :: deltaz		! thickness

		! extract velocity, density and thickness information from parent array
		if (option==1) then			! velocities only
			do k=1,nparam
				vel(k)=parent(k,j)
			end do
			rho=1.
			deltaz=dz*FT2M
		else if (option==2) then	! velocities and densities
			do k=1,nparam
				kk=1+k/2
				if (MOD(k,2)) then
					rho(kk)=parent(k,j)
				else 
					vel(kk)=parent(k,j)
				end if
			end do
			deltaz=dz*FT2M
		else if (option==3) then	! velocities, densities and thicknesses	
			do k=1,nparam
				kk=1+k/2
				kkk=1+k/3
				if (MOD(k,3)) then
					deltaz(kkk)=parent(k,j)
				else if (MOD(k,2)) then
					rho(kkk)=parent(k,j)
				else
					vel(kkk)=parent(k,j)
				end if
			end do
		end if
		
		! create a trace
        call create_trace(nt,nlyr,dt,freq,vel,rho,deltaz,trace)

		! if convergence has been achieved, output the trace
		if (present(flag)) then

			!! write info to history file of solution trace
			call to_history("n1",nt,"trace")
			call to_history("n2",1,"trace")
			call to_history("n3",1,"trace")
			call to_history("o1",0.,"trace")
			call to_history("d1",dt,"trace")
			call to_history("esize",4,"trace")

			!! write the trace out
			call srite("trace",trace,4*size(trace))
		end if

		return
	
	end subroutine synthetic_seismogram
	!-----------------------------------------------------------------------
		

	subroutine create_trace(nt,nlyr,dt,freq,vel,rho,deltaz,trace)	
	!-----------------------------------------------------------------------
	! subroutine to create a seismic trace
	!-----------------------------------------------------------------------
		
		!! subroutine arguments	
		integer, intent(in) :: nt 
		integer, intent(in) :: nlyr 
		real, intent(in)    :: freq		! peak frequency for Ricker wavelet
		real, intent(in)    :: dt
		real,dimension(nlyr),intent(in) :: vel			! velocity
		real,dimension(nlyr),intent(in) :: rho			! density
		real,dimension(nlyr),intent(in) :: deltaz		! thickness
		real,dimension(:),intent(out)	:: trace		! output trace

		!! local variables
		integer :: k,iw,lw
		real	:: tmp1,tmp2
		real,dimension(:),pointer	:: wavelet	! Ricker wavelet
		real,dimension(nlyr)		:: rc		! reflection coefficients
		real,dimension(nlyr)		:: t		! 2-way vertical travel time
		real,dimension(nt)		:: tout		! array of equisampled output times
		real,dimension(nt)		:: aux_trace! array for interpolated trace 

		! compute normal incidence reflection coefficients and
		! 2-way normal-incidence times 
		call depth_to_time(nlyr,vel,rho,deltaz,rc,t)

		! create array of equisampled times for output trace
		do k=1,nt
			tout(k)=(k-1)*dt
		end do

		! interpolate times to uniform sampling interval
		call ints8r(nlyr,dt,t,rc,0.0,0.0,nt,tout,aux_trace)

		! compute ricker wavelet
		call rickersu (freq,dt,lw,iw,wavelet)

		! convolve reflection amplitudes with ricker wavelet
		! to form the synthetic trace
		call convsu(lw,iw+1,wavelet,nt,1,aux_trace,nt,1,trace)

		return

	end subroutine create_trace

	subroutine depth_to_time(nlyr,vel,rho,deltaz,rc,t)
	! compute normal incidence reflection coefficients and
	! 2-way normal-incidence times 

		! subroutine arguments
		integer, intent(in)				:: nlyr		! number of layers
		real,dimension(nlyr),intent(in) :: vel		! velocity (m/s)
		real,dimension(nlyr),intent(in) :: rho		! density (g/cc)
		real,dimension(nlyr),intent(in) :: deltaz	! thickness (m)
		real,dimension(nlyr),intent(out):: rc		! reflection coefficients
		real,dimension(nlyr),intent(out):: t		! 2-way vertical travel time

		! local variables
		real 	:: tmp1,tmp2
		integer :: k

		!! initialize arrays
		tmp1=rho(1)*vel(1)		
		rc(1)=0.0
		t(1)=2.0*deltaz(1)/vel(1)

		!! iterate over layers
		do k=2,nlyr
			tmp2=rho(k)*vel(k)
			rc(k)=(tmp2-tmp1)/(tmp2+tmp1)
			tmp1=tmp2
			t(k)=t(k-1)+2.0*deltaz(k)/vel(k)
		end do

		return

	end subroutine depth_to_time

end module ga_Mauro_function
