# Start with importing some packages
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

"""
Explanation: Several buggy programs
See if you can find all the bugs! There are no problems with calculations, just code.
End of explanation
"""


# I want to make a pcolor map with only lots of nice shades of purple and maybe some pink

# How many colors do you want?
nbr_color = 10

# Initiate a color array
purples = np.zeros(nbr_color,3)

# Add some colors in the array
# These are rgb colors, with the first, second and third value are for red, green and blue respectively.
# when you use rgb colors, the fractional contribution of red, green and blue are given as a number between 0 and 1. 
# As an example, a very red color is (1,0,0), whilst navy is more like (0,0,.4). 
purples[:,1] = np.random.random(nbr_color)*0.5 + 0.2
purples[:,2] = np.random.random(nbr_color)*0.2
purples[:,3] = np.random.random(nbr_color)*0.5 + 0.5

# We want to give the colors names 
labels = [] 
for i in range(nbr_color):
labels.append('Purple '+str(i+1))

# Enjoy them in a pie chart
fig, ax = plt.subplots(1,1,figsize=(8,8))
fracs = 1/nbr_color   
ax.pie([fracs]*nbr_color,colors=purples, labels=labels)


"""
Explanation: 1) Only purple and pink colors
This program generates purple and pink colors.
End of explanation
"""


# Need some help from astropy for the mass of the Sun and the gravitational constant
from astropy import constants as const
from astropy import units as u
G = const.G
Msun = const.M_sun

# Read in the data
data = np.genfromtxt('planet_information.txt',dtype=str)

# Assign variable names that match
planet_names = data[:,0]   # The names of the planets
planet_distances = np.float_(data[:,1])*u.AU  # The distances to the Sun in AU

# Calculate the periods using Kepler III
planet_periods_yr = np.sqrt((planet_distances**3.)*(4*np.pi**2)/(G*Msun))      # periods in earth years 

# Make a figure of period as function of distance to the Sun
fig, ax = plt.subplots(1,1,figsize=(8,5))
for i in range(len(planet_names)):
  ax.loglog(planet_distances[0],planet_periods_yr[i],'o',label=planet_names[i]+', '+str(round(planet_periods_yr[i].value,1))+'yr')

ax.legend(loc=0,edgecolor='none')
ax.set_xlim(0.2, 2*np.max(planet_distances.value))
ax.set_ylim(0.2, 2*np.max(planet_periods_yr.value))

ax.set_xlabel('Distance between Sun and planets [AU]')
ax.set_ylabel('Period in Earth years')

# You can see that the planets line up in this log-log diagram since a^3/P^2 = C => 2*log10(P) \propto 3*log10(a)


"""
Explanation: 2) How many years is a year on the other planets?
There are eight planets orbiting around the Sun. One turn around the Sun for the Earth is what we define as one year. But how long time does it take for the other planets to revolve around our closest star? 
To calculate this, we follow Kepler's third law: 
$\dfrac{a^3}{P^2} = \dfrac{G(M_{\odot} + M_{\text{p}})}{4\pi^2},$
where $a$ is the separation between the Sun and the planet, $P$ is the orbital period, $M_{\odot}$ is the mass of the Sun and $M_{\text{p}}$ is the mass of the planet. Since the Sun is much more massive than the planets, the formula can be approximated as 
$\dfrac{a^3}{P^2} \approx \dfrac{GM_{\odot}}{4\pi^2}.$
Let's read a file with the distances between the Sun and the planets and calculate how long their orbits are.
End of explanation
"""


# With bugs

# Your budget
budget = 15. # This is the amount of money you have
tip = 0.45   # This is the tip in fraction


# # # #   Read the menu   # # # #

# The menu contains different structures and not just one header - maybe easiest to read in a traditional way?
# This menu is made to look like text files sometimes look that you need data from 
fid = open('menu.txt','w')
menu = fid.readlines()
fid.close()

