import numpy as np
import astropy.modeling.blackbody as bb 
import astropy.constants as const
from astropy.io import fits
from scipy.interpolate import interp2d

class FaiglerMazehFit():

	def __init__(self, P_orb, inc, R_star, M_star, T_star, A_ellip=False, A_beam=False,
				 R_p=False, a=False, u=False, g=0.65, logg=None, tele='TESS', M_p=False,
				 K=False):
		self.P_orb = P_orb 	   		 # orbital period in days
		self.inc = inc * np.pi / 180 # inclination converted to radians
		self.R_star = R_star  		 # radius of the star in solar units
		self.M_star = M_star  		 # mass of the star in solar units
		self.T_star = T_star		 # temperature of the star [K]
		self.A_ellip = A_ellip 		 # ellipsoidal amplitude in ppm
		self.A_beam = A_beam 		 # beaming amplitude in ppm

		self.g = g			  		 # gravity-darkening coefficient, expected range is 0.3-1.0 
		self.logg = logg 			 # log surface gravity of the star [cm s^-2]
		self.tele = tele.lower() 	 # observation instrument used, default is TESS. Only other
									 # other option (for now) is Kepler.
		self.R_p = R_p 				 # radius of the planet in jupiter radii
		self.a = a
		self.u = u 					 # the limb-darkening coefficient, range is 0-1
		self.g = g
		self.M_p = M_p
		self.K = K 
	
		# get the mass from the ellipsoidal amplitude, if given.
		# u is the limb-darkening coefficient, range is 0-1
		if not M_p and not not A_ellip and not not logg:
			self.u = self.LDC()
			self.M_p = self.m_from_ellip()

		# star-planet separation [au] assuming a circular orbit
		if not a and not not M_p:
			self.a = get_a(self.P_orb * 86400, self.M_star * const.M_sun.value, \
						   self.M_p * const.M_jup.value) / const.au.value


	def alpha_ellip(self):
		if not self.u:
			self.u = self.LDC()
		if not self.g:
			self.g = self.GDC()

		a = 15 + self.u
		b = 1 + self.g
		c = 3 - self.u
		return 0.15 * a * b / c

	def RV_amp(self):
		"""
		Returns the radial velocity amplitude [m/s] of the star given a companion mass.
		"""
		return 27 / 40 * const.c.value \
			   * self.M_star ** (-2/3) \
			   * self.P_orb ** (-1/3) \
			   * self.M_p * np.sin(self.inc) 

	def doppler_shift(self, K):
		"""
		Returns the shift in wavelength for a given radial velocity amplitude.
		"""
		return K / const.c.value

	def response_convolution(self, lambdas, response):		
		return response * bb.blackbody_lambda(lambdas, self.T_star).value

	def alpha_beam(self, K):
		""" 
		Returns the factor that accounts for the flux lost when a star gets Doppler shifted
		in and out of the observer's bandpass. 
		"""
		print(K)
		rest_lambdas, response = response_func(self.tele)

		flux_rest = np.trapz(self.response_convolution(rest_lambdas, response), \
							 x=rest_lambdas)

		blueshifted_lambdas = rest_lambdas - self.doppler_shift(K=K)
		flux_blueshift = np.trapz(self.response_convolution(blueshifted_lambdas, response), \
								  x=rest_lambdas)

		redshifted_lambdas = rest_lambdas + self.doppler_shift(K=K)
		flux_redshift = np.trapz(self.response_convolution(redshifted_lambdas, response), \
								 x=rest_lambdas)

		alpha_blue = abs( (flux_rest - flux_blueshift) / flux_rest )
		alpha_red = abs( (flux_rest - flux_redshift) / flux_rest )

		return 1 - np.mean([alpha_red, alpha_blue])


	def m_from_ellip(self):
		return self.A_ellip \
			   * self.R_star ** (-3) \
			   * self.M_star ** 2 \
			   * self.P_orb ** 2 \
			   / (12.8 * self.alpha_ellip() * np.sin(self.inc) ** 2)

	def ellip_from_m(self):
		return self.M_p * 12.8 * self.alpha_ellip() * np.sin(self.inc) ** 2 \
			   * self.R_star ** 3 \
			   * self.M_star ** (-2) \
			   * self.P_orb ** (-2)

	def m_from_beam(self, K=False, alpha_beam=False):
		if not alpha_beam and not K and not not self.M_p:
			alpha_beam = self.alpha_beam(K=self.RV_amp())
		elif not alpha_beam and not not K:
			alpha_beam = self.alpha_beam(K=K)
		elif not not K and not not alpha_beam:
			raise ValueError("Please only specify either K or alpha_beam, not both.")
		elif not K and not alpha_beam:
			raise ValueError("Please specify a radial velocity (K) or alpha_beam parameter")
	
		return self.A_beam \
			   * self.M_star ** (2/3) \
			   * self.P_orb ** (1/3) \
			   / (alpha_beam * np.sin(self.inc) * 2.7)


	def beam_from_m(self):
		"""
		Returns the expected Doppler beaming amplitude [ppm] for a given mass.
		"""
		if not self.M_p:
			raise ValueError("Argument 'M_p' must be specified if you're trying to " +  
							 "derive a beaming amplitude from a mass.")

		if not self.K:
			K=self.RV_amp()
				  
		return 2.7 * self.alpha_beam(K=self.K) \
			   * self.M_star ** (-2/3) \
			   * self.P_orb ** (-1/3) \
			   * self.M_p * np.sin(self.inc)

	def Ag_from_thermref(self, A_thermref):
		"""
		Return the geometric albedo derived from the thermal + ref amplitude.
		"""
		return  A_thermref * (self.R_p / self.a) ** -2 * (const.au / const.R_jup) ** 2

	def mass(self, derived_from=None, K=False, alpha_beam=False):

		if derived_from == "ellip":
			return self.m_from_ellip()

		elif derived_from == "beam":
			return self.m_from_beam(K=K, alpha_beam=alpha_beam)

		else:
			raise ValueError("derived_from must equal either 'ellip' or 'beam'")

	def nearest_neighbors(self, value, array, max_difference):
		""" 
		Returns a set of nearest neighbor indices of the given array.
		"""
		return set(list((np.where(abs(array - value) < max_difference))[0]))

	def correct_maxdiff(self, value, array, guess):
		while len(self.nearest_neighbors(value, array, guess)) > 0:
			guess -= 0.01 * guess
		return guess

	def shared_neighbor(self, value1, array1, max_diff1, value2, array2, max_diff2):
		set1 = self.nearest_neighbors(value1, array1, max_diff1)
		set2 = self.nearest_neighbors(value2, array2, max_diff2)
		nearest = list(set1.intersection(set2))

		# if len(nearest) > 1:
		# newmax_diff1 = self.correct_maxdiff(value1, array1, max_diff1)
		# newmax_diff2 = self.correct_maxdiff(value2, array2, max_diff2)
		# print(newmax_diff1, newmax_diff2)
		# if newmax_diff2 > newmax_diff1:
		# 	max_diff2 = newmax_diff2
		# else:
		# 	max_diff1 = newmax_diff1

		# set1 = self.nearest_neighbors(value1, array1, max_diff1)
		# set2 = self.nearest_neighbors(value2, array2, max_diff2)
		# nearest = list(set1.intersection(set2))
		# print(nearest)

		# # if len(nearest) > 1:

		# # 	raise ValueError("Multiple shared nearest neighbors, indices = ", nearest)
		# # else:
		# # 	return nearest[0]

		return nearest[0]


	def tess_warning(self):
		if self.tele != 'tess':
			raise ValueError("This function is only appropriate for observations done with " +
							 "the TESS satellite")

	def claret_LDC(self):
		"""
		Returns the mu coefficient and the four-parameters used in the Claret four-parameter 
		limb-darkening law (Claret 2000). These are obtained by finding the nearest neighbor
		in the model limb-darkening of TESS from Claret 2018. 
		"""
		# print("claret_LDC is still garbage, sorry. Quitting now...")
		# exit()
		self.tess_warning()

		logg, Teff, a1, a2, a3, a4, mu, mod = np.genfromtxt('../claret_ldc.dat', 
														usecols=(0,1,4,5,6,7,8,10),
														unpack=True)
		mod = np.genfromtxt('../claret_ldc.dat', usecols=(10,), dtype='str')

		if self.T_star <= 3000:
			# the PC model is meant for cool stars, and if we break it up this way we can do an 
			# easier 2D interpolation.
			mask = mod == 'PD'
		else:
			mask = mod == 'PC'
		logg = logg[mask]
		Teff = Teff[mask]
		a1 = a1[mask]
		a2 = a2[mask]
		a3 = a3[mask]
		a4 = a4[mask]
		mu = mu[mask]

		nearest = self.shared_neighbor(self.T_star, Teff, 100, self.logg, logg, 0.25)

		mu = mu[nearest]
		a_coeffs = [a1[nearest], a2[nearest], a3[nearest], a4[nearest]]

		return mu, a_coeffs

	def GDC(self):
		"""
		Returns the gravity-darkening coefficient from the Claret 2017 model
		"""
		self.tess_warning()

		logg, log_Teff, g = np.genfromtxt('../claret_gdc.dat', usecols=(2,3,4), unpack=True)
		
		nearest = self.shared_neighbor(np.log10(self.T_star), log_Teff, .01, self.logg, 
					   				   logg, 0.25)
		return g[nearest]

	def LDC(self):
		"""
		Returns the limb-darkening coefficient of the host star.
		"""
		mu, a_coeffs = self.claret_LDC()
		return 1 - sum([a_coeffs[k] * (1 - mu ** ((k+1) / 2)) for k in range(4)])

