import numpy as np
from scipy import interpolate
import matplotlib.pyplot as plt
from matplotlib.ticker import AutoMinorLocator
import pandas as pd
from astropy.io import fits
from astropy import units as u
import sys
import os
import warnings
import copy
import smart
warnings.filterwarnings("ignore")

class Spectrum():
	"""
	The spectrum class for reading the reduced Keck/NIRSPEC data by NSDRP and SDSS/APOGEE data.

	Parameters
	----------
	name : str
	       The data filename before the order number.
	       Ex. name='jan19s0022'
	order: int
	       The order of the spectra.
	path : str
	       The directory where the reduced data is.

	Returns
	-------
	flux : numpy.ndarray
	       The flux of the spectrum.
	wave : numpy.ndarray
	       The wavelength of the spectrum
	noise: numpy.ndarray
	       The noise of the spectrum
	sky  : numpy.ndarray
	       The sky emission
	plot : matplotlib plot
	       plot the spectrum (with noise on/off option)

	Examples
	--------
	>>> import smart
	>>> path = '/path/to/reducedData'
	>>> data = smart.Spectrum(name='jan19s0022', order=33, path=path)
	>>> data.plot()

	"""
	def __init__(self, **kwargs):
		self.instrument = kwargs.get('instrument','nirspec')
		if self.instrument == 'nirspec':
			self.name      = kwargs.get('name')
			self.order     = kwargs.get('order')
			self.path      = kwargs.get('path')
			self.apply_sigma_mask = kwargs.get('apply_sigma_mask',False)
			#self.manaulmask = kwargs('manaulmask', False)

			if self.path == None:
				self.path = './'

			fullpath = self.path + '/' + self.name + '_' + str(self.order) + '_all.fits'

			hdulist = fits.open(fullpath, ignore_missing_end=True)

			#The indices 0 to 3 correspond to wavelength, flux, noise, and sky
			self.header = hdulist[0].header
			self.wave   = hdulist[0].data
			self.flux   = hdulist[1].data
			self.noise  = hdulist[2].data
			try:
				self.sky = hdulist[3].data
			except IndexError:
				print("No sky line data.")
				self.sky = np.zeros(self.wave.shape)

			self.mask  = []

			# define a list for storing the best wavelength shift
			self.bestshift = []

			# store the original parameters
			self.oriWave  = hdulist[0].data
			self.oriFlux  = hdulist[1].data
			self.oriNoise = hdulist[2].data

		elif self.instrument == 'apogee':
			self.name      = kwargs.get('name')
			self.path      = kwargs.get('path')
			self.datatype  = kwargs.get('datatype','aspcap')
			self.apply_sigma_mask = kwargs.get('apply_sigma_mask',False)
			self.applytell = kwargs.get('applytell', False)
			self.chip      = kwargs.get('chip', 'all')

			hdulist        = fits.open(self.path)

			if self.datatype == 'aspcap':
				crval1         = hdulist[1].header['CRVAL1']
				cdelt1         = hdulist[1].header['CDELT1']
				naxis1         = hdulist[1].header['NAXIS1']
				self.header4   = hdulist[4].header
				self.param     = hdulist[4].data['PARAM']
				self.wave      = np.array(pow(10, crval1 + cdelt1 * np.arange(naxis1)))
				self.oriWave   = np.array(pow(10, crval1 + cdelt1 * np.arange(naxis1)))
				self.flux      = np.array(hdulist[1].data)
				self.noise     = np.array(hdulist[2].data)

			elif self.datatype == 'ap1d':
				self.header4   = hdulist[4].header
				self.header5   = hdulist[5].header
				# use aspcap data as wavelength calibrators
				self.wave      = np.array(hdulist[4].data)
				self.flux      = np.array(hdulist[1].data)
				self.noise     = np.array(hdulist[2].data)
				# store the original parameters
				self.oriWave   = np.array(hdulist[4].data)
				self.oriFlux  = np.array(hdulist[1].data)
				self.oriNoise = np.array(hdulist[2].data)

			elif self.datatype == 'apvisit':
				self.header1   = hdulist[1].header
				self.header2   = hdulist[2].header
				self.header3   = hdulist[3].header
				self.header4   = hdulist[4].header
				self.header5   = hdulist[5].header
				self.header6   = hdulist[6].header
				self.header7   = hdulist[7].header
				self.header8   = hdulist[8].header
				self.header9   = hdulist[9].header
				self.header10  = hdulist[10].header

				# read the bitmask
				self.bitmask   = hdulist[3].data

				#import bitmask
				# chip a
				if self.chip == 'all' or self.chip == 'a':
					mask_0 = []
					for i in range(len(hdulist[3].data[0])):
						bitmask = smart.bits_set(hdulist[3].data[0][i])
						if (0 in bitmask) or (1 in bitmask) or (2 in bitmask) or \
						(3 in bitmask) or (4 in bitmask) or (5 in bitmask) or \
						(6 in bitmask) or (12 in bitmask) or (14 in bitmask):
							mask_0.append(i)

				# chip b
				if self.chip == 'all' or self.chip == 'b':
					mask_1 = []
					for i in range(len(hdulist[3].data[1])):
						bitmask = smart.bits_set(hdulist[3].data[1][i])
						if (0 in bitmask) or (1 in bitmask) or (2 in bitmask) or \
						(3 in bitmask) or (4 in bitmask) or (5 in bitmask) or \
						(6 in bitmask) or (12 in bitmask) or (14 in bitmask):
							mask_1.append(i)

				# chip c
				if self.chip == 'all' or self.chip == 'c':
					mask_2 = []
					for i in range(len(hdulist[3].data[2])):
						bitmask = smart.bits_set(hdulist[3].data[2][i])
						if (0 in bitmask) or (1 in bitmask) or (2 in bitmask) or \
						(3 in bitmask) or (4 in bitmask) or (5 in bitmask) or \
						(6 in bitmask) or (12 in bitmask) or (14 in bitmask):
							mask_2.append(i)

				if self.chip == 'all':
					self.wave      = np.array(list(np.delete(hdulist[4].data[0], mask_0))+list(np.delete(hdulist[4].data[1], mask_1))+list(np.delete(hdulist[4].data[2], mask_2)))
					self.flux      = np.array(list(np.delete(hdulist[1].data[0], mask_0))+list(np.delete(hdulist[1].data[1], mask_1))+list(np.delete(hdulist[1].data[2], mask_2)))
					self.noise     = np.array(list(np.delete(hdulist[2].data[0], mask_0))+list(np.delete(hdulist[2].data[1], mask_1))+list(np.delete(hdulist[2].data[2], mask_2)))
					self.sky       = np.array(list(hdulist[5].data[0])+list(hdulist[5].data[1])+list(hdulist[5].data[2]))
					self.skynoise  = np.array(list(hdulist[6].data[0])+list(hdulist[6].data[1])+list(hdulist[6].data[2]))
					self.tell      = np.array(list(np.delete(hdulist[7].data[0], mask_0))+list(np.delete(hdulist[7].data[1], mask_1))+list(np.delete(hdulist[7].data[2], mask_2)))
					self.tellnoise = np.array(list(np.delete(hdulist[8].data[0], mask_0))+list(np.delete(hdulist[8].data[1], mask_1))+list(np.delete(hdulist[8].data[2], mask_2)))

					# store the original parameters
					self.oriWave   = np.array(list(hdulist[4].data[0])+list(hdulist[4].data[1])+list(hdulist[4].data[2]))
					self.oriFlux   = np.array(list(hdulist[1].data[0])+list(hdulist[1].data[1])+list(hdulist[1].data[2]))
					self.oriNoise  = np.array(list(hdulist[2].data[0])+list(hdulist[2].data[1])+list(hdulist[2].data[2]))

				elif self.chip == 'a':
					self.wave      = np.array(list(np.delete(hdulist[4].data[0], mask_0)))
					self.flux      = np.array(list(np.delete(hdulist[1].data[0], mask_0)))
					self.noise     = np.array(list(np.delete(hdulist[2].data[0], mask_0)))
					self.sky       = np.array(list(hdulist[5].data[0]))
					self.skynoise  = np.array(list(hdulist[6].data[0]))
					self.tell      = np.array(list(np.delete(hdulist[7].data[0], mask_0)))
					self.tellnoise = np.array(list(np.delete(hdulist[8].data[0], mask_0)))

					# store the original parameters
					self.oriWave   = np.array(list(hdulist[4].data[0]))
					self.oriFlux   = np.array(list(hdulist[1].data[0]))
					self.oriNoise  = np.array(list(hdulist[2].data[0]))

				elif self.chip == 'b':
					self.wave      = np.array(list(np.delete(hdulist[4].data[1], mask_1)))
					self.flux      = np.array(list(np.delete(hdulist[1].data[1], mask_1)))
					self.noise     = np.array(list(np.delete(hdulist[2].data[1], mask_1)))
					self.sky       = np.array(list(hdulist[5].data[1]))
					self.skynoise  = np.array(list(hdulist[6].data[1]))
					self.tell      = np.array(list(np.delete(hdulist[7].data[1], mask_1)))
					self.tellnoise = np.array(list(np.delete(hdulist[8].data[1], mask_1)))

					# store the original parameters
					self.oriWave   = np.array(list(hdulist[4].data[1]))
					self.oriFlux   = np.array(list(hdulist[1].data[1]))
					self.oriNoise  = np.array(list(hdulist[2].data[1]))

				elif self.chip == 'c':
					self.wave      = np.array(list(np.delete(hdulist[4].data[2], mask_2)))
					self.flux      = np.array(list(np.delete(hdulist[1].data[2], mask_2)))
					self.noise     = np.array(list(np.delete(hdulist[2].data[2], mask_2)))
					self.sky       = np.array(list(hdulist[5].data[2]))
					self.skynoise  = np.array(list(hdulist[6].data[2]))
					self.tell      = np.array(list(np.delete(hdulist[7].data[2], mask_2)))
					self.tellnoise = np.array(list(np.delete(hdulist[8].data[2], mask_2)))

					# store the original parameters
					self.oriWave   = np.array(list(hdulist[4].data[2]))
					self.oriFlux   = np.array(list(hdulist[1].data[2]))
					self.oriNoise  = np.array(list(hdulist[2].data[2]))

				if self.applytell:
					self.flux *= self.tell
				self.wavecoeff = hdulist[9].data
				self.lsfcoeff  = hdulist[10].data

				if self.wave[0] > self.wave[-1]:
					self.wave      = self.wave[::-1]
					self.flux      = self.flux[::-1]
					self.noise     = self.noise[::-1]
					self.sky       = self.sky[::-1]
					self.skynoise  = self.skynoise[::-1]
					self.tell      = self.tell[::-1]
					self.tellnoise = self.tellnoise[::-1]
					self.oriWave   = self.oriWave[::-1]
					self.oriFlux   = self.oriFlux[::-1]
					self.oriNoise  = self.oriNoise[::-1]

				# to separate the continuum end points
				self.oriWave0  = hdulist[4].data
				self.oriFlux0  = hdulist[1].data

				## APOGEE APVISIT has corrected the telluric absorption; the forward-modeling routine needs to put it back
				#self.flux     *= self.tell

			elif self.datatype == 'apstar':
				crval1         = hdulist[0].header['CRVAL1']
				cdelt1         = hdulist[0].header['CDELT1']
				naxis1         = hdulist[0].header['NWAVE']
				self.header4   = hdulist[4].header
				self.header5   = hdulist[5].header
				self.header6   = hdulist[6].header
				self.header7   = hdulist[7].header
				self.header8   = hdulist[8].header
				self.header9   = hdulist[9].header

				#print(hdulist)
				#print(hdulist[1])
				#print(hdulist[1].data.shape)
				#sys.exit()

				self.wave      = np.array(pow(10, crval1 + cdelt1 * np.arange(1, naxis1+1)))
				self.flux      = hdulist[1].data
				self.noise     = hdulist[2].data
				self.sky       = hdulist[4].data
				self.skynoise  = hdulist[5].data
				self.tell      = hdulist[6].data
				self.tellnoise = hdulist[7].data
				self.lsfcoeff  = hdulist[8].data
				self.binary    = hdulist[9].data

				# store the original parameters
				self.oriWave   = np.array(pow(10, crval1 + cdelt1 * np.arange(1, naxis1+1)))
				self.oriFlux   = hdulist[1].data
				self.oriNoise  = hdulist[2].data

				## APOGEE APVISIT has corrected the telluric absorption; the forward-modeling routine needs to put it back
				#self.flux     *= self.tell

			self.header   = hdulist[0].header
			self.header1  = hdulist[1].header
			self.header2  = hdulist[2].header
			self.header3  = hdulist[3].header

			self.model    = np.array(hdulist[3].data)
			self.mask     = []

		elif self.instrument == 'igrins':
			self.name      = kwargs.get('name')
			self.order     = kwargs.get('order')
			self.path      = kwargs.get('path')
			self.apply_sigma_mask = kwargs.get('apply_sigma_mask',False)
			#self.manaulmask = kwargs('manaulmask', False)

			if self.path == None:
				self.path = './'

			fullpath = self.path + '/' + self.name + '_' + str(self.order) + '.fits'

			hdulist = fits.open(fullpath, ignore_missing_end=True)

			#The indices 0 to 3 correspond to wavelength, flux, noise, and sky
			self.header = hdulist[0].header
			self.wave   = hdulist[0].data * 10000.0 # convert to Angstrom
			self.flux   = hdulist[1].data
			self.noise  = hdulist[2].data

		if self.apply_sigma_mask:
			# set up masking criteria
			self.avgFlux = np.mean(self.flux)
			self.stdFlux = np.std(self.flux)

			self.smoothFlux = self.flux
			# set the outliers as the flux below
			#self.smoothFlux[self.smoothFlux <= self.avgFlux - 2 * self.stdFlux] = 0
			#self.smoothFlux[ np.abs(self.smoothFlux - self.avgFlux ) <= 2 * self.stdFlux] = 0

			self.mask  = np.where(np.abs(self.flux - self.avgFlux ) >= 3. * self.stdFlux)
			#print(self.mask)

			if self.instrument == 'apogee':
				#self.mask = np.union1d(self.mask[0],np.where(self.noise >= self.flux)[0])
				noise_median = np.median(self.noise)
				self.mask = np.union1d(self.mask[0], np.where(self.noise >= 3. * noise_median)[0])
			self.wave  = np.delete(self.wave, list(self.mask))
			self.flux  = np.delete(self.flux, list(self.mask))
			self.noise = np.delete(self.noise, list(self.mask))

			if self.instrument == 'nirspec':
				self.sky   = np.delete(self.sky, list(self.mask))
			self.mask  = self.mask[0]

	def mask_custom(self, custom_mask):
		"""
		Mask the pixels by a self-defined list.
		"""
		## combine the list and remove the duplicates
		self.mask  =  list(set().union(self.mask, custom_mask))

		self.wave  = np.delete(self.oriWave, list(self.mask))
		self.flux  = np.delete(self.oriFlux, list(self.mask))
		self.noise = np.delete(self.oriNoise, list(self.mask))

		return self

	def maskBySigmas(self, sigma=2):
		"""
		Mask the outlier data points by sigmas.
		"""
		# set up masking criteria
		self.avgFlux = np.mean(self.flux)
		self.stdFlux = np.std(self.flux)

		self.smoothFlux = self.flux
		# set the outliers as the flux below
		self.smoothFlux[self.smoothFlux <= self.avgFlux - sigma * self.stdFlux] = 0

		self.mask  = np.where(self.smoothFlux <= 0)
		self.wave  = np.delete(self.wave, list(self.mask))
		self.flux  = np.delete(self.flux, list(self.mask))
		self.noise = np.delete(self.noise, list(self.mask))
		self.sky   = np.delete(self.sky, list(self.mask))
		self.mask  = self.mask[0]

	def maskByModel(self, model, sigma=3, pixel_start=30, pixel_end=-10):
		"""
		Mask the data by a forward model.
		"""
		pixel        = np.delete(np.arange(len(self.oriWave)),self.mask)[pixel_start: pixel_end]
		custom_mask2 = pixel[np.where(np.abs(self.flux-model.flux) > sigma*np.std(self.flux-model.flux))]
		print(custom_mask2)
		custom_mask2 = np.append(custom_mask2, np.array(self.mask))
		custom_mask2.sort()
		custom_mask2 = custom_mask2.tolist()
		self.mask_custom(custom_mask2)

	def plot(self, **kwargs):
		"""
		Plot the spectrum.
		"""
		#xlim   = kwargs.get('xrange', [self.wave[0], self.wave[-1]])
		#ylim   = kwargs.get('yrange', [min(self.flux)-.2, max(self.flux)+.2])
		items  = kwargs.get('items', ['spec','noise'])
		title  = kwargs.get('title')
		mask   = kwargs.get('mask', True)
		save   = kwargs.get('save', False)
		output = kwargs.get('output', str(self.name) + '.png')

		plt.figure(figsize=(16,6))
		plt.rc('font', family='sans-serif')
		## Plot masked spectrum
		if ('spectrum' in items) or ('spec' in items):
			if "_" in self.name:
				plot_name = self.name.split("_")[0]
			else:
				plot_name = self.name
			if mask:
				plt.plot(self.wave, self.flux, color='k',
					alpha=.8, linewidth=1,
					label="{} O{}".format(plot_name,self.order))
			if not mask:
				plt.plot(self.oriWave, self.oriFlux, color='k',
					alpha=.8, linewidth=1,
					label="{} O{}".format(plot_name,self.order))

		## Plot spectrum noise
		if 'noise' in items:
			if mask:
				plt.fill_between(self.wave, -self.noise, self.noise,
					color='gray', linewidth=1, alpha=.6)
			elif not mask:
				plt.fill_between(self.oriWave, -self.oriNoise,
					self.oriNoise,
					color='gray', linewidth=1, alpha=.6)

		plt.legend(fontsize=12)
		#plt.xlim(xlim)
		#plt.ylim(ylim)

		plt.xlabel('Wavelength [$\AA$]', fontsize=18)
		plt.ylabel('Flux (cnts/s)', fontsize=18)
		plt.minorticks_on()
		plt.tick_params(axis='both', labelsize=18)

		if title != None:
			plt.title(title, fontsize=20)

		if save == True:
			plt.savefig(output)

		plt.show()
		plt.close()

	def writeto(self, save_to_path, method='ascii',
		tell_sp=None):
		"""
		Save the data as an ascii or a fits file.

		Parameters
		----------
		save_to_path 	:	str
							the path to save the output file

		method 			: 	'ascii' or 'fits'
							the output file format, either in
							a single ascii file or several fits
							files labeled in the order of
							wavelength

		Optional Parameters
		-------------------
		tell_sp 		: 	Spectrum object
							the telluric data for the corresponding
							wavelength calibration

		Returns
		-------
		ascii or fits 	: 	see the method keyword
							The wavelength is in microns

		"""
		#pixel = np.delete(np.arange(1024),list(self.mask))
		pixel = np.arange(len(self.oriWave))
		## create the output mask array 0=good; 1=bad
		if (self.apply_sigma_mask) or (self.mask != []):
			mask = np.zeros((len(self.oriWave),),dtype=int)
			np.put(mask,self.mask,int(1))
		else:
			mask = np.zeros((len(self.oriWave),),dtype=int)

		if method == 'fits':
			#fullpath = self.path + '/' + self.name + '_' + str(self.order) + '_all.fits'
			#hdulist = fits.open(fullpath, ignore_missing_end=True)
			#hdulist.writeto(save_to_path)
			#hdulist.close()
			if self.header['NAXIS1'] == 1024:
				save_to_path2 = save_to_path + self.header['FILENAME'].split('.')[0]\
				+ '_O' + str(self.order)
			else:
				save_to_path2 = save_to_path + self.header['OFNAME'].split('.')[0]\
				+ '_O' + str(self.order)
			## wavelength
			hdu1 = fits.PrimaryHDU(self.wave/10000, header=self.header)
			save_to_path2_1 = save_to_path2 + '_wave.fits'
			hdu1.writeto(save_to_path2_1)
			## flux
			hdu2 = fits.PrimaryHDU(self.flux, header=self.header)
			save_to_path2_2 = save_to_path2 + '_flux.fits'
			hdu2.writeto(save_to_path2_2)
			## uncertainty
			hdu3 = fits.PrimaryHDU(self.noise, header=self.header)
			save_to_path2_3 = save_to_path2 + '_uncertainty.fits'
			hdu3.writeto(save_to_path2_3)
			## pixel
			hdu4 = fits.PrimaryHDU(pixel, header=self.header)
			save_to_path2_4 = save_to_path2 + '_pixel.fits'
			hdu4.writeto(save_to_path2_4)
			## mask
			hdu5 = fits.PrimaryHDU(mask, header=self.header)
			save_to_path2_5 = save_to_path2 + '_mask.fits'
			hdu5.writeto(save_to_path2_5)

			if tell_sp is not None:
				tell_sp2 = copy.deepcopy(tell_sp)
				# the telluric standard model
				wavelow = tell_sp2.wave[0] - 20
				wavehigh = tell_sp2.wave[-1] + 20
				tell_mdl = smart.getTelluric(wavelow=wavelow,wavehigh=wavehigh)
				# continuum correction for the data
				tell_sp2 = smart.continuumTelluric(data=tell_sp2,
					model=tell_mdl,order=tell_sp2.order)
				# telluric flux
				hdu6 = fits.PrimaryHDU(tell_sp.flux, header=tell_sp.header)
				save_to_path2_6 = save_to_path2 + '_telluric_flux.fits'
				hdu5.writeto(save_to_path2_6)
				# telluric uncertainty
				hdu7 = fits.PrimaryHDU(tell_sp.noise, header=tell_sp.header)
				save_to_path2_7 = save_to_path2 + '_telluric_uncertainty.fits'
				hdu5.writeto(save_to_path2_7)
				# telluric model
				hdu8 = fits.PrimaryHDU(tell_mdl.flux, header=tell_sp.header)
				save_to_path2_8 = save_to_path2 + '_telluric_model.fits'
				hdu5.writeto(save_to_path2_8)

		elif method == 'ascii':
			if '.txt' not in save_to_path:
				if self.header['NAXIS1'] == 1024:
					save_to_path2 = save_to_path + self.header['FILENAME'].split('.')[0]\
					+ '_O' + str(self.order) + '.txt'
				else:
					save_to_path2 = save_to_path + self.header['OFNAME'].split('.')[0]\
					+ '_O' + str(self.order) + '.txt'
			else:
				save_to_path2 = save_to_path

			if tell_sp is None:
				df = pd.DataFrame(data={'wavelength':list(self.oriWave/10000),
					'flux':list(self.oriFlux),
					'uncertainty':list(self.oriNoise),
					'pixel':list(pixel),
					'mask':list(mask)})
				df.to_csv(save_to_path2, index=None, sep='\t', mode='a',
					header=True, columns=['wavelength', 'flux', 'uncertainty',
					'pixel', 'mask'])

			elif tell_sp is not None:
				tell_sp2 = copy.deepcopy(tell_sp)
				tell_sp2 = smart.continuumTelluric(data=tell_sp2
					,order=self.order)
				lsf0 = smart.getLSF(tell_sp2)
				tell_sp2.flux = tell_sp2.oriFlux
				tell_sp2.wave = tell_sp2.oriWave
				tell_mdl = smart.convolveTelluric(lsf0, tell_sp2)

				print(len(self.oriWave), len(self.oriFlux), len(self.oriNoise), len(tell_sp.oriFlux),
					len(tell_sp.oriNoise), len(tell_mdl.flux), len(pixel), len(mask))

				df = pd.DataFrame(data={'wavelength':list(self.oriWave/10000),
					'flux':list(self.oriFlux),
					'uncertainty':list(self.oriNoise),
					'telluric_flux':list(tell_sp.oriFlux),
					'telluric_uncertainty':list(tell_sp.oriNoise),
					'telluric_model':list(tell_mdl.flux),
					'pixel':list(pixel),
					'mask':list(mask)})

				df.to_csv(save_to_path2, index=None, sep='\t', mode='a',
					header=True, columns=['wavelength', 'flux', 'uncertainty',
					'telluric_flux', 'telluric_uncertainty', 'telluric_model',
					'pixel', 'mask'])

	def coadd(self, sp, method='pixel'):
		"""
		Coadd individual extractions, either in pixel space or
		wavelength space.

		Parameters
		----------
		sp 		: 	Spectrum object
					spectrum to be coadded

		method 	: 	'pixel' or 'wavelength'
					coadd based on adding pixels or wavelength
					If 'wavelength', the second spectrum would be
					10x supersample and then cross correlated
					to be optimally shifted and coadded

		Returns
		-------
		self 	: 	Spectrum object
					coadded spectra

		"""
		if method == 'pixel':
			w1 = 1/self.oriNoise**2
			w2 = 1/sp.oriNoise**2
			self.oriFlux = (self.oriFlux*w1 + sp.oriFlux*w2)/(w1 + w2)
			self.oriNoise = np.sqrt(1/(w1 + w2))
			## set up masking criteria
			self.avgFlux = np.mean(self.oriFlux)
			self.stdFlux = np.std(self.oriFlux)
			self.smoothFlux = self.oriFlux
			## set the outliers as the flux below
			if self.apply_sigma_mask:
				self.smoothFlux[self.smoothFlux <= self.avgFlux-2*self.stdFlux] = 0
				self.mask = np.where(self.smoothFlux <= 0)
			else:
				self.mask = []
			self.wave  = np.delete(self.oriWave, list(self.mask))
			self.flux  = np.delete(self.oriFlux, list(self.mask))
			self.noise = np.delete(self.oriNoise, list(self.mask))

		elif method == 'wavelength':
			self_supers = copy.deepcopy(self)
			g = interpolate.interp1d(self.wave, self.flux)
			sp_supers = copy.deepcopy(sp)
			f = interpolate.interp1d(sp.wave, sp.flux)
			## 10x supersample the average difference of
			## the wavelength
			#step0 = np.mean(np.diff(self.wave))/10
			#self_supers.wave = np.arange(self.wave[0],
			#	self.wave[-1],step0)
			self_supers.flux = g(self_supers.wave)
			self_supers.oriWave = np.arange(self.oriWave[0],
				self.oriWave[-1],(self.oriWave[-1]-self.oriWave[0])/10240)
			g1 = interpolate.interp1d(self.oriWave, self.oriFlux)
			self_supers.oriFlux = g1(self_supers.oriWave)

			#step = np.mean(np.diff(sp.wave))/10
			#sp_supers.wave = np.arange(sp.wave[0],sp.wave[-1],step)
			#sp_supers.flux = f(sp_supers.wave)
			sp_supers.oriWave = np.arange(sp.oriWave[0],
				sp.oriWave[-1],(sp.oriWave[-1]-sp.oriWave[0])/10240)
			f1 = interpolate.interp1d(sp.oriWave, sp.oriFlux)
			sp_supers.oriFlux = f1(sp_supers.oriWave)

			## calculate the max cross correlation value
			def xcorr(a0,b0,shift):
				"""
				Shift is the index number after supersampling
				both of the spectra.
				"""
				a = copy.deepcopy(a0)
				b = copy.deepcopy(b0)

				## shift the wavelength of b
				length = b.oriFlux.shape[0]
				if shift >= 0:
					mask_a = np.arange(0,shift,1)
					a.oriFlux = np.delete(a.oriFlux,mask_a)
					mask_b = np.arange(length-1,length-shift-1,-1)
					b.oriFlux = np.delete(b.oriFlux,mask_b)

				elif shift < 0:
					mask_a = np.arange(length-1,length+shift-1,-1)
					a.oriFlux = np.delete(a.oriFlux,mask_a)
					mask_b = np.arange(0,-shift,1)
					b.oriFlux = np.delete(b.oriFlux,mask_b)

				## shift the wavelength of b
				#b.wave += shift * step
				## discard the points where the wavelength values
				## are larger
				#condition = (a.wave > b.wave[0]) & (a.wave < b.wave[-1])

				#a.flux = a.flux[np.where(condition)]
				#a.wave = a.wave[np.where(condition)]
				## resampling the telluric model
				#b.flux = np.array(smart.integralResample(xh=b.wave,
				#	yh=b.flux, xl=a.wave))

				return np.inner(a.oriFlux, b.oriFlux)/\
				(np.average(a.oriFlux)*np.average(b.oriFlux))/a.oriFlux.shape[0]

			xcorr_list = []
			## mask the ending pixels
			self_supers2 = copy.deepcopy(self_supers)
			sp_supers2 = copy.deepcopy(sp_supers)
			self_supers2.wave = self_supers2.wave[1000:-1000]
			self_supers2.flux = self_supers2.flux[1000:-1000]
			sp_supers2.wave = sp_supers2.wave[1000:-1000]
			sp_supers2.flux = sp_supers2.flux[1000:-1000]
			for shift in np.arange(-10,10,1):
				xcorr_list.append(xcorr(self_supers2,sp_supers2,shift))

			## dignostic plot for cc result
			fig, ax = plt.subplots()
			ax.plot(np.arange(-10,10,1),np.array(xcorr_list),'k-')
			plt.show()
			plt.close()

			step = np.absolute(np.mean(np.diff(sp_supers.wave)))
			bestshift = np.arange(-10*step,10*step,step)[np.argmax(xcorr_list)]
			sp_supers.oriWave += bestshift
			## discard the points where the wavelength values
			## are larger
			condition = (self.oriWave > sp_supers.oriWave[0])\
			& (self.oriWave < sp_supers.oriWave[-1])

			self.oriFlux = self.oriFlux[np.where(condition)]
			self.oriWave = self.oriWave[np.where(condition)]
			self.oriNoise = self.oriNoise[np.where(condition)]
			sp_supers.oriNoise = sp_supers.oriNoise[np.where(condition)]
			sp_supers.oriFlux = np.array(smart.integralResample(xh=sp_supers.oriWave,
				yh=sp_supers.oriFlux, xl=self.oriWave))

			w1 = 1/self.oriNoise**2
			w2 = 1/sp_supers.oriNoise**2
			self.oriFlux = (self.oriFlux*w1 + sp_supers.oriFlux*w2)/(w1 + w2)
			self.oriNoise = np.sqrt(1/(w1 + w2))
			## set up masking criteria
			self.avgFlux = np.mean(self.oriFlux)
			self.stdFlux = np.std(self.oriFlux)
			self.smoothFlux = self.oriFlux
			## set the outliers as the flux below
			self.smoothFlux[self.smoothFlux <= self.avgFlux-2*self.stdFlux] = 0
			self.mask = np.where(self.smoothFlux <= 0)
			self.wave  = np.delete(self.oriWave, list(self.mask))
			self.flux  = np.delete(self.oriFlux, list(self.mask))
			self.noise = np.delete(self.oriNoise, list(self.mask))

		return self

	def updateWaveSol(self, tell_sp):
		"""
		Return a new wavelength solution given a wavelength
		calibrated telluric spectrum.

		Parameters
		----------
		tell_sp 	: 	Spectrum object
						the calibrated telluric spectra
		"""
		wfit0 = tell_sp.header['WFIT0NEW']
		wfit1 = tell_sp.header['WFIT1NEW']
		wfit2 = tell_sp.header['WFIT2NEW']
		wfit3 = tell_sp.header['WFIT3NEW']
		wfit4 = tell_sp.header['WFIT4NEW']
		wfit5 = tell_sp.header['WFIT5NEW']
		c3    = tell_sp.header['c3']
		c4    = tell_sp.header['c4']

		length1 = tell_sp.header['NAXIS1']

		self.wave = np.delete(smart.waveSolution(np.arange(length1),
			wfit0,wfit1,wfit2,wfit3,wfit4,wfit5,c3,c4, order=self.order), list(self.mask))
		self.oriWave = smart.waveSolution(np.arange(length1),
			wfit0,wfit1,wfit2,wfit3,wfit4,wfit5,c3,c4, order=self.order)

		return self


