#!/usr/bin/env Python3
import numpy as np
from pdb import set_trace
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from matplotlib.colors import LogNorm
import time

'''
	Description:
	An attempt at creating a classic lyot coronograph.
	Not sure that it worked entirely as advertised... But 
	it does produce a reasonable result.
	A bright star is blocked, while a faint companion, orbiting
	on an elliptical orbit, is imaged as it goes through its 
	orbital phase.
'''

def keplerIII_period_to_semimajor_axis(orbital_period):
    #########################################################
    # Units: orbital period [yr], separation [au]           #
    # Function from 
    # https://github.com/katiebreivik/Keplers_Laws/blob/master/Graphing%20Kepler's%20Laws%20with%20Python.ipynb
    #########################################################
    semimajor_axis_cubed = orbital_period**2
    semimajor_axis = semimajor_axis_cubed**(1./3.)
    
    return semimajor_axis

def make_kepler_orbit(eccentricity,orbital_period, nStep=500):
    ##########################################################
    # Units: orbital period [years]                          #
    # returns: 500 true anomaly values throughout the orbit  #
    # Function from 
    # https://github.com/katiebreivik/Keplers_Laws/blob/master/Graphing%20Kepler's%20Laws%20with%20Python.ipynb
    ##########################################################
    tRange=np.linspace(0.0,orbital_period,nStep)
    theta=[]
    for time in tRange:
        PsiDiff=1.0
        M=2*np.pi*time/orbital_period
        PsiOld=M
        theta0old=180.0
        while PsiDiff>1e-10:
            PsiNew=M+eccentricity*np.sin(PsiOld)
            PsiDiff=PsiNew-PsiOld
            PsiOld=PsiNew
        theta0=2*np.arctan(((1+eccentricity)/(1-eccentricity))**(0.5)*np.tan(PsiOld/2.))
        theta.append(theta0)
    return theta

def orbit(semimajor_axis,eccentricity,true_anomaly): 
    ##############################################
    # Units: separation [au] #
    # Function from 
    # https://github.com/katiebreivik/Keplers_Laws/blob/master/Graphing%20Kepler's%20Laws%20with%20Python.ipynb
    ##############################################
    
    # define the shape equation
    r_orbit = semimajor_axis*(1 - eccentricity**2)/(1 + eccentricity*np.cos(true_anomaly))
    x_orbit = r_orbit*np.cos(true_anomaly)
    y_orbit = r_orbit*np.sin(true_anomaly)   
    return x_orbit,y_orbit