def get_response_specs(tele):
	if tele=="tess":
		return "../tess-response-function-v1.0.csv", ',', 1e1
	elif tele=="kepler":
		return "../kepler_hires.dat", '\t', 1e4

def response_func(tele):
	file, delimiter, to_AA = get_response_specs(tele)
	lambdas, response = np.genfromtxt(file, delimiter=delimiter, usecols=(0,1), unpack=True)
	return lambdas * to_AA, response

def get_a(P, M_star, M_p):
	"""
	Use Kepler's third law to derive the star-planet separation.
	"""
	return (P ** 2 * const.G.value * (M_star + M_p) / (4 * np.pi ** 2)) ** (1/3)

[((266, 21, 266, 89), 'numpy.genfromtxt', 'np.genfromtxt', (), '', True, 'import numpy as np\n'), ((213, 40, 215, 26), 'numpy.genfromtxt', 'np.genfromtxt', (), '', True, 'import numpy as np\n'), ((216, 8, 216, 70), 'numpy.genfromtxt', 'np.genfromtxt', (), '', True, 'import numpy as np\n'), ((245, 22, 245, 86), 'numpy.genfromtxt', 'np.genfromtxt', (), '', True, 'import numpy as np\n'), ((61, 19, 61, 35), 'numpy.sin', 'np.sin', ({(61, 26, 61, 34): 'self.inc'}, {}), '(self.inc)', True, 'import numpy as np\n'), ((94, 13, 94, 45), 'numpy.mean', 'np.mean', ({(94, 21, 94, 44): '[alpha_red, alpha_blue]'}, {}), '([alpha_red, alpha_blue])', True, 'import numpy as np\n'), ((140, 19, 140, 35), 'numpy.sin', 'np.sin', ({(140, 26, 140, 34): 'self.inc'}, {}), '(self.inc)', True, 'import numpy as np\n'), ((247, 33, 247, 54), 'numpy.log10', 'np.log10', ({(247, 42, 247, 53): 'self.T_star'}, {}), '(self.T_star)', True, 'import numpy as np\n'), ((70, 20, 70, 61), 'astropy.modeling.blackbody.blackbody_lambda', 'bb.blackbody_lambda', ({(70, 40, 70, 47): 'lambdas', (70, 49, 70, 60): 'self.T_star'}, {}), '(lambdas, self.T_star)', True, 'import astropy.modeling.blackbody as bb\n'), ((102, 37, 102, 53), 'numpy.sin', 'np.sin', ({(102, 44, 102, 52): 'self.inc'}, {}), '(self.inc)', True, 'import numpy as np\n'), ((123, 22, 123, 38), 'numpy.sin', 'np.sin', ({(123, 29, 123, 37): 'self.inc'}, {}), '(self.inc)', True, 'import numpy as np\n'), ((105, 48, 105, 64), 'numpy.sin', 'np.sin', ({(105, 55, 105, 63): 'self.inc'}, {}), '(self.inc)', True, 'import numpy as np\n')]