# -*- encoding: utf-8 -*-
"""
@File : tools.py
@Time : 2019/03/04 08:35:46
@Author : wangtf
@Version : 1.0
@Desc : None
"""

# here put the import lib
import os
import glob
import numpy as np

def voc_ap(recall, precision, use_07_metric=False):
    """
    ap = voc_ap(recall, precision, [use_07_metric])

    Compute VOC AP given precision and recall.
    If use_07_metric is true, uses  the
    VOC 07 11 point method (default: False).
    Please make shure that recall and precison are sorted by scores.

    Args:
        recall: the shape of (n,) ndarray;
        precision: the shape of (n,) ndarray;
        use_07_metric: if true, the 11 points method will be used.
    Returns:
        the float number result of average precision.
    """
    if use_07_metric:
        # 11 point metric
        ap = 0.
        for t in np.arange(0., 1.1, 0.1):
            if np.sum(recall >= t) == 0:
                p = 0
            else:
                p = np.max(precision[recall >= t])
            ap = ap + p / 11.
    else:
        # correct AP calculation
        # first append sentinel values at the end
        mrec = np.concatenate(([0.], recall, [1.]))
        mpre = np.concatenate(([0.], precision, [0.]))

        # compute the precision envelope
        for i in range(mpre.size - 1, 0, -1):
            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

        # to calculate area under PR curve, look for points
        # where X axis (recall) changes value
        i = np.where(mrec[1:] != mrec[:-1])[0]

        # and sum (\Delta recall) * prec
        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap

