!******************************************************************************
!
!    pspectra.f95 (part of doptomog package v2.0)
!    Copyright (C) 2015  Enrico J. Kotze ejk@saao.ac.za
!
!    This program is free software: you can redistribute it and/or modify
!    it under the terms of the GNU General Public License as published by
!    the Free Software Foundation, either version 3 of the License, or
!    (at your option) any later version.
!
!    This program is distributed in the hope that it will be useful,
!    but WITHOUT ANY WARRANTY; without even the implied warranty of
!    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!    GNU General Public License for more details.
!
!    You should have received a copy of the GNU General Public License
!    along with this program.  If not, see <http://www.gnu.org/licenses/>.
!
!******************************************************************************
program pspectra
!******************************************************************************

!******************************************************************************
	use mconstants
	use mfunctions
	use mspectra
	implicit none
	
	integer											::	i
	character(len=255)								::	arg
	
	! get current working directory
	call getcwd(cwd)

	! get input parameters from command line 
	do i = 1, iargc()
		call getarg(i, arg)
		if (i == 1) read(arg, '(i2)') stype
		if (i == 2) read(arg, '(a)') srcpath
		! ... that's it for now...
	end do

	select case (stype)
		case (prepare)
			! get input from 'specprepare.in'
			call inputPrepare()
			! get input spectra
			call inputSpectraPrepare()
			! process spectra
			call processPrepare()
			! output spectra
			call outputPrepare()
		case (extract)
			! get input from 'specextract.in'
			call inputExtract()
			! get input spectra
			call inputSpectraExtract()
			! process spectra
			call processExtract()
			! output spectra
			call outputExtract()
	end select
	
!******************************************************************************
contains	
!******************************************************************************

!******************************************************************************