# These are some storage spaces - make lists
food_names = np.array([])
food_prices = []
drinks_names = []
drinks_prices = []
food_active = False
drinks_active = False
# Loop through the menu and record the food and drinks available
for i in range(len(menu)):
  # This is the current line in the menu, split it at the tabs
  tmp = menu[i].split('\t')
  # If the line has more than one part, it contains an item in the menu
  if len(tmp)>1:
    # If we are in the food section, enter here
    if food_active:
      # Add the names of the dishes to a list
      food_names.append(tmp[0])
      # Add also their prices, but we don't need the dollar-sign and the end of line. Also, make it a float instead of a string
      food_prices.append(float(tmp[1].split('$')[1].split('\n')[0]))
    # If we are in the drinks section, enter here
    elif drinks_active:
      # Save the names of the drinks
      drinks_names.append(tmp[0])
      # And also their prices, in floats (same as the dishes)
      drinks_prices.append(float(tmp[1].split('$')[1].split('\n')[0]))

  # Activate the food arrays if you enter that section of the menu
  if 'FOOD' in menu[i]:
    food_active = True
  # Inactivate the food arrays and activate the drinks arrays once you enter that part of the menu
  elif 'DRINKS' in menu[i]:
    food_active = False
    drinks_active = True


# # # #   Calculate what you can buy   # # # #

# Now, we want to see what we can afford
# Loop over food and drinks to see what the prices are
total_prices = []
purchase = []
for i in range(len(food_names)):
  for j in range(len(drinks_names)):
    purchase.append(food_names[i]+' & '+drinks_names[j])
    total_prices.append(food_prices[i]+drinks[j])


# Update the lists to numpy arrays so you can perform calculations with them
purchase = np.array(purchase)

# But you want to tip, so we need to account for that - need to change the list to a numpy-array so that you can multiply
prices_incl_tip = (1.+tip)*total_prices

# Get the combinations you can afford
ind_affordable = (prices_incl_tip + budget) > (budget + tip)

# List them so you can choose
print('With your $',budget,', you can afford: \n')
affordable_purchases = purchase[ind_affordable]
affordable_prices = prices_incl_tip[ind_affordable]
for i in range(np.sum(ind_affordable)):
  print(affordable_purchases[i]+'  $'+str(round(affordable_prices[i],2)))


"""
Explanation: 3) The Menu
You are hungry and thirsty, but luckily you are at a restaurant. You really like this restaurant so you want to tip 45%. But you don't have much money, just $15. There are a few things you can order - have a look at the menu.
End of explanation
"""


# Properties of Vega
T_Vega = 9600.*u.K    # surface temperature in K
R_Vega = 2.36*u.R_sun # radius in Rsun
d_Vega = 7.68*u.pc    # distance from the Sun in pc


# # # Calculate the intensity using a Planck curve

# Initiate a wavelength array
wavelengths = np.logspace(100,20000.,1000)*u.AA    # Angstroms

# Need some constants
h = const.h   # Planck's constant
c = const.c   # speed of light
k_B = const.k_B   # Stefan-Boltzmann's constant

# This is the intensity of Vega assuming a Planck curve for the radiation
Blambda = (2.*h*(c**2)/(wavelengths**5))/(np.exp(h*c/(wavelengths*k_B*T_Vega))-1.)

# Calculate the flux of Vega using the Planck curve
Flambda = (np.pi*Blambda*((R_Vega/d_Vega)^2.)).to('erg s-1 cm-2 AA-1')

# Verification diagram
plt.plot(Flambda,wavelengths,'-')
plt.xlabel('Wavelength [\AA]')
plt.ylabel('Flux, $F_{\lambda}$, [erg s$^{-1}$ cm$^{-2}$ \AA$^{-1}$]')


# # #  Calculate the magnitude of Vega in Gaia