def compute_overlaps(boxes, one_box):
    """
    iou = compute_overlaps(boxes, one_box)

    compute intersection over union of ndarray.
    The format of one_box is [xmin, ymin, xmax, ymax].

    Args:
        boxes: the (n, 4) shape ndarray, ground truth boundboxes;
        bb: the (4,) shape ndarray, detected boundboxes;
    Returns:
        a (n, ) shape ndarray.
    """
    # compute overlaps
    # intersection
    ixmin = np.maximum(boxes[:, 0], one_box[0])
    iymin = np.maximum(boxes[:, 1], one_box[1])
    ixmax = np.minimum(boxes[:, 2], one_box[2])
    iymax = np.minimum(boxes[:, 3], one_box[3])
    iw = np.maximum(ixmax - ixmin + 1., 0.)
    ih = np.maximum(iymax - iymin + 1., 0.)
    inters = iw * ih

    # union
    boxes_area = (boxes[:, 2] - boxes[:, 0] + 1.) * (boxes[:, 3] -
                                                     boxes[:, 1] + 1.)
    one_box_area = (one_box[2] - one_box[0] + 1.) * (one_box[3] - one_box[1] +
                                                     1.)
    iou = inters / (one_box_area + boxes_area - inters)

    return iou

def voc_eval(class_recs: dict,
             detect: dict,
             iou_thresh: float = 0.5,
             use_07_metric: bool = False):
    """
    recall, precision, ap = voc_eval(class_recs, detection,
                                [iou_thresh],
                                [use_07_metric])

    Top level function that does the PASCAL VOC evaluation.
    Please make sure that the class_recs only have one class annotations.

    precision = tp / (tp + fp)
    recall = tp / (tp + fn)

    Args:
        class_recalls: recalls dict of a class
            class_recs[image_name]={'bbox': []}.
        detection: Path to annotations
            detection={'image_ids':[], bbox': [], 'confidence':[]}.
        [iou_thresh]: Overlap threshold (default = 0.5)
        [use_07_metric]: Whether to use VOC07's 11 point AP computation
            (default False)
    Returns:
        a dict of result including true_positive_number, false_positive_number,
        recall, precision and average_precision.
    Raises:
        TypeError: the data format is not np.ndarray.
    """
    # format data
    # class_rec data load
    npos = 0
    for imagename in class_recs.keys():
        if not isinstance(class_recs[imagename]['bbox'], np.ndarray):
            raise TypeError
        detected_num = class_recs[imagename]['bbox'].shape[0]
        npos += detected_num
        class_recs[imagename]['det'] = [False] * detected_num

    # detections data load
    image_ids = detect['image_ids']
    confidence = detect['confidence']
    BB = detect['bbox']
    if not isinstance(confidence, np.ndarray):
        raise TypeError
    if not isinstance(BB, np.ndarray):
        raise TypeError

    # sort by confidence
    sorted_ind = np.argsort(-confidence)
    BB = BB[sorted_ind, :]
    image_ids = [image_ids[x] for x in sorted_ind]

    # go down dets and mark TPs and FPs
    nd = len(image_ids)
    tp = np.zeros(nd)
    fp = np.zeros(nd)
    for d in range(nd):
        R = class_recs[image_ids[d]]
        bb = BB[d, :].astype(float)
        iou_max = -np.inf
        BBGT = R['bbox'].astype(float)

        if BBGT.size > 0:
            overlaps = compute_overlaps(BBGT, bb)
            iou_max = np.max(overlaps)
            iou_max_index = np.argmax(overlaps)

        if iou_max > iou_thresh:
            if not R['det'][iou_max_index]:
                tp[d] = 1.
                R['det'][iou_max_index] = 1
            else:
                fp[d] = 1.
        else:
            fp[d] = 1.

    # compute precision recall
    fp = np.cumsum(fp)
    tp = np.cumsum(tp)
    true_positive_number = tp[-1]
    false_positive_number = fp[-1]

    recall = tp / float(npos)
    # avoid divide by zero in case the first detection matches
    # a difficult ground truth
    precision = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)
    average_precision = voc_ap(recall, precision, use_07_metric)

    result = {}
    result['true_positive_number'] = true_positive_number
    result['false_positive_number'] = false_positive_number
    result['recall'] = recall
    result['precision'] = precision
    result['average_precision'] = average_precision
    return result