if __name__ == '__main__':
	plt.rc('font', family='serif')
	plt.rc('axes', labelsize='large')

	# Start time counter
	start = time.time()

	# Pick a wavelength
	l = 30e-3

	# Initiate Fourier Optics
	fou = FourierOptics(l=l)
	
	# Pick sampling parameters
	dpix = 1/50
	n, m = 3001, 3001
	print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix))
	
	# Create an input grid
	x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m))
	
	# Initiate the profiles class and pick a profile (tophat, obscured tophat or spider)
	prof = Profiles()

	###############
	# from https://github.com/katiebreivik/Keplers_Laws/blob/master/Graphing%20Kepler's%20Laws%20with%20Python.ipynb
	orbital_period = 1.0
	eccentricity = 0.23
	list1 = make_kepler_orbit(eccentricity,orbital_period, 360)
	semimajor_axis = keplerIII_period_to_semimajor_axis(orbital_period)
	xOrbit, yOrbit = orbit(semimajor_axis,eccentricity,list1)
	###############

	xOrbit *= 8
	yOrbit *= 8

	# Create lens system
	L1 = LensObject('N-BK7', -50, 1e99, 1, 15, name='L1')
	L2 = LensObject('N-BK7',  50, 1e99, 1, 15, name='L2')
	L3 = LensObject('N-BK7',  50, 1e99, 1, 15, name='L3')
	L4 = LensObject('N-BK7',  50, 1e99, 1, 15, name='L4')

	# Calculate corresponding phase delays
	t_l1 = L1.phase_delay(x,y,l, paraxial=True)
	t_l2 = L2.phase_delay(x,y,l, paraxial=True)
	t_l3 = L3.phase_delay(x,y,l, paraxial=True)
	t_l4 = L4.phase_delay(x,y,l, paraxial=True)

	# Verify that the focal lengths are correct
	print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f))
	print('Focal length of lens {} is {:.1f} mm.'.format(L2.name, L2.f))
	print('Focal length of lens {} is {:.1f} mm.'.format(L3.name, L3.f))
	print('Focal length of lens {} is {:.1f} mm.'.format(L4.name, L4.f))

	# Other parameters
	oC 			= [x.min(), x.max(), y.min(), y.max()]
	savename 	= False

	# The coronographic mask
	coroMask 	= np.ones((n,m)) - prof.top_hat(2, n, m, dpix=dpix, dx=0, dy=0)
	coroMask = np.ones((n,m))

	# The Lyot stop
	lyotStop	= prof.obscured_top_hat(15, 3, n, m, dpix=dpix)
	lyotStop = np.ones((n,m))

	for i in range(len(list1)):
		title = f"{i+1:03d}"
		print (' Calculating {0}/{1}'.format(title, f"{len(list1):03d}"), end="\r")
		intenseSource = prof.top_hat(1, n, m, dpix=dpix, oversample=1)
		weakCompanion = prof.top_hat(0.1, n, m, dpix=dpix, oversample=1, dx=xOrbit[i], dy=yOrbit[i])
		A = intenseSource + 0.01*weakCompanion
		A[A == 0] = 1e-99
		# plt.figure()
		# plt.imshow(np.abs(A)**2, cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()], norm=LogNorm(vmin=1e-6, vmax=1))
		# plt.xlim(-20, 20)
		# plt.ylim(-20, 20)
		# plt.colorbar()
		# plt.savefig('./figures/high_contrast_imaging/input_{:02d}.png'.format(i+1), bbox_inches='tight', transparent=True)
		# plt.close()

		# The input distribution gets a curved wavefront after the lens
		U_in1 		= A * t_l1

		# Performing the propagation using Fresnel approximation to the focal plane.
		Exyz,_,_ 	= fou.fresnel_propagation(U_in1, x, y, dz=-L1.f, show=False, save=savename, outputCoordinates=oC)

		# Focussing the beam again
		U_in2 		= Exyz * t_l2 * t_l2
		Exyz2,_,_ 	= fou.fresnel_propagation(U_in2, x, y, dz=L2.f, show=False, save=savename, outputCoordinates=oC)

		# Now we add a coronographic mask, following a gaussian
		U_in3 		= Exyz2.real * coroMask + Exyz2.imag * coroMask * 1j

		# Let this propagate again
		Exyz3,_,_ 	= fou.fresnel_propagation(U_in3, x, y, dz=L2.f, show=False, save=savename, outputCoordinates=oC)

		# Collimating the beam
		U_in4 		= Exyz3 * t_l3

		# Addition of a lyot stop
		U_in4 		= U_in4.real * lyotStop + U_in4.imag * lyotStop * 1j

		# Making an image
		U_in5 		= U_in4 * t_l4
		Exyz5,_,_ 	= fou.fresnel_propagation(U_in5, x, y, dz=L4.f, show=False, save=savename, outputCoordinates=oC)	
		# imageMask 	= np.ones((n,m)) - prof.top_hat(5.7, n, m, dpix=dpix, dx=0, dy=0)
		# imageMask[ imageMask == 0] = 1e-99
		# imageData = np.abs(Exyz5.real * imageMask + Exyz5.imag * imageMask * 1j)**2 
		imageData = np.abs(Exyz5)**2

		plt.figure()
		plt.imshow(imageData, cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()], norm=LogNorm(vmin=1e-11))
		plt.xlim(-20, 20)
		plt.ylim(-20, 20)
		plt.colorbar()
		plt.savefig('./figures/high_contrast_imaging/no_mask_image_{:03d}.png'.format(i+1), bbox_inches='tight', transparent=True)
		# plt.show()
		plt.close()
		

	# End time counter
	end = time.time()

	print('The program took {:.3f} seconds!'.format((end-start)))