!******************************************************************************
	subroutine inputPrepare()

		! get input from 'specprepare.in'
		inpath = trim(srcpath)//'in/'
		open(unit = 1, file = trim(inpath)//'specprepare.in')
		! name of file containing list of spectra
		read(1, '(a20)') specList; specList = trim(specList)
		! name of folder containing list of spectra and spectra files
		read(1, '(a20)') specDir; specDir = trim(specDir)
		! period fold indicator (0 = orbital, 1 = spin, 2 = extract spin)
		read(1, *) pmod
		! column nr of spectrum file name in spectra list
		read(1, *) fcol
		! column nr of phase value
		read(1, *) pcol
		! column nr of phase interval (0 if not present)
		read(1, *) dcol
		! column nr of wavelength in spectrum file
		read(1, *) wcol
		! column nr of flux in spectrum file
		read(1, *) scol
		! fixed phase interval for the exposure times (0 if not relevant)
		read(1, *) dpha
	5	continue
		close(1)

		return

	end subroutine inputPrepare
!******************************************************************************

!******************************************************************************
	subroutine inputExtract()

		! get input from 'specextract.in'
		inpath = trim(srcpath)//'in/'
		open(unit = 1, file = trim(inpath)//'specextract.in')
		! gamma velocity [m/s]
		read(1, *) gam
		! rest wavelength of line [Angstroms]
		read(1, *) w0
		! wavelength range around w0 to keep (0 to keep everything)
		read(1, *) cw
		! wavelength range around w0 (w0 - dw < w < w0 + dw)
		read(1, *) dw
		! dw factor for inner range
		read(1, *) df1
		! dw factor for outer range
		read(1, *) df2
		! if abl = 1, invert line (for absorption lines)
		read(1, *) abl
		! if atm = 1, take out atmospheric absorption
		read(1, *) atm
		! nr of half orbit spectra to extract
		read(1, *, end = 5) nho
	5	continue
		close(1)

		return

	end subroutine inputExtract
!******************************************************************************

!******************************************************************************

!******************************************************************************
	subroutine inputSpectraPrepare()

		character(len=80)							::	line = ''
		integer										::	i, j

		! update 'inpath' with 'specDir'
		inpath = trim(inpath)//trim(specDir)//'/'

		! get number of lines in specList, i.e., nr of phases entries (npe)
		open(unit = 1, file = trim(inpath)//trim(specList))
		do
			read(1, *, iostat=ios) line
			if (ios /= 0) exit
			npe = npe + 1
		end do
		close(1)
		! allocate dimensions to applicable arrays
		allocate(sfile(npe), phase(npe), dphase(npe))
		phase = 0d0; dphase = 0d0; dpham = 0d0
		! set maximum number of columns to be read from specList
		mcol = max(fcol, pcol, dcol)
		! allocate dimensions to entries array
		allocate(entries(mcol))
		! process specList
		open(unit = 2, file = trim(inpath)//trim(specList))
		do i = 1, npe
			! read a line
			read(2, '(a)') line
			! read entries from the line
			read(line, *) (entries(j), j = 1, mcol)
			! get file name
			sfile(i) = trim(entries(fcol))
			! read phases into phase array
			read(entries(pcol), *) phase(i)
			! read phase intervals into dphase array (if present)
			if (dcol /= 0) read(entries(dcol), *) dphase(i)
			if (dcol /= 0) dpham = dpham + dphase(i)
		end do
		close(2)
		! determine phase intervals (if not present)
		if(dcol == 0) dphase = 1d0 / dble(npe)
		! determine average phase interval
		if(dcol /= 0) dpham = dpham / dble(npe)
		! if a fixed phase interval was specified
		if(dpha /= 0d0) dphase = dpha
		! determine average phase interval
		if(dpha /= 0d0) dpham = dpha
		! clean-up
		! de-allocate work arrays
		deallocate(entries)

		return

	end subroutine inputSpectraPrepare
!******************************************************************************

!******************************************************************************
	subroutine inputSpectraExtract()

		integer										::	i, j
		character(len=80)							::	errMessage

		inpath = trim(srcpath)//'out/spectra/'
		open(unit = 1, file = trim(inpath)//'specprepare.out')
		! read period fold indicator
		read(1, *) pmod
		! read array sizes
		read(1, *) npe, nws
		! allocate dimension to spec
		allocate(spec(npe, nws))
		! allocate dimension to phase, dphase, wl, wt
		allocate(phase(npe), dphase(npe), wl(nws)) !, wt(npe))
		! read phases and phase bins
		read(1, *) (phase(i), i = 1, npe)
		read(1, *) (dphase(i), i = 1, npe)
		! read wavelengths
		read(1, *) (wl(i), i = 1, nws)
		if ((w0 + dw) > wl(nws) .or. (w0 - dw) < wl(1)) then
			write(errMessage,'(a,f8.2,a,f6.2,a,2f8.2)') &
				'w0 =', w0, ' with dw =', dw, ' is out of range:', wl(1), wl(nws)
			call nrerror('inputSpectraExtract: '//trim(errMessage))
		end if
		if ((w0 + df2 * dw) > wl(nws) .or. (w0 - df2 * dw) < wl(1)) then
			write(errMessage,'(a,f8.2,a,f7.2,a,2f8.2)') &
				'w0 =', w0, ' with df2*dw =', df2 * dw, ' is out of range:', wl(1), wl(nws)
			call nrerror('inputSpectraExtract: '//trim(errMessage))
		end if
		! read spectra
		read(1, *) ((spec(i,j), i = 1, npe), j = 1, nws)
		! read file name containing list of spectra
		read(1, '(a20)') specList
		close(1)
		npr = npe
		! phase index master array (maximum index = # of phase values)
		allocate(pIndexM(0:npe))
		do i = 0, npe
			pIndexM(i) = dble(i)
		end do

		return

	end subroutine inputSpectraExtract
!******************************************************************************

!******************************************************************************
	subroutine processPrepare()

		character(len=80)							::	line = ''
		integer										::	i, j, k

		! set maximum number of columns to be read from each spectrum file
		mcol = max(wcol, scol)
		! allocate dimensions to entries arrays
		allocate(entries(mcol))
		do i = 1, npe
			nws = 0
			! get number of lines in spectrum file, i.e., nr of wavelength and
			! spectrum entries (nws)
			open(unit = 1, file = trim(inpath)//sfile(i))
			do
				read(1, *, iostat=ios) line
				if (ios /= 0) exit
				nws = nws + 1
			end do
			close(1)
			! allocate dimensions to applicable arrays
			if (i > 1) then
				deallocate(wlj, specj, spec1)
			end if
			allocate(wlj(nws), specj(0:nws), spec1(0:nws))
			specj = 0d0; spec1 = 0d0
			! read spectrum file
			open(unit = 2, file = trim(inpath)//sfile(i))
			do j = 1, nws
				! read a line
				read(2, '(a)') line
				! read entries from the line
				read(line, *) (entries(k), k = 1, mcol)
				! read wavelength into wlj array
				read(entries(wcol), *) wlj(j)
				! read spectrum into specj array
				read(entries(scol), *) specj(j)
			end do
			close(2)
			if (i == 1) then
				! allocate dimensions to applicable arrays
				allocate(wl1(nws), spec(npe,nws))
				! store first set of wavelengths
				wl1 = wlj
				spec1 = specj
			else
				! interpolate spectrum to scale of first file
				call interpolate(wlj, specj, nws, wl1, spec1, nws)
			end if
			! read spectrum into spec array
			do j = 1, nws
				spec(i,j) = spec1(j)
			end do
		end do
		allocate(wl(nws))
		wl = wl1
		! clean-up
		! de-allocate work arrays
		deallocate(entries, sfile, wlj, wl1, specj, spec1)

		return

	end subroutine processPrepare
!******************************************************************************

!******************************************************************************
	subroutine processExtract()

		! clean spectra
		call cleanSpectra()
		! sort phases
		call sortPhases()
		! convert to rest wavelength (gamma)
		call restWavelengthGamma()
		! determine windowed spectrum
		call windowedSpectrum()
		! remove atmospheric absorption
		call atmosphericAbsorption()
		! determine Doppler velocities
		call dopplerVelocities()
		! extract spectra
		call extractSpectra()
		! clean-up
		! de-allocate work arrays
		deallocate(bl, indx, ind, iw, trb, trbc, clev, trw, trwc)

		return

	end subroutine processExtract
!******************************************************************************

!******************************************************************************
	subroutine cleanSpectra()

		! clean spectra
		integer										::	i, j, nc
		real										::	sumc

		if(cw > 0.) then
			do i = 1, npe
				nc = 0
				sumc = 0.
				do j = 1, nws
					if(wl(j) < (w0 - cw) .or. wl(j) > (w0 + cw)) then
						nc = nc + 1
						sumc = sumc + spec(i,j)
					end if
				end do
				do j = 1, nws
					if(wl(j) < (w0 - cw) .or. wl(j) > (w0 + cw)) then
						spec(i,j) = sumc / nc
					end if
				end do
			end do
		end if

		return

	end subroutine cleanSpectra
!******************************************************************************

!******************************************************************************
	subroutine sortPhases()

		integer										::	i
		real, dimension(:), allocatable				::	phasef

		! sort phases
		allocate(phasef(npe))
		phasef = mod(phase, 1.)
		do i = 1, npe
			if (phasef(i) < 0.) phasef(i) = phasef(i) + 1.
		end do
		phase = phasef

		return

	end subroutine sortPhases
!******************************************************************************

!******************************************************************************
	subroutine restWavelengthGamma()

		! convert to rest wavelength (gamma)
		real(kind(1d0))								::	gamadj

		! subtract gamma velocity
		if (gam /= 0d0) then
			gamadj = (1d0 + gam / cSI) * sqrt(1d0 - (gam / cSI)**2)
			wl = wl / gamadj
			! apply relativity correction to spectra
			spec = spec * gamadj
		end if

		return

	end subroutine restWavelengthGamma
!******************************************************************************

!******************************************************************************
	subroutine windowedSpectrum()

		integer										::	i
		real										::	f
		real(kind(1d0)), dimension(:), allocatable	::	indt

		! allocate dimensions to index arrays
		allocate(indx(0:nws), iw(0:2)); indx = 0d0; iw = 0d0
		! master index array (maximum index = # of wavelength values)
		do i = 1, nws
			indx(i) = dble(i)
		end do
		! specified wavelengths
		wr(1) = w0 - df2 * dw
		wr(2) = w0 + df2 * dw
		! interpolate wavelength indices for specified wavelengths
		call interpolate(wl, indx, nws, wr, iw, 2)
		i0 = max(1, int(iw(1)))
		i1 = min(nws, int(iw(2)))
		! get the wavelengths between the indices
		wl = wl(i0:i1+1)
		nwl = (i1 + 1 - i0) + 1
		! allocate dimensions to windowed spectrum and other work arrays
		allocate(trw(npe,nwl), trwc(npe,nwl), trb(nwl), trbc(nwl), ind(nwl))
		allocate(clev(npe))
		! windowed spectrum
		trw = spec(1:npe,i0:i1+1)
		trwc = trw
		! phase averaged windowed spectrum
		trb = sum(trw, dim=1) / dble(npe)
		! corresponding wavelength index array (without zero index)
		ind = indx(1:nwl)
		! sort phase averaged windowed spectrum ascending
		call piksrt(nwl, trb, ind)
		! pick 25% of the indices around the median
		f = 0.25
		allocate(indt(int((1.-f)*0.5*(i1-i0))+1:int((1.+f)*0.5*(i1-i0))+1))
		indt = ind(int((1.-f)*0.5*(i1-i0))+1:int((1.+f)*0.5*(i1-i0))+1)
		clev = sum(trw(1:npe,int(indt)), dim=2) / dble(size(indt))
		! phase averaged
		cla = sum(clev) / dble(npe)
		deallocate(indt)

		return

	end subroutine windowedSpectrum
!******************************************************************************

!******************************************************************************
	subroutine atmosphericAbsorption()

		integer										::	i

		! remove atmospheric absorption
		if (atm == 1) then
			do i = 1, i1 - i0
				trw(1:npe,i) = trw(1:npe,i) / clev
			end do
			allocate(tra(nwl))
			tra = sum(trw, dim=1) / dble(npe)
			do i = 1, npe
				trw(i,1:nwl) = trw(i,1:nwl) - tra + 1d0
			end do
			do i = 1, i1 - i0
				trw(1:npe,i) = trw(1:npe,i) * clev
			end do
		end if
		do i = 1, npe
			trwc(i,1:nwl) = trw(i,1:nwl) - clev(i) + cla
		end do

		return

	end subroutine atmosphericAbsorption
!******************************************************************************

!******************************************************************************
	subroutine dopplerVelocities()

		integer										::	i, ic0, ic1, ic2, ic3

		! phase averaged windowed spectrum
		trbc = sum(trwc, dim=1) / dble(npe)
		! corresponding index array (with zero index)
		deallocate(ind); allocate(ind(0:nwl))
		ind = indx(0:nwl)
		! set wavelength ranges for continuum
		wc1(1) = w0 - df2 * dw
		wc1(2) = w0 - df1 * dw
		wc2(1) = w0 + df1 * dw
		wc2(2) = w0 + df2 * dw
		! interpolate wavelength indices for 1st range of continuum wavelengths
		call interpolate(wl, ind, nwl, wc1, iw, 2)
		ic0 = max(1, int(iw(1)))
		ic1 = min(nws, int(iw(2)))
		! interpolate wavelength indices for 2nd range of continuum wavelengths
		call interpolate(wl, ind, nwl, wc2, iw, 2)
		ic2 = max(1, int(iw(1)))
		ic3 = min(nws, int(iw(2)))
		! wavelength - average continuum
		ca1 = 0d0; ca2 = 0d0
		if (ic1 - ic0 /= 0d0) ca1 = sum(trbc(ic0:ic1-1)) / dble(ic1 - ic0)
		if (ic3 - ic2 /= 0d0) ca2 = sum(trbc(ic2:ic3-1)) / dble(ic3 - ic2)
		! average wavelength
		w1 = (wc1(1) + wc1(2)) / 2d0
		w2 = (wc2(1) + wc2(2)) / 2d0
		! Doppler velocity
		v1 = ((w1 / w0)**2 - 1d0) * cSI / ((w1 / w0)**2 + 1d0)
		v2 = ((w2 / w0)**2 - 1d0) * cSI / ((w2 / w0)**2 + 1d0)

		return

	end subroutine dopplerVelocities
!******************************************************************************

!******************************************************************************
	subroutine extractSpectra()

		integer										::	i

		! re-determine i0, i1:
		! specified wavelengths
		wr(1) = w0 - dw
		wr(2) = w0 + dw
		! interpolate wavelength indices for specified wavelengths
		call interpolate(wl, ind, nwl, wr, iw, 2)
		i0 = max(1, int(iw(1)))
		i1 = min(nws, int(iw(2)))
		! allocate dimension to spout
		nve = i1 - i0
		allocate(spout(npe,nve)); spout = 0d0
		allocate(wp(nve), vp(nve), bl(nve)); wp = 0d0; vp = 0d0; bl = 0d0
		wp = wl(i0+1:i1)
		! phase index array
		deallocate(ind); allocate(ind(1:npe))
		ind = indx(1:npe)
		! build spout
		do i = 1, npe
			spout(i,1:nve) = trwc(int(ind(i)),i0+1:i1)
		end do
		! calculate Doppler velocities for wavelengths in requested range
		vp = cSI * ((wl(i0+1:i1) / w0)**2 - 1d0) / ((wl(i0+1:i1) / w0)**2 + 1d0)
		! set velocity scale
		vs = 10.**(dble(int(dlog10(vp(nve))))) / vf
		! construct linearly varying continuum
		bl = (ca2 - ca1) / (v2 - v1) * (vp - v1) + ca1
		! subtract from spectrum
		do i = 1, npe
			spout(i,1:nve) = spout(i,1:nve) - bl
		end do
		if (abl /= 0) spout = -spout

		return

	end subroutine extractSpectra
!******************************************************************************

!******************************************************************************
	subroutine outputPrepare()

		integer										::	i, j

		! write output to 'specprepare.out'
		outpath = trim(srcpath)//'out/spectra/'
		open(unit = 1, file = trim(outpath)//'specprepare.out')
		write(1, '(i3)') pmod
		write(1, '(2i5)') npe, nws
		write(1, '(1p6es14.6)') phase(1:npe)
		write(1, '(1p6es14.6)') dphase(1:npe)
		write(1, '(1p6es14.6)') wl(1:nws)
		write(1, '(1p6es14.6)') ((spec(i,j), i = 1, npe), j = 1, nws)
		write(1, *) trim(specList)
		! clean-up
		! de-allocate work arrays
		deallocate(phase, dphase, wl, spec)

		return

	end subroutine outputPrepare
!******************************************************************************

!******************************************************************************
	subroutine outputExtract()

		integer										::	i, j, k, tpe
		character(len=4)							::	suffix

		outpath = trim(srcpath)//'out/spectra/'
		! allocate dimensions to temporary arrays
		allocate(tphase(npe), tdphase(npe), tspout(npe,nve))
		if (nho > 0) then
			! write phased output (if any) to 'specphased.out'
			open(unit = 1, file = trim(outpath)//'specphased.out')
			write(1, '(i3)') nho
		end if
		! loop for nr of half orbits
		do k = 0, nho
			! determine phase extracted stuff
			call phaseExtract(k, tpe, suffix)
			! write output to 'specextract.out'
			open(unit = 2, file = trim(outpath)//'specextract'//trim(suffix)//'.out')
			write(2, '(i3, es10.2)') pmod, vf
			write(2, '(2es10.2, f10.2, 2i3, a)') vs, gam, w0, abl, atm, '  '//trim(specList)
			write(2, '(2i5)') tpe, nve
			write(2, '(1p6es14.6)') (tphase(i), i = 1, tpe)
			write(2, '(1p6es14.6)') (tdphase(i), i = 1, tpe)
			write(2, '(1p6es14.6)') (vp(j), j = 1, nve)
			do j = 1, nve
				write(2, '(1p6es14.6)') tspout(1:tpe,j)
			end do
			close(2)
		end do
		if (nho > 0) then
			close(1)
		end if
		! clean-up
		! de-allocate work arrays
		deallocate(phase, dphase, vp, spout, tphase, tdphase, tspout, wl, spec)

	68 	format (a, a)
	69 	format (a, 'integer, parameter	::	', a, i3)

		return

	end subroutine outputExtract
!******************************************************************************

!******************************************************************************
	subroutine phaseExtract(k, tpe, suffix)

		integer, intent(in)							::	k
		integer, intent(out)						::	tpe
		character(len=4), intent(out)				::	suffix
		real(kind(1d0)), dimension(0:4)				::	pIndex
		real(kind(1d0)), dimension(4)				::	pRange
		real(kind(1d0))								::	pStep
		integer, dimension(4)						::	ip
		integer										::	i, j, n

		if (k == 0) then
			suffix = ''
			ip(1) = 0
			ip(2) = 0
			ip(3) = 0
			ip(4) = npe
		else if (k > 0) then
			pStep = 1d0 / dble(nho)
			! write output file suffix
			write(suffix, '(i3)') k - 1
			suffix = '.'//trim(adjustl(suffix))
			! determine half orbit range
			if (pStep * (k - 1) < 0.5d0) then
				pRange(1) = 0d0
				pRange(2) = 0d0
				pRange(3) = pStep * (k - 1)
				pRange(4) = pStep * (k - 1) + 0.5d0
				write(1, '(2f10.4)') pRange(3), pRange(4)
			else
				pRange(1) = 0d0
				pRange(2) = pStep * (k - 1) - 0.5d0
				pRange(3) = pStep * (k - 1)
				pRange(4) = 1d0
				write(1, '(2f10.4)') pRange(3), pRange(2)
			end if
			! interpolate phase indices for specified half orbit range
			call interpolate(phase, pIndexM, npe, pRange, pIndex, 4)
			ip(1) = int(pIndex(1))
			ip(2) = int(pIndex(2))
			ip(3) = int(pIndex(3))
			ip(4) = int(pIndex(4))
			if (ip(3) > 0 .and. ip(4) == 0) ip(4) = ip(4) + npe
		end if
		tpe = ip(4) - ip(3) + ip(2) - ip(1)
		n = 0
		if (ip(2) > 0) then
			do i = ip(1) + 1, ip(2)
				n = n + 1
				tphase(n) = phase(i)
				tdphase(n) = dphase(i)
				tspout(n,1:nve) = spout(i,1:nve)
			end do
		end if
		do i = ip(3) + 1, ip(4)
			n = n + 1
			tphase(n) = phase(i)
			tdphase(n) = dphase(i)
			tspout(n,1:nve) = spout(i,1:nve)
		end do

		return

	end subroutine phaseExtract
!******************************************************************************

!******************************************************************************

!******************************************************************************
end program pspectra
!******************************************************************************