def voc_eval_files(class_recs_dir,
                   detect_file,
                   label_id,
                   iou_thresh=0.5,
                   use_07_metric=False):
    """
    recall, precision, ap = voc_eval(class_recs, detection,
                                [iou_thresh],
                                [use_07_metric])

    Top level function that does the PASCAL VOC evaluation.
    Please make sure that the class_recs only have one class annotations.

    precision = tp / (tp + fp)
    recall = tp / (tp + fn)

    Args:
        class_recalls: recalls dict of a class
            class_recs[image_name]={'bbox': []}.
        detection: Path to annotations
            detection={'image_ids':[], bbox': [], 'confidence':[]}.
        [iou_thresh]: Overlap threshold (default = 0.5)
        [use_07_metric]: Whether to use VOC07's 11 point AP computation
            (default False)
    Returns:
        a dict of result including true_positive_number, false_positive_number,
        recall, precision and average_precision.
    Raises:
        IOError: can not find the path.
    """
    if not os.path.exists(class_recs_dir):
        raise IOError
    if not os.path.exists(detect_file):
        raise IOError

    class_recs = {}
    recs_list = glob.glob(os.path.join(class_recs_dir, '*.txt'))
    for path in recs_list:
        image_id = os.path.basename(path)[:-4]
        with open(path) as f:
            data = f.read().strip().split('\n')
            bboxes = []
            for line in data:
                label, xmin, ymin, xmax, ymax = line.strip().split(' ')
                if label == str(label_id):
                    bboxes.append([xmin, ymin, xmax, ymax])
            bboxes = np.array(bboxes)
            class_recs[image_id] = {'bbox': bboxes}

    detection = {'image_ids': [], 'bbox': [], 'confidence': []}
    with open(detect_file) as f:
        data = f.read().strip().split('\n')
        for line in data:
            image_id, confidence, xmin, ymin, xmax, ymax = line.strip().split()
            detection['image_ids'].append(image_id)
            detection['confidence'].append(confidence)
            detection['bbox'].append([xmin, ymin, xmax, ymax])
    detection['image_ids'] = np.array(detection['image_ids'])
    detection['confidence'] = np.array(detection['confidence'])
    detection['bbox'] = np.array(detection['bbox'])

    result = voc_eval(class_recs,
                      detection,
                      iou_thresh=iou_thresh,
                      use_07_metric=use_07_metric)
    return result


# Copyright 2020 Arthur Coqué, Guillaume Morin, Pôle OFB-INRAE ECLA, UR RECOVER
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""This module gathers wc algorithms used for estimating Chl-a concentrations.

Each class of this module correspond to one algorithm. An algorithm can have
several calibrations (a calibration is a set of parameters), either
packaged within SISPPEO (these default calibrations are located in
'resources/wc_algo_calibration') or provided by the user.
Before its utilisation, an algorithm has to be instantiate with specific
settings like the product_type of further input products, the calibration
used, the band used (if needed), etc.

Example:

    algo1 = CHLAGons('S2_GRS', 'Gons_2004')
    out_array1 = algo1(red_array, rededge_array, nir_array, 'rho')

    algo2 = CHLAGittelson('L8_GRS', '3_bands', 'Gitelson_2008')
    out_array2 = algo2(red_array, rededge_array, nir_array, 'rrs')
