%matplotlib inline
%config InlineBackend.figure_format = 'retina'

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import datetime
from astropy import coordinates as coord, units as u, time, constants as const
import logging
import sys
import os
from scipy import interpolate, stats
import ephem
import multiprocessing as mp

logger = logging.getLogger('window-function')
logger.setLevel(logging.DEBUG)

np.random.seed(42)

colours = sns.color_palette(n_colors=5)

sns.set(style='white')

"""
Explanation: Can we simulate the period window for a typical planet that NGTS is interested in and compute for each period, what is the sensitivity to said period averaged over all phases and planet periods.
Process

Pick an observing strategy (every day, or every two days for two years)
Define a weather fraction
From the start point of the simulation, for two years remove days based on weather fraction
Pick a planet class (large neptune, or small neptune) - for now, until we draw random classes
Pick a random period from 1 to 30 days
Pick a random starting epoch, after the start point of the simulation
Pick a random ra, and declination between 36 and -84
For each day that is not knocked out by weather:
find out the beginning/end of astronomical night
decide if a transit is contained within this range,
and that it's visible above airmass 2 (altitude 30)
Compute the number of transits detected divided by the target number for belief (probably 3)
For each period, sum the fractions to get the period window
End of explanation
"""


def datetime_to_jd(dt):
    datetime_object = datetime.datetime(
        dt.year, dt.month, dt.day, 0, 0, 1)
    return time.Time(datetime_object.isoformat(), format='isot', scale='utc').jd

os.environ['NPLANETS'] = '100'
try:
    n_planets_to_draw = int(os.environ['NPLANETS'])
except KeyError:
    raise RuntimeError("Set NPLANETS environment variable to choose the number of planets")
    

print('Computing for %d planets' % n_planets_to_draw)
weather_fraction = 0.7
ndraws = int(os.environ.get('NDRAWS', 10))
n_requested_transits = 3.
start_date, end_date = datetime.date(2015, 1, 1), datetime.date(2018, 1, 1)
ndays = (end_date - start_date).days
dates = np.asarray([start_date + datetime.timedelta(days=i) for i in range(ndays)])
site_coordinates = coord.EarthLocation.from_geodetic(70.4042 * u.degree, -24.6272 * u.degree, 2400 * u.m)
ra_choices = np.linspace(175.5, 184.5, 24) * u.degree
dec_choices = np.linspace(-56, -44, 24) * u.degree
start_jd, end_jd = datetime_to_jd(start_date), datetime_to_jd(end_date)
half_an_hour = (0.5 * u.hour).to(u.d)

def altitude(position):
    return site_coordinates.itrs.separation(position)

"""
Explanation: First we'll start with the simple observing as often as we can when a field is up. We only need to consider one year in this case. Start with some constants:
End of explanation
"""


ind = np.ones(ndays, dtype=bool)
days = np.arange(ndays)
bad_days = np.random.choice(days, replace=False, size=(1 - weather_fraction) * ndays)
logger.info('Removing %d bad days', bad_days.size)
ind[bad_days] = False
valid_dates = set(dates[ind])

"""
Explanation: Now we can remove the days in which the weather is bad.
End of explanation
"""


class DrawPeriods(object):
    
    def __init__(self, values):
        self.periods = np.asarray([2, 3.4, 5.9, 10, 17])
        self.values = values
        self.interp = interpolate.interp1d(self.values, self.periods)
        
    def draw(self):
        chosen_value = np.random.uniform(self.values.min(), self.values.max())
        return float(self.interp(chosen_value))
    
    @property
    def min_period(self):
        return self.periods.min()
    
    @property
    def max_period(self):
        return self.periods.max()
    
    @classmethod
    def large_neptunes(cls):
        values = np.asarray([0.004, 0.01, 0.12, 0.21, 0.49]) / 100.
        return cls(values)
    
    @classmethod
    def small_neptunes(cls):
        values = np.asarray([0.035, 0.22, 0.95, 2.88, 6.55]) / 100.
        return cls(values)

"""
Explanation: Now a class to draw periods, as the setup is relatively expensive.
End of explanation
"""


d = DrawPeriods.large_neptunes()
periods = [d.draw() for _ in range(1000)]
_, _, lines = plt.hist(periods, bins=100, normed=True, cumulative=True, histtype='step', lw=2., color=colours[0])
plt.ylim(0, 1)
plt.twinx()
prob_dist = plt.plot(d.periods, d.values / d.values.max(), color=colours[1])
plt.ylim(0, 1)
plt.legend([lines[0], prob_dist[0]], ['Planet distribution', 'Probability distribution'], loc='best')
None