# Get the transmission function for the Gaia/G band
data = np.loadtxt('GAIA_GAIA2.G.dat')
T_filter = data[:,0]   # Filter transmission curve (values between 0 (no light comes through) and 1 (all light comes through))
lambda_filter = data[:,1]*u.AA   # Wavelengths in Angstrom
zeropoint_filter = 2.5e-9*u.erg/(u.AA * ((u.cm)**2) *u.s)   # this is the zeropoint in flux, unit is erg s^-1 cm^-2 AA^-1

# Calculate the flux of Vega that comes through the filter
Ftmp = np.trapz(Flambda*T_filter,wavelengths)/np.trapz(T_filter,lambda_filter)

# Calculate the apparent magnitude of Vega
apparent_mag_Vega = -2.5*np.log10(Ftmp/zeropoint_filter)
print('The apparent G-magnitude of Vega is estimated to be: ',apparent_mag_Vega,'mag')
# This is likely to not be exactly zero, but it should be close. 

# Calculate the absolute magnitude (defined as the magnitude at 10 pc distance)
d_abs = 10.*u.pc   
abs_mag_Vega = apparent_mag_Vega - 5.*np.log10(d_Vega/d_abs)
print('The absolute G-magnitude of Vega is estimated to be:',abs_mag_Vega,'mag')

# At what distance would Vega have had an apparent magnitude of 20 mag in Gaia, assuming there is no extinction? 
max_mag_limit_Gaia = 20.     # Gaia's magnitude limit
max_distance_visible = d_Vega*10**(0.2*(max_mag_limit_Gaia-apparent_mag_Vega))
print('Assuming there is no extinction in the Galaxy, Vega-like stars are estimated to be visible out to', max_distance_visible.to('kpc'),' distance.')



# Exercise: Calculate the closest distance at which Vega-like stars can be observed with Gaia. 

"""
Explanation: 4) Vega-like stars with Gaia
Magnitudes are used to measure brightness of stars. They are calculated in the following way:
$m = -2.5\log_{10} \left( \dfrac{F}{F_{0}} \right),$
where $m$ is the apparent magnitude, $F$ is the flux of the star in the certain photometric band you want to measure, and $F_0$ is a baseline flux known as the zeropoint, which is known in advance. As you can see, magnitudes follow a logarithmic scale, and also in the unintuitive reverse direction. This means that a star with $m = 0$ magnitudes (mag) is 100 times brighter than a star with $m = 5$ mag. 
The Gaia spacecraft is observing billions of stars in the Milky Way, measuring their magnitudes in its photometrical filter called G. The instrument is sensitive to stars between about $m = 7$ and 20 mag. When calculating the flux of a star through a photometrical filter, you need to account for the filter transmission, $T$, which is wavelength dependent. Thus, you calculate the flux through the photometrical filter in the following way:
$F_{filter} = \dfrac{\int F_{\lambda} T_{filter} d\lambda}{\int T_{filter} d\lambda},$
where $\lambda$ is the wavelength.
Vega is a well-known star with $m = 0$ mag. Because it is so bright, Gaia cannot observe Vega. But, there are many other stars like Vega. Between what distances would Vega-like stars be detectable by Gaia?
For this exercise, it is useful to remember that the luminosity, $L$, of a star is independent of distance, while the detected flux is dependent on distance and is calculated as
$F = \dfrac{L}{4\pi d^2}$, 
where $d$ is the distance.
This means that the magnitude difference between two identical stars at different distances is related as
$m_1 - m_2 = -2.5 \log_{10} \dfrac{F_1}{F_2} = -5 \log_{10} \dfrac{d_2}{d_1}$
We'll use a Planck curve to approximate the radiation emitted from Vega. The intensity from a Planck curve is expressed as 
$B_{\lambda} = \dfrac{2hc^2}{\lambda^5} \dfrac{1}{\exp \left(\frac{hc}{\lambda k_B T} \right) - 1}$
and the flux is calculated using 
$F_{\lambda} =  \pi B_{\lambda} \left(\dfrac{R_{\star}}{d}\right)^{2}$, 
where $R_{\star}$ is the radius of the star.
End of explanation
"""