"""

from pathlib import Path
from typing import Optional, Union

import numpy as np
import xarray as xr

from sisppeo.utils.algos import load_calib, producttype_to_sat
from sisppeo.utils.config import wc_algo_config as algo_config, wc_calib
from sisppeo.utils.exceptions import InputError

# pylint: disable=invalid-name
# Ok for a custom type.
P = Union[str, Path]
N = Union[int, float]

class CHLAGons:
    """Chlorophyll-a concentration (in mg/m3) from 3 red bands after Gons et al., 1999, 2002, 2004

    Red edge algorithm to retrieve Chlorophyll-a concentration (in mg/m3) from
    surface reflectances (rho, unitless) or remote sensing reflectances (Rrs,
    in sr-1) at 665nm B4 MSI, 704nm B5 MSI and 783nm B7 MSI.
    This algorithm was published in Gons et al., 1999, 2002, 2004

    Attributes:
        name: The name of the algorithm used. This is the key used by
            L3AlgoBuilder and that you must provide in config or when using
            the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: A dict of metadata (calibration name, model coefficients, etc).
    """
    _default_calibration_file = wc_calib / 'chla-gons.yaml'
    _default_calibration_name = 'Gons_2004'
    name = 'chla-gons'

    def __init__(self,
                 product_type: str,
                 calibration: Optional[P] = None,
                 **_ignored) -> None:
        """Inits an 'CHLAGons' instance with specific settings.

        Args:
            product_type: The type of the input satellite product (e.g.
                S2_ESA_L2A or L8_USGS_L1GT)
            calibration: Optional; The calibration (set of parameters) used by
                the algorithm (default=_default_calibration_name).
            **_ignored: Unused kwargs sent to trash.
        """
        try:
            self.requested_bands = algo_config[self.name][
                producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with {self.name}'
            raise InputError(msg) from invalid_product
        calibration_dict, calibration_name = load_calib(
            calibration,
            self._default_calibration_file,
            self._default_calibration_name
        )
        self._valid_limit = calibration_dict['validity_limit']
        try:
            params = calibration_dict[producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with this calibration'
            raise InputError(msg) from invalid_product
        self.__dict__.update(params)
        self.meta = {'calibration': calibration_name,
                     'validity_limit': self._valid_limit,
                     **params}

    def __call__(self,
                 ref_red: xr.DataArray,
                 ref_rededge: xr.DataArray,
                 ref_nir: xr.DataArray,
                 data_type: str,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('ref').

        Args:
            ref_red: An array (dimension 1 * N * M) of 'data_type'.
            ref_redegde: An array (dimension 1 * N * M) of 'data_type'.
            ref_nir: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'ref' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of chl-a (in mg/m3).
        """

        if data_type == 'rho':
            ref_red = ref_red / np.pi
            ref_rededge = ref_rededge / np.pi
            ref_nir = ref_nir / np.pi

        np.warnings.filterwarnings('ignore')
        ref_red = ref_red.where(ref_red >= 0)
        ref_rededge = ref_rededge.where(ref_red >= 0)
        ref_nir = ref_red.where(ref_nir >= 0)

        bb783 = ref_nir.where(ref_nir >= 0).copy()
        # pylint: disable=no-member
        # Loaded in __init__ whit "__dict__.update".
        bb783 = (self.a * bb783) / (0.082 - 0.6 * bb783)
        aphy = ref_rededge / ref_red * (self.aw705 + bb783) - self.aw665 \
            - np.power(bb783, self.p)
        chla = aphy / self.aphy_star
        chla = chla.where((chla >= 0) & (chla <= self._valid_limit))
        return chla