"""
Explanation: As a test, we shall see what the cumulative distribution of periods looks like.
End of explanation
"""


def separation_from_period(period):
    period = (period * u.d).to(u.s)
    mass = (1. * u.astrophys.M_sun).to(u.kg)
    norm = ((const.G * mass / (4. * np.pi ** 2))) ** (1. / 3.)
    return (norm * period ** (2. / 3.)).to(u.m)

def transit_probability(period):
    radius = 1. * u.astrophys.R_sun.to(u.m)
    return (radius / separation_from_period(period)).value

def compute_twighlight_limits(dt):
    observer = ephem.Observer()
    observer.date = datetime.datetime(dt.datetime.year,
                                      dt.datetime.month,
                                      dt.datetime.day,
                                      0, 0, 1)
    observer.lat = site_coordinates.latitude.value
    observer.long = site_coordinates.longitude.value    
    observer.horizon = '-18'
    
    return (observer.previous_rising(ephem.Sun(), use_center=True).datetime(),
            observer.next_setting(ephem.Sun(), use_center=True).datetime())

def transit_is_visible(dt, sky_position, minimum_elevation=30. * u.degree):
    night = dt.datetime.date()  
    
    if night not in valid_dates:
        return False
    
    twighlight_limits = compute_twighlight_limits(dt)
    if dt < twighlight_limits[0] or dt > twighlight_limits[1]:
        return False
    
    if altitude(sky_position) < minimum_elevation:
        return False
    
    return True  

def run_simulation(index):
    periods, fraction = [], []
    i = 0
    while True:
        period = d.draw()
#         if np.random.uniform() > transit_probability(period):
#             continue

        sky_position = coord.SkyCoord(np.random.choice(ra_choices), np.random.choice(dec_choices), unit=(u.degree, u.degree))
        epoch = time.Time(np.random.uniform(start_jd, end_jd), format='jd')

        n_visible_transits = 0
        for transit in np.arange(365.25 / period):
            midpoint = epoch + time.TimeDelta(transit * period, format='jd')
            if transit_is_visible(midpoint, sky_position):
                n_visible_transits += 1

        if n_visible_transits > 0:
            i += 1
        periods.append(period)
        fraction.append(float(n_visible_transits) / n_requested_transits)

        if i >= n_planets_to_draw:
            break

    periods, fraction = [np.asarray(val) for val in [periods, fraction]]
    ind = periods.argsort()
    periods, fraction = [data[ind] for data in [periods, fraction]]
    return periods, fraction

"""
Explanation: Drawing planets
Now we can start drawing planets and determining if they are detectable. First we define the "is visible" function:
End of explanation
"""


pool = mp.Pool()
draws = np.arange(ndraws)
results = pool.map(run_simulation, draws)

periods = np.hstack([row[0] for row in results])
fraction = np.hstack([row[1] for row in results])

"""
Explanation: Run multiple copies of the simulation function in parallel, for speed, and then combine the results
End of explanation
"""


out = np.array(list(zip(periods, fraction)), dtype=[('periods', float), ('fraction', float)])
np.save('results', out)

"""
Explanation: Save the results for later analysis
End of explanation
"""


by, bx, _ = stats.binned_statistic(np.log10(periods), fraction, bins=25)
norm, _, _ = stats.binned_statistic(np.log10(periods), periods, bins=bx)
by = by.astype(float) / norm
by[by != by] = 0.

"""
Explanation: These are a series of delta functions, so we need to bin up in the x axis. We also need to normalise by the number of planets in each period window.
End of explanation
"""


l, = plt.semilogx(10 ** bx[:-1], by, drawstyle='steps-post')
ticks = [1., 2., 5., 10., 20]
plt.xticks(ticks, ticks)
plt.xlim(2, 20)
target = plt.axhline(n_requested_transits, color='k', ls=':')
plt.legend([l, target], ['Window function', 'Target'], loc='best')
plt.xlabel(r'Orbital period / days')
plt.ylabel(r'Detected transits')
plt.tight_layout()
for extension in 'png', 'pdf':
    plt.savefig('window-function.{}'.format(extension))

"""
Explanation: Now plot the window function
End of explanation
"""