class CHLAGitelson:
    """Chlorophyll-a concentration (in mg/m3) from 3 red bands after Gitelson et al., 2008

    Red edge algorithm to retrieve Chlorophyll-a concentration (in mg/m3) from
    surface reflectances (rho, unitless) or remote sensing reflectances (Rrs,
    in sr-1) at 665nm B4 MSI, 705nm B5 MSI and 740nm B6 MSI.
    This algorithm was published in Gitelson et al., 2008

    Attributes:
        name: The name of the algorithm used. This is the key used by
            L3AlgoBuilder and that you must provide in config or when using
            the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: A dict of metadata (calibration name, model coefficients, etc).
    """
    _default_calibration_file = wc_calib / 'chla-gitelson.yaml'
    _default_calibration_name = 'Gitelson_2008'
    _default_design = '3_bands'
    name = 'chla-gitelson'

    def __init__(self,
                 product_type: str,
                 design: str = _default_design,
                 calibration: Optional[P] = None,
                 **_ignored) -> None:
        """Inits an 'CHLAGitelson' instance with specific settings.

        Args:
            product_type: The type of the input satellite product (e.g.
                S2_ESA_L2A or L8_USGS_L1GT)
            calibration: The calibration (set of parameters) used by the
                algorithm (default=_default_calibration_name).
            **_ignored: Unused kwargs sent to trash.
        """
        self._design = design
        try:
            self.requested_bands = algo_config[self.name][
                producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with {self.name}'
            raise InputError(msg) from invalid_product
        calibration_dict, calibration_name = load_calib(
            calibration,
            self._default_calibration_file,
            self._default_calibration_name
        )
        self._valid_limit = calibration_dict['validity_limit']
        try:
            params = calibration_dict[producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with this calibration'
            raise InputError(msg) from invalid_product
        self.__dict__.update(params)
        self.meta = {'calibration': calibration_name,
                     'design': design,
                     'validity_limit': self._valid_limit,
                     **params}

    def __call__(self,
                 ref_red: xr.DataArray,
                 ref_rededge: xr.DataArray,
                 ref_nir: xr.DataArray,
                 data_type: str,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('ref').

        Args:
            ref_red: An array (dimension 1 * N * M) of 'data_type'.
            ref_redegde: An array (dimension 1 * N * M) of 'data_type'.
            ref_nir: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'rho' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of chl-a (in mg/m3).
        """

        if data_type == 'rho':
            ref_red = ref_red / np.pi
            ref_rededge = ref_rededge / np.pi
            ref_nir = ref_nir / np.pi

        np.warnings.filterwarnings('ignore')
        ref_red = ref_red.where(ref_red >= 0)
        ref_rededge = ref_rededge.where(ref_red >= 0)
        ref_nir = ref_red.where(ref_nir >= 0)
        print(self._design, self._valid_limit)
        if self._design == '3_bands':
            print('3 bands selected')
            # pylint: disable=no-member
            # Loaded in __init__ whit "__dict__.update".
            chla = self.a_3bands + self.b_3bands \
                * (1 / ref_red - 1 / ref_rededge) * ref_nir
        else:
            print('2 bands selected')
            # pylint: disable=no-member
            # Loaded in __init__ whit "__dict__.update".
            chla = self.a_2bands + self.b_2bands * (1 / ref_red) * ref_nir
        chla = chla.where((chla >= 0) & (chla <= self._valid_limit))
        return chla

class CHLAGurlin:
    """Chlorophyll-a concentration (in mg/m3) from 3 red bands after Gurlin et al., 2011

    Red edge algorithm to retrieve Chlorophyll-a concentration (in mg/m3) from
    surface reflectances (rho, unitless) or remote sensing reflectances (Rrs,
    in sr-1) at 665nm B4 MSI, 704nm B5 MSI and 783nm B7 MSI.
    This algorithm was published in Gons et al., 1999, 2002, 2004

    Attributes:
        name: The name of the algorithm used. This is the key used by
            L3AlgoBuilder and that you must provide in config or when using
            the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: A dict of metadata (calibration name, model coefficients, etc).
    """
    _default_calibration_file = wc_calib / 'chla-gurlin.yaml'
    _default_calibration_name = 'Gurlin_2011'
    name = 'chla-gurlin'

    def __init__(self,
                 product_type: str,
                 calibration: Optional[P] = None,
                 **_ignored) -> None:
        """Inits an 'CHLAGurlin' instance with specific settings.

        Args:
            product_type: The type of the input satellite product (e.g.
                S2_ESA_L2A or L8_USGS_L1GT)
            calibration: The calibration (set of parameters) used by the
                algorithm (default=_default_calibration_name).
            **_ignored: Unused kwargs sent to trash.
        """
        try:
            self.requested_bands = algo_config[self.name][
                producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with {self.name}'
            raise InputError(msg) from invalid_product
        calibration_dict, calibration_name = load_calib(
            calibration,
            self._default_calibration_file,
            self._default_calibration_name
        )
        self._valid_limit = calibration_dict['validity_limit']
        try:
            params = calibration_dict[producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with this calibration'
            raise InputError(msg) from invalid_product
        self.__dict__.update(params)
        self.meta = {'calibration': calibration_name,
                     'validity_limit': self._valid_limit,
                     **params}

    def __call__(self,
                 ref_red: xr.DataArray,
                 ref_rededge: xr.DataArray,
                 data_type: str,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('ref').

        Args:
            ref_red: An array (dimension 1 * N * M) of 'data_type'.
            ref_redegde: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'rho' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of chl-a (in mg/m3).
        """
        np.warnings.filterwarnings('ignore')
        ref_red = ref_red.where(ref_red >= 0)
        ref_rededge = ref_rededge.where(ref_red >= 0)
        # pylint: disable=no-member
        # Loaded in __init__ whit "__dict__.update".
        chla = self.a * pow(ref_rededge / ref_red, 2) + self.b \
            * (ref_rededge / ref_red) + self.c
        chla = chla.where((chla >= 0) & (chla <= self._valid_limit))
        return chla

class CHLAOC:
    """Chlorophyll-a concentration (in mg/m3) from polynomial maximum band ratio by O'Reilly et al., 1998 and updates

    Blue/green algorithm to retrieve Chlorophyll-a concentration (in mg/m3) from
    surface reflectances (rho, unitless) or remote sensing reflectances (Rrs,
    in sr-1)
    This algorithm was published in O'Reilly 1998, 2000
    calibration OC2 for OLI from Franz et al., 2015, OC3 for OLI O'Reilly and Werdell, 2019
    MSI Pahlevan et al., 2020 after O'Reilly and Werdell, 2019

    Attributes:
        name: The name of the algorithm used. This is the key used by
            L3AlgoBuilder and that you must provide in config or when using
            the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: A dict of metadata (calibration name, model coefficients, etc).
    """
    _default_calibration_file = wc_calib / 'chla-oc.yaml'
    _default_calibration_name = 'OC3'
    name = 'chla-oc'

    def __init__(self,
                 product_type: str,
                 calibration: Optional[P] = None,
                 **_ignored) -> None:
        """Inits an 'CHLAOC' instance with specific settings.

        Args:
            product_type: The type of the input satellite product (e.g.
                S2_ESA_L2A or L8_USGS_L1GT)
            calibration: The calibration (set of parameters) used by the
                algorithm (default=_default_calibration_name).
            **_ignored: Unused kwargs sent to trash.
        """
        try:
            self.requested_bands = algo_config[self.name][
                producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with {self.name}'
            raise InputError(msg) from invalid_product
        calibration_dict, calibration_name = load_calib(
            calibration,
            self._default_calibration_file,
            self._default_calibration_name
        )
        self._valid_limit = calibration_dict['validity_limit']
        self._version = calibration_name
        try:
            params = calibration_dict[producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with this calibration'
            raise InputError(msg) from invalid_product
        self.__dict__.update(params)
        self.meta = {'calibration': self._version,
                     'validity_limit': self._valid_limit,
                     **params}

    def __call__(self,
                 ref_violet: xr.DataArray,
                 ref_blue: xr.DataArray,
                 ref_green: xr.DataArray,
                 data_type: str,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('ref').

        Args:
            ref_violet: An array (dimension 1 * N * M) of 'data_type'.
            ref_blue: An array (dimension 1 * N * M) of 'data_type'.
            ref_green: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'ref' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of chl-a (in mg/m3).
        """

        np.warnings.filterwarnings('ignore')
        if data_type == 'rho':
            ref_violet = ref_violet / np.pi
            ref_blue = ref_blue / np.pi
            ref_green = ref_green / np.pi

        if self._version == 'OC3':
            print(f'{self._version} is used')
            max_ratio = np.log(np.maximum(ref_violet.values, ref_blue.values)
                               / ref_green)
            # np.log(max(Rrs_B1, Rrs_B2) / Rrs_B3))
        else:   # self._version == 'OC2'
            print(f'{self._version} is used')
            max_ratio = np.log(ref_blue.values / ref_green)
        # pylint: disable=no-member
        # Loaded in __init__ whit "__dict__.update".
        chla = np.power(10, self.a0 + self.a1 * max_ratio + self.a2
                        * np.power(max_ratio, 2) + self.a3
                        * np.power(max_ratio, 3) + self.a4
                        * np.power(max_ratio, 4))
        chla = chla.where((chla >= 0) & (chla <= self._valid_limit))
        return chla

class CHLALins:
    """Chlorophyll-a concentration (in mg/m3) from NIR/Red bands ratio after Lins et al., 2017

    Red edge algorithm to retrieve Chlorophyll-a concentration (in mg/m3) from
    surface reflectances (rho, unitless) or remote sensing reflectances (Rrs,
    in sr-1) at 665nm B4 MSI, 705nm B5 MSI
    This algorithm was published in Lins et al., 2017

    Attributes:
        name: The name of the algorithm used. This is the key used by
          L3AlgoBuilder and that you must provide in config or when using
          the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: A dict of metadata (calibration name, model coefficients, etc).
    """
    _default_calibration_file = wc_calib / 'chla-lins.yaml'
    _default_calibration_name = 'Lins_2017'
    name = 'chla-lins'

    def __init__(self,
                 product_type: str,
                 calibration: Optional[P] = None,
                 **_ignored) -> None:
        """Inits an 'CHLALins' instance with specific settings.

        Args:
            product_type: The type of the input satellite product (e.g.
                S2_ESA_L2A or L8_USGS_L1GT)
            calibration: The calibration (set of parameters) used by the
                algorithm (default=_default_calibration_name).
            **_ignored: Unused kwargs sent to trash.
        """
        try:
            self.requested_bands = algo_config[self.name][
                producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with {self.name}'
            raise InputError(msg) from invalid_product
        calibration_dict, calibration_name = load_calib(
            calibration,
            self._default_calibration_file,
            self._default_calibration_name
        )
        self._valid_limit = calibration_dict['validity_limit']
        try:
            params = calibration_dict[producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with this calibration'
            raise InputError(msg) from invalid_product
        self.__dict__.update(params)
        self.meta = {'calibration': calibration_name,
                     'validity_limit': self._valid_limit,
                     **params}

    def __call__(self,
                 ref_red: xr.DataArray,
                 ref_rededge: xr.DataArray,
                 data_type: str,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('ref').

        Args:
            ref_red: An array (dimension 1 * N * M) of 'data_type'.
            ref_redegde: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'rho' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of chl-a (in mg/m3).
        """
        np.warnings.filterwarnings('ignore')
        ref_red = ref_red.where(ref_red >= 0)
        ref_rededge = ref_rededge.where(ref_red >= 0)
        # pylint: disable=no-member
        # Loaded in __init__ whit "__dict__.update".
        chla = self.p * (ref_rededge / ref_red) + self.q
        chla = chla.where((chla >= 0) & (chla <= self._valid_limit))
        return chla

class NDCI:
    """Normalized Difference Chlorophyll Index

    NDCI from surface reflectances (rho, unitless) or remote sensing reflectances (Rrs,
    in sr-1) at 665nm B4 MSI, 704nm B5 MSI

    Attributes:
        name: The name of the algorithm used. This is the key used by
            L3AlgoBuilder and that you must provide in config or when using
            the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: An empty dict, since there is no parametrisation for NDWI.
    """
    name = 'ndci'

    def __init__(self, product_type: str, **_ignored) -> None:
        """Inits an 'Ndci' instance for a given 'product_type'.

        Args:
            product_type: The type of the input satellite product (e.g.
              S2_ESA_L2A or L8_USGS_L1GT)
            **_ignored: Unused kwargs send to trash.
        """
        try:
            self.requested_bands = algo_config[self.name][
                producttype_to_sat(product_type)]
        except KeyError as unvalid_product:
            msg = f'{product_type} is not allowed with {self.name}'
            raise InputError(msg) from unvalid_product
        self.meta = {}

    def __call__(self,
                 ref_red: xr.DataArray,
                 ref_nir: xr.DataArray,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('ref').

        Args:
            ref_red: An array (dimension 1 * N * M) of 'data_type'.
            ref_nir: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'rho' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of NDCI values.
        """
        np.warnings.filterwarnings('ignore')
        red = ref_red.where(ref_red >= 0)
        nir = ref_nir.where(ref_nir >= 0)

        return (nir - red) / (nir + red)


# -*- coding: utf-8 -*-
#
# Copyright (c) 2020, the cclib development team
#
# This file is part of cclib (http://cclib.github.io) and is distributed under
# the terms of the BSD 3-Clause License.

"""Calculation of DDEC charges based on data parsed by cclib."""
import copy
import random
import numpy
import logging
import math
import os
import sys

from cclib.method.calculationmethod import Method
from cclib.method.volume import electrondensity_spin
from cclib.parser.utils import convertor
from cclib.parser.utils import find_package

from typing import List

class MissingInputError(Exception):
    pass

class DDEC6(Method):
    """DDEC6 charges."""

    # All of these are required for DDEC6 charges.
    required_attrs = ("homos", "mocoeffs", "nbasis", "gbasis")

    def __init__(
        self, data, volume, proatom_path=None, progress=None, loglevel=logging.INFO, logname="Log"
    ):
        # Inputs are:
        # data -- ccData object that describe target molecule.
        # volume -- Volume object that describe target Cartesian grid.
        # proatom_path -- path to proatom densities
        #      (directory containing atoms.h5 in horton or c2_001_001_000_400_075.txt in chargemol)
        super(DDEC6, self).__init__(data, progress, loglevel, logname)

        self.volume = volume
        self.fragresults = None
        self.proatom_path = proatom_path

        if numpy.sum(self.data.coreelectrons) != 0:
            # TODO: Pseudopotentials should be added back
            pass

        # Check whether proatom_path is a valid directory or not.
        assert os.path.isdir(
            proatom_path
        ), "Directory that contains proatom densities should be added as an input."

        # Read in reference charges.
        self.proatom_density = []
        self.radial_grid_r = []
        for atom_number in self.data.atomnos:
            density, r = self._read_proatom(proatom_path, atom_number, 0)
            self.proatom_density.append(density)
            self.radial_grid_r.append(r)

    def __str__(self):
        """Return a string representation of the object."""
        return "DDEC6 charges of {}".format(self.data)

    def __repr__(self):
        """Return a representation of the object."""
        return "DDEC6({})".format(self.data)

    def _check_required_attributes(self):
        super(DDEC6, self)._check_required_attributes()

    def _cartesian_dist(self, pt1, pt2):
        """ Small utility function that calculates Euclidian distance between two points
            pt1 and pt2 are numpy arrays representing a point in Cartesian coordinates. """
        return numpy.sqrt(numpy.dot(pt1 - pt2, pt1 - pt2))

    def _read_proatom(
        self, directory, atom_num, charge  # type = str  # type = int  # type = float
    ):
        # type: (...) -> numpy.ndarray, numpy.ndarray
        """Return a list containing proatom reference densities."""
        # TODO: Treat calculations with psuedopotentials
        # TODO: Modify so that proatom densities are read only once for horton
        #       [https://github.com/cclib/cclib/pull/914#discussion_r464039991]
        # File name format:
        #   ** Chargemol **
        #       c2_[atom number]_[nuclear charge]_[electron count]_[cutoff radius]_[# shells]
        #   ** Horton **
        #       atoms.h5
        # File format:
        #   Starting from line 13, each line contains the charge densities for each shell
        # If `charge` is not an integer, proatom densities have to be linearly interpolated between
        # the densities of the ion/atom with floor(charge) and ceiling(charge)
        charge_floor = int(math.floor(charge))
        charge_ceil = int(math.ceil(charge))

        chargemol_path_floor = os.path.join(
            directory,
            "c2_{:03d}_{:03d}_{:03d}_500_100.txt".format(
                atom_num, atom_num, atom_num - charge_floor
            ),
        )
        chargemol_path_ceil = os.path.join(
            directory,
            "c2_{:03d}_{:03d}_{:03d}_500_100.txt".format(
                atom_num, atom_num, atom_num - charge_ceil
            ),
        )
        horton_path = os.path.join(directory, "atoms.h5")

        if os.path.isfile(chargemol_path_floor) or os.path.isfile(chargemol_path_ceil):
            # Use chargemol proatom densities
            # Each shell is .05 angstroms apart (uniform).
            # *scalefactor* = 10.58354497764173 bohrs in module_global_parameter.f08
            if atom_num <= charge_floor:
                density_floor = numpy.array([0])
            else:
                density_floor = numpy.loadtxt(chargemol_path_floor, skiprows=12, dtype=float)
            if atom_num >= charge_ceil:
                density_ceil = numpy.array([0])
            else:
                density_ceil = numpy.loadtxt(chargemol_path_ceil, skiprows=12, dtype=float)

            density = (charge_ceil - charge) * density_floor + (
                charge - charge_floor
            ) * density_ceil
            radiusgrid = numpy.arange(1, len(density) + 1) * 0.05

        elif os.path.isfile(horton_path):
            # Use horton proatom densities
            assert find_package("h5py"), "h5py is needed to read in proatom densities from horton."

            import h5py

            with h5py.File(horton_path, "r") as proatomdb:
                if atom_num <= charge_floor:
                    density_floor = numpy.array([0])
                    radiusgrid = numpy.array([0])
                else:
                    keystring_floor = "Z={}_Q={:+d}".format(atom_num, charge_floor)
                    density_floor = numpy.asanyarray(list(proatomdb[keystring_floor]["rho"]))

                    # gridspec is specification of integration grid for proatom densities in horton.
                    # Example -- ['PowerRTransform', '1.1774580743206259e-07', '20.140888089596444', '41']
                    #   is constructed using PowerRTransform grid
                    #   with rmin = 1.1774580743206259e-07
                    #        rmax = 20.140888089596444
                    #   and  ngrid = 41
                    # PowerRTransform is default in horton-atomdb.py.
                    gridtype, gridmin, gridmax, gridn = (
                        proatomdb[keystring_floor].attrs["rtransform"].split()
                    )
                    gridmin = convertor(float(gridmin), "bohr", "Angstrom")
                    gridmax = convertor(float(gridmax), "bohr", "Angstrom")
                    gridn = int(gridn)
                    # Convert byte to string in Python3
                    if sys.version[0] == "3":
                        gridtype = gridtype.decode("UTF-8")

                    # First verify that it is one of recognized grids
                    assert gridtype in [
                        "LinearRTransform",
                        "ExpRTransform",
                        "PowerRTransform",
                    ], "Grid type not recognized."

                    if gridtype == "LinearRTransform":
                        # Linear transformation. r(t) = rmin + t*(rmax - rmin)/(npoint - 1)
                        gridcoeff = (gridmax - gridmin) / (gridn - 1)
                        radiusgrid = gridmin + numpy.arange(1, gridn + 1) * gridcoeff
                    elif gridtype == "ExpRTransform":
                        # Exponential transformation. r(t) = rmin*exp(t*log(rmax/rmin)/(npoint - 1))
                        gridcoeff = math.log(gridmax / gridmin) / (gridn - 1)
                        radiusgrid = gridmin * numpy.exp(numpy.arange(1, gridn + 1) * gridcoeff)
                    elif gridtype == "PowerRTransform":
                        # Power transformation. r(t) = rmin*t^power
                        # with  power = log(rmax/rmin)/log(npoint)
                        gridcoeff = math.log(gridmax / gridmin) / math.log(gridn)
                        radiusgrid = gridmin * numpy.power(numpy.arange(1, gridn + 1), gridcoeff)

                if atom_num <= charge_ceil:
                    density_ceil = numpy.array([0])
                else:
                    keystring_ceil = "Z={}_Q={:+d}".format(atom_num, charge_ceil)
                    density_ceil = numpy.asanyarray(list(proatomdb[keystring_ceil]["rho"]))

                density = (charge_ceil - charge) * density_floor + (
                    charge - charge_floor
                ) * density_ceil

                del h5py

        else:
            raise MissingInputError("Pro-atom densities were not found in the specified path.")

        if charge == charge_floor:
            density = density_floor

        return density, radiusgrid

    def calculate(self, indices=None, fupdate=0.05):
        """
        Calculate DDEC6 charges based on doi: 10.1039/c6ra04656h paper.
        Cartesian, uniformly spaced grids are assumed for this function.
        """

        # Obtain charge densities on the grid if it does not contain one.
        if not numpy.any(self.volume.data):
            self.logger.info("Calculating charge densities on the provided empty grid.")
            if len(self.data.mocoeffs) == 1:
                self.chgdensity = electrondensity_spin(
                    self.data, self.volume, [self.data.mocoeffs[0][: self.data.homos[0]]]
                )
                self.chgdensity.data *= 2
            else:
                self.chgdensity = electrondensity_spin(
                    self.data,
                    self.volume,
                    [
                        self.data.mocoeffs[0][: self.data.homos[0]],
                        self.data.mocoeffs[1][: self.data.homos[1]],
                    ],
                )
        # If charge densities are provided beforehand, log this information
        # `Volume` object does not contain (nor rely on) information about the constituent atoms.
        else:
            self.logger.info("Using charge densities from the provided Volume object.")
            self.chgdensity = self.volume

        # STEP 1
        # Carry out step 1 of DDEC6 algorithm [Determining ion charge value]
        # Refer to equations 49-57 in doi: 10.1039/c6ra04656h
        self.logger.info("Creating first reference charges.")
        ref, loc, stock = self.calculate_refcharges()
        self.refcharges = [ref]
        self._localizedcharges = [loc]
        self._stockholdercharges = [stock]

        # STEP 2
        # Load new proatom densities.
        self.logger.info("Creating second reference charges.")
        self.proatom_density = []
        self.radial_grid_r = []
        for i, atom_number in enumerate(self.data.atomnos):
            density, r = self._read_proatom(
                self.proatom_path, atom_number, float(self.refcharges[0][i])
            )
            self.proatom_density.append(density)
            self.radial_grid_r.append(r)

        # Carry out step 2 of DDEC6 algorithm [Determining ion charge value again]
        ref, loc, stock = self.calculate_refcharges()
        self.refcharges.append(ref)
        self._localizedcharges.append(loc)
        self._stockholdercharges.append(stock)

        # STEP 3
        # Load new proatom densities.
        self.proatom_density = []
        self.radial_grid_r = []
        for i, atom_number in enumerate(self.data.atomnos):
            density, r = self._read_proatom(
                self.proatom_path, atom_number, float(self.refcharges[1][i])
            )
            self.proatom_density.append(density)
            self.radial_grid_r.append(r)

        # Carry out step 3 of DDEC6 algorithm [Determine conditioned charge density and tau]
        self.logger.info("Conditioning charge densities.")
        self.condition_densities()

    def calculate_refcharges(self):
        """ Calculate reference charges from proatom density and molecular density
            [STEP 1 and 2]
        """
        # Generator object to iterate over the grid
        xshape, yshape, zshape = self.chgdensity.data.shape
        atoms = len(self.data.atomnos)
        indices = (
            (i, x, y, z)
            for i in range(atoms)
            for x in range(xshape)
            for y in range(yshape)
            for z in range(zshape)
        )

        stockholder_w = numpy.zeros((atoms, xshape, yshape, zshape))
        localized_w = numpy.zeros((atoms, xshape, yshape, zshape))
        self.closest_r_index = numpy.zeros((atoms, xshape, yshape, zshape), dtype=int)

        for atomi, xindex, yindex, zindex in indices:
            # Distance of the grid from atom grid
            dist_r = self._cartesian_dist(
                self.data.atomcoords[-1][atomi],
                self.chgdensity.coordinates([xindex, yindex, zindex]),
            )
            self.closest_r_index[atomi][xindex][yindex][zindex] = numpy.abs(
                self.radial_grid_r[atomi] - dist_r
            ).argmin()

            # Equation 54 in doi: 10.1039/c6ra04656h
            stockholder_w[atomi][xindex][yindex][zindex] = self.proatom_density[atomi][
                self.closest_r_index[atomi][xindex][yindex][zindex]
            ]

        # Equation 55 in doi: 10.1039/c6ra04656h
        localized_w = numpy.power(stockholder_w, 4)

        # Equation 53 in doi: 10.1039/c6ra04656h
        stockholder_bigW = numpy.sum(stockholder_w, axis=0)
        localized_bigW = numpy.sum(localized_w, axis=0)

        refcharges = numpy.zeros((atoms))
        localizedcharges = numpy.zeros((atoms))
        stockholdercharges = numpy.zeros((atoms))

        for atomi in range(atoms):
            # Equation 52 and 51 in doi: 10.1039/c6ra04656h
            localizedcharges[atomi] = self.data.atomnos[atomi] - self.chgdensity.integrate(
                weights=(localized_w[atomi] / localized_bigW)
            )
            stockholdercharges[atomi] = self.data.atomnos[atomi] - self.chgdensity.integrate(
                weights=(stockholder_w[atomi] / stockholder_bigW)
            )

            # In DDEC6, weights of 1/3 and 2/3 are assigned for stockholder and localized charges.
            # (Equation 50 and 58 in doi: 10.1039/c6ra04656h)
            refcharges[atomi] = (stockholdercharges[atomi] / 3.0) + (
                localizedcharges[atomi] * 2.0 / 3.0
            )

        return refcharges, localizedcharges, stockholdercharges

    def condition_densities(self):
        """ Calculate conditioned densities
            [STEP 3]
        """
        # Generator object to iterate over the grid
        xshape, yshape, zshape = self.chgdensity.data.shape
        atoms = len(self.data.atomnos)
        indices = (
            (i, x, y, z)
            for i in range(atoms)
            for x in range(xshape)
            for y in range(yshape)
            for z in range(zshape)
        )

        self._rho_ref = numpy.zeros((xshape, yshape, zshape))

        for atomi, xindex, yindex, zindex in indices:
            # rho_ref -- Equation 41 in doi: 10.1039/c6ra04656h
            self._rho_ref[xindex][yindex][zindex] += self.proatom_density[atomi][
                self.closest_r_index[atomi][xindex][yindex][zindex]
            ]

        self._candidates_bigPhi = []
        self._candidates_phi = []

        # Initial conditions are detailed in Figure S1 in doi: 10.1039/c6ra04656h
        phiAI = numpy.zeros_like(self.data.atomnos, dtype=float)
        bigphiAI = numpy.zeros_like(self.data.atomnos, dtype=float)
        self._y_a = []
        self._cond_density = []

        for atomi in range(len(self.data.atomnos)):
            # y_a -- equation 40 in doi: 10.1039/c6ra04656h
            self._y_a.append(self._ya(self.proatom_density, atomi))
            # rho_A^cond -- equation S97 in doi: 10.1039/c6ra04656h
            self._cond_density.append(
                self._y_a[atomi] + bigphiAI[atomi] * numpy.sqrt(self._y_a[atomi])
            )
            # Monotonic Decrease Condition (as expressed in equation S99)
            self._cond_density[atomi] = numpy.minimum.accumulate(self._cond_density[atomi])
            # phi_A^I -- Equation S100 in doi: 10.1039/c6ra04656h
            phiAI[atomi] = (
                self._integrate_from_radial([self._cond_density[atomi]], [atomi])
                - self.data.atomnos[atomi]
                + self.refcharges[-1][atomi]
            )

            self._candidates_bigPhi.append([bigphiAI[atomi]])
            self._candidates_phi.append([phiAI[atomi]])

            fitphi = True  # when convergence is reached, this is modified to False.

            while phiAI[atomi] <= 0:
                # Iterative algorithm until convergence
                # Refer to S101 in doi: 10.1039/c6ra04656h
                bigphiAI[atomi] = 2 * bigphiAI[atomi] - phiAI[atomi] / self._integrate_from_radial(
                    [numpy.sqrt(self._y_a[atomi])], [atomi]
                )

                # When Phi is updated, related quantities are updated as well
                # Refer to S100 in doi: 10.1039/c6ra04656h
                phiAI[atomi], self._cond_density[atomi] = self._phiai(
                    self._y_a[atomi], bigphiAI[atomi], atomi
                )

                self._candidates_phi[atomi].append(phiAI[atomi])
                self._candidates_bigPhi[atomi].append(bigphiAI[atomi])

            # lowerbigPhi is largest negative Phi.
            # upperbigPhi is smallest positive Phi.
            if fitphi:
                self._candidates_phi[atomi] = numpy.array(self._candidates_phi[atomi], dtype=float)
                self._candidates_bigPhi[atomi] = numpy.array(
                    self._candidates_bigPhi[atomi], dtype=float
                )
                if numpy.count_nonzero(self._candidates_phi[atomi] < 0) > 0:
                    lower_ind = numpy.where(
                        self._candidates_phi[atomi]
                        == self._candidates_phi[atomi][self._candidates_phi[atomi] < 0].max()
                    )[0][0]
                    lowerbigPhi = self._candidates_bigPhi[atomi][lower_ind]
                    lowerphi = self._candidates_phi[atomi][lower_ind]
                else:  # assign some large negative number
                    lowerbigPhi = numpy.NINF
                    lowerphi = numpy.NINF
                if numpy.count_nonzero(self._candidates_phi[atomi] > 0) > 0:
                    upper_ind = numpy.where(
                        self._candidates_phi[atomi]
                        == self._candidates_phi[atomi][self._candidates_phi[atomi] > 0].min()
                    )[0][0]
                    upperbigPhi = self._candidates_bigPhi[atomi][upper_ind]
                    upperphi = self._candidates_phi[atomi][upper_ind]
                else:  # assign some large positive number
                    upperbigPhi = numpy.PINF
                    upperphi = numpy.PINF

            iter = 0
            while fitphi and iter < 50:
                # Flow diagram on Figure S1 in doi: 10.1039/c6ra04656h details the procedure.
                iter = iter + 1
                midbigPhi = (lowerbigPhi + upperbigPhi) / 2.0
                midphi = self._phiai(self._y_a[atomi], midbigPhi, atomi)[0]
                # Exit conditions
                if abs(lowerphi) < 1e-10:
                    bigphiAI[atomi] = lowerbigPhi
                    fitphi = False
                elif abs(upperphi) < 1e-10:
                    bigphiAI[atomi] = upperbigPhi
                    fitphi = False
                elif abs(midphi) < 1e-10:
                    bigphiAI[atomi] = midbigPhi
                    fitphi = False
                else:
                    # Parabolic fitting as described on Figure S1 in doi: 10.1039/c6ra04656h
                    # Type casting here converts from size 1 numpy.ndarray to float
                    xpts = numpy.array(
                        [float(lowerbigPhi), float(midbigPhi), float(upperbigPhi)], dtype=float
                    )
                    ypts = numpy.array(
                        [float(lowerphi), float(midphi), float(upperphi)], dtype=float
                    )
                    fit = numpy.polyfit(xpts, ypts, 2)
                    roots = numpy.roots(fit)  # max two roots (bigPhi)

                    belowphi = self._phiai(self._y_a[atomi], roots.min(), atomi)[0]
                    abovephi = self._phiai(self._y_a[atomi], roots.max(), atomi)[0]

                    if abs(abovephi) < 1e-10:
                        bigphiAI[atomi] = roots.min()
                        fitphi = False
                    elif abs(belowphi) < 1e-10:
                        bigphiAI[atomi] = roots.max()
                        fitphi = False
                    else:
                        if 3 * abs(abovephi) < abs(belowphi):
                            corbigPhi = roots.max() - 2.0 * abovephi * (
                                roots.max() - roots.min()
                            ) / (abovephi - belowphi)
                        elif 3 * abs(belowphi) < abs(abovephi):
                            corbigPhi = roots.min() - 2.0 * belowphi * (
                                roots.max() - roots.min()
                            ) / (abovephi - belowphi)
                        else:
                            corbigPhi = (roots.max() + roots.min()) / 2.0
                        # New candidates
                        corphi = self._phiai(self._y_a[atomi], corbigPhi, atomi)[0]
                        self._candidates_bigPhi[atomi] = numpy.array(
                            [
                                lowerbigPhi,
                                midbigPhi,
                                upperbigPhi,
                                roots.max(),
                                roots.min(),
                                corbigPhi,
                            ],
                            dtype=float,
                        )
                        self._candidates_phi[atomi] = numpy.array(
                            [lowerphi, midphi, upperphi, abovephi, belowphi, corphi], dtype=float
                        )

                        # Update upperphi and lowerphi
                        lower_ind = numpy.where(
                            self._candidates_phi[atomi]
                            == self._candidates_phi[atomi][self._candidates_phi[atomi] < 0].max()
                        )[0][0]
                        upper_ind = numpy.where(
                            self._candidates_phi[atomi]
                            == self._candidates_phi[atomi][self._candidates_phi[atomi] > 0].min()
                        )[0][0]

                        lowerphi = self._candidates_phi[atomi][lower_ind]
                        upperphi = self._candidates_phi[atomi][upper_ind]

                        if abs(lowerphi) < 1e-10:
                            bigphiAI[atomi] = self._candidates_bigPhi[atomi][lower_ind]
                            fitphi = False
                        elif abs(upperphi) < 1e-10:
                            bigphiAI[atomi] = self._candidates_bigPhi[atomi][upper_ind]
                            fitphi = False
                        else:
                            # Fitting needs to continue in this case.
                            lowerbigPhi = self._candidates_bigPhi[atomi][lower_ind]
                            lowerphi = self._candidates_phi[atomi][lower_ind]
                            upperbigPhi = self._candidates_bigPhi[atomi][upper_ind]
                            upperphi = self._candidates_phi[atomi][upper_ind]

            assert not fitphi, "Iterative conditioning failed to converge."

            # Set final conditioned density using chosen Phi
            self._cond_density[atomi] = self._phiai(self._y_a[atomi], bigphiAI[atomi], atomi)[1]

        self.logger.info("Calculating tau and combined conditioned densities.")

        # Calculate tau(r) and rho^cond(r)
        # Refer to equation 65 and 66 in doi: 10.1039/c6ra04656h
        # Assign rho^cond on grid using generator object
        self.rho_cond = copy.deepcopy(self.chgdensity)
        self.rho_cond.data = numpy.zeros_like(self.rho_cond.data, dtype=float)
        rho_cond_sqrt = numpy.zeros_like(self.rho_cond.data, dtype=float)

        # Generator object to iterate over the grid
        xshape, yshape, zshape = self.chgdensity.data.shape
        atoms = len(self.data.atomnos)
        indices = (
            (i, x, y, z)
            for i in range(atoms)
            for x in range(xshape)
            for y in range(yshape)
            for z in range(zshape)
        )

        self._leftterm = numpy.zeros((atoms, xshape, yshape, zshape), dtype=float)
        self.tau = []

        # rho_cond -- equation 65 in doi: 10.1039/c6ra04656h
        for atomi, xindex, yindex, zindex in indices:
            self.rho_cond.data[xindex][yindex][zindex] += self._cond_density[atomi][
                self.closest_r_index[atomi][xindex][yindex][zindex]
            ]

        rho_cond_sqrt = numpy.sqrt(self.rho_cond.data)

        for atomi in range(len(self.data.atomnos)):
            self.tau.append(numpy.zeros_like(self.proatom_density[atomi], dtype=float))
            grid = ((x, y, z) for x in range(xshape) for y in range(yshape) for z in range(zshape))
            for xindex, yindex, zindex in grid:
                # leftterm is the first spherical average term in equation 66.
                # <rho^cond_A(r_A) / sqrt(rho^cond(r))>
                self._leftterm[atomi][xindex][yindex][zindex] = self._cond_density[atomi][
                    self.closest_r_index[atomi][xindex][yindex][zindex]
                ]
            self._leftterm[atomi] = self._leftterm[atomi] / rho_cond_sqrt
            for radiusi in range(len(self.tau[atomi])):
                grid_filter = self.closest_r_index[atomi] == radiusi
                num_grid_filter = numpy.count_nonzero(grid_filter)
                if num_grid_filter < 1:
                    self.tau[atomi][radiusi] = 0.0
                else:
                    leftaverage = numpy.sum(grid_filter * self._leftterm[atomi]) / num_grid_filter
                    rightaverage = numpy.sum(grid_filter * rho_cond_sqrt) / num_grid_filter
                    if leftaverage < 1e-20:
                        self.tau[atomi][radiusi] = 0.0
                    else:
                        self.tau[atomi][radiusi] = numpy.divide(
                            leftaverage,
                            rightaverage,
                            out=numpy.zeros_like(leftaverage),
                            where=rightaverage != 0.0,
                        )
            # Make tau monotonic decreasing
            self.tau[atomi] = numpy.maximum.accumulate(self.tau[atomi][::-1])[::-1]

    def _phiai(self, ya, bigphiAI, atomi):
        """ Evaluate phi_A^I based on equation S100
        """
        # Re-evaluate cond_density
        if isinstance(bigphiAI, float) and not numpy.isinf(bigphiAI):
            cond_density = ya + bigphiAI * numpy.sqrt(ya)
        else:
            cond_density = ya

        # Monotonic Decrease Condition
        cond_density = numpy.minimum.accumulate(cond_density)

        # Re-evaluate phi_AI
        phiAI = (
            self._integrate_from_radial([cond_density], [atomi])
            - self.data.atomnos[atomi]
            + self.refcharges[-1][atomi]
        )

        return phiAI, cond_density

    def _ya(self, proatom_density, atomi):
        # Function that calculates Y_a^avg
        # See Eq. 40-41 in doi: 10.1039/c6ra04656h
        rho_ref = self._rho_ref
        # Y_a^avg -- Equation 40 in doi: 10.1039/c6ra04656h
        ya = numpy.zeros_like(proatom_density[atomi], dtype=float)
        weights = self.chgdensity.data / rho_ref
        for radiusi in range(len(ya)):
            grid_filter = self.closest_r_index[atomi] == radiusi
            num_grid_filter = numpy.count_nonzero(grid_filter)
            if num_grid_filter < 1:
                ya[radiusi] = 0.0
            else:
                spherical_avg = numpy.sum(grid_filter * weights) / num_grid_filter
                ya[radiusi] = proatom_density[atomi][radiusi] * spherical_avg

        # Make y_a monotonic decreasing
        # Refer to module_reshaping_functions.f08::77-79
        ya = numpy.maximum.accumulate(ya[::-1])[::-1]
        ya = numpy.minimum.accumulate(ya)

        # Normalize y_a (see module_DDEC6_valence_iterator.f08::284)
        nelec = self._integrate_from_radial([ya], [atomi])
        ya *= (self.data.atomnos[atomi] - self.refcharges[-1][atomi]) / nelec

        return ya

    def _integrate_from_radial(self, radial_density_list, atom_list):
        # Function that reads in list of radial densities, projects it on Cartesian grid,
        # and returns integrated value
        grid = copy.deepcopy(self.chgdensity)
        grid.data = numpy.zeros_like(grid.data)

        xshape, yshape, zshape = self.chgdensity.data.shape
        indices = ((x, y, z) for x in range(xshape) for y in range(yshape) for z in range(zshape))

        for density, atomi in zip(radial_density_list, atom_list):
            for x, y, z in indices:
                grid.data[x][y][z] = (
                    grid.data[x][y][z] + density[self.closest_r_index[atomi][x][y][z]]
                )

        return grid.integrate()