# Import python modules.
from __future__ import division
from itertools import izip
import time
import datetime

import numpy as np
import matplotlib.pyplot as plt
import astropy.units as u
import astropy.table as table
import astropy.io.ascii as ascii
from PyAstronomy.pyasl import KeplerEllipse

%matplotlib inline

"""
Explanation: How likely is it to detect a certain Cepheid?
In this short notebook, I want to calculate how likely it is that orbital motion of a Cepheid can be detected in radial velocity (RV) measurements, given an observational uncertainty and an observing sequence.
Clearly, it is easier to see the RV modulation if the companion is more massive and the orbit is shorter, but even in the best case, there is still a certain probability that no RV varibility is detected, because the orbital plane conincides with the plane of the sky.
In this notebook, I keep the code for the calculation together with some plots that illustrate the results and some notes to help you understand what I did and to help me remember how I did it.
End of explanation
"""


def generate_LOSs(n):
    '''Generate a random distribution of Line of sights
    
    Generate a random distribution of sightlines in cartesian coordinates,
    with an isotropic distribution (no clustering around the poles).
    Analytically, a truely isotropic distribution could be generated using
    geodesic grids, but this is much faster to program and (for large n) just as good.

    Parameters
    ----------
    n : integer
        number of sightlines to be generated

    Returns
    -------
    los : np.array of dim [3,m] with m < n
        x,y,z cartesian coordinates for each sightline. Only a subset of all generated
        sightlines is returend, selected in such a way that they are (within the limits
        of their randomness) isotropically distributed over the sphere.
    '''
    # Make points in box with x,y,z in [-1,1].
    los = np.random.random((3, n)) * 2 - 1.
    # Throw out all points beyond a unit sphere.
    # Yes, this is a wasteful implementation, but it's fast enough that I don't care.
    r = np.sqrt((los * los).sum(axis=0))
    ind = ((r <= 1) & (r >= 1e-6))  # Throw out inner values where round-off errors matter.
    return los[:, ind] / r[ind]

def get_LOSs(n = 1000):
    '''Generate a random distribution of Line of sights
    
    Generate a random distribution of sightlines in cartesian coordinates,
    with an isotropic distribution (no clustering around the poles).
    Analytically, a truely isotropic distribution could be generated using
    geodesic grids, but this is much faster to program and (for large n) just as good.

    Parameters
    ----------
    n : integer
        number of sightlines to be generated

    Returns
    -------
    los : np.array of dim [3,m] with m < n
        x,y,z cartesian coordinates for each sightline. 
    '''
    while True:
        los = generate_LOSs(4 * n) # 2 should be big enough that this usually succeeds.
        if los.shape[1] >= n:
            break
    return los[:, :n]

def calc_many_vobs(times, a, p, e, n_tau=50, n_los=1000, los=None):
    '''Calculate radial velocities for one orbit and many LOS.
    
    For one orbit with given semi-major axis, period, and eccentricity calculate
    the radial velocity (RV) signal for a given set of observaton times. 
    This calculation is done for larger number of sightlines to the system and for
    different starting points in the orbit.
    
    Parameters
    ----------
    a : float
        semi-major axis (in AU)
    p : float
        period (in years)
    e : float
        eccentricity of orbit
    n_tau : integer
        The calculation will be done for ``n_tau`` different stating points
        regularly distributed over the orbit, because it does matter if a 
        star is observed close to periastron or apastron.
    n_los : integer
        number of lines-of-sight to be evaluated
    los : None or np.array of dim [3, n]
        If ``None`` then ``n_los`` sightlines will be randomnly generated.
        Otherwise a defined set of sightlines can be passed in as an array.
        The elements of the array have to be the cartesian coordinates of 
        points on the unit sphere.
    
    Returns
    -------
    v_obs : array of astropy.quantity with dim [n_tau, len(times), n_los]
        This holds the calculated RV values that would be observed.
    ''' 
    
    if los is None:
        los = get_LOSs(n_los)
    else:
        n_los = los.shape[1]
        
    taus = np.linspace(0,p, n_tau, endpoint=False)
    v_obs = np.zeros((n_tau, len(times), n_los))
    for j, tau in enumerate(taus):
        ke = KeplerEllipse(a, p, e, tau)
        vxyz = ke.xyzVel(times)
        for i in range(len(times)):
            v_obs[j, i,:] = np.dot(vxyz[i,:], los)

    return v_obs *u.AU / u.year

def calc_maxdv(*args, **kwargs):
    '''Run ``calc_many_vobs`` and simplyfy its output.
    
    See ``calc_many_vobs`` for accepted parameters.
    
    Returns
    -------
    maxdv : astropy.quantity
        Maximum differences between two RV values that would be 
        observed for the given orbital parameters for each LOS and tau.
    '''
    v_obs = calc_many_vobs(*args, **kwargs)
    maxdv = np.max(v_obs, axis=1) - np.min(v_obs, axis=1)
    return maxdv.flatten()

def prob_to_detect(dv_sensitivity, *args, **kwargs):
    '''Calculate the probability to detect binarity.
    
    Parameters
    ----------
    dv_sensitivity : astropy.quantity
        Minimum delta(RV) required for detection.
        
    See ``calc_many_vobs`` for the remaining  parameters.
    
    Returns
    -------
    prob : float
        Probability of detection
    
    Example
    -------
    
    >>> times = np.arange(0,1,.1)
    >>> prob_to_detect((20.*u.km/u.s), times, 1 , 1, 0)
    0.94
    (approximate result, since based on random numbers)
    '''
    maxdv = calc_maxdv(*args, **kwargs)
    return (maxdv > dv_sensitivity).sum() / maxdv.size

"""
Explanation: The observed magnitude of the RV variation depends crucially on the angle between the line-of-sight (LOS) and the orbit, particularly the inclination and the longitude of periastron. They way I handle that here is that for each orbit I calculate the observed RV values for a large number of sightlines. I then check for which fraction of those sightlines the RV change is large enough to be detected. 
To obtain a fair estimate of the probability of detection I need an isotropic distribution of sightlines. There are schemes in the literature to tesselate a sphere in such a way that each tesselate has the same area ("geodesic grid") and then I could use the mid-point of this tesselate for my sightlines. However, all those schemes are a pain to implement (or if implemented already, a pain to use) due to all the sin and cos involved. It's not impossible, but for now I just generate a large number of random LOS.
End of explanation
"""


# Observation cadence. Here: 10 years, once per year, regular intervals.
times = np.arange(0.,10.,1.)

# Grid for semi-major axis between the two stars (in AU)
a = np.logspace(0.1,2)
# Masses (in M_sun)
m1 = 5   # primary
m2 = np.array([0.02, 0.05, 0.08, 0.1,0.5, 0.8,1,2,3,4,5])  # secondary
M = m1 + m2  # total mass of system
agrid, m2grid = np.meshgrid(a, m2)
# semi-major axis of primary
a1grid = agrid * m2grid / (m1 + m2grid)
# Period
Pgrid = (m1 + m2grid) * np.sqrt(a1grid**3 / m2grid**3)

fig = plt.figure(figsize=(10,5))
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)

cs = ax1.contour(Pgrid, [0.1,0.2,1,5,25,50,100,300])
ax1.clabel(cs)
ax1.set_title('Period grid [yr]')
cs = ax2.contour(m2grid/m1)
ax2.clabel(cs)
ax2.set_title('mass ratio grid')

"""
Explanation: Set up the grid of parameters. I chose the secondary mass $m_2$ and the total semi-major axis $a$ as parameters. From this I calculate the semi-major axis of the primary and the period assuming a constant primary mass of $m_1 = 6 M_{\odot}$.
End of explanation
"""


cs = plt.contour(agrid, m2grid, Pgrid, [10,20,30,60,120,240.])
plt.clabel(cs)
plt.xlabel('Semi-major axis [AU]')
plt.ylabel('mass of secondary [$M_{\odot}$]')
plt.title('Orbital period [years]')

"""
Explanation: This first plot shows the orbital period in years. This shall serve as a comparison for the following plots.
End of explanation
"""


prop_e0 = np.zeros_like(Pgrid)
for x, y in np.ndindex(Pgrid.shape):
    prop_e0[x,y] = prob_to_detect(2.*u.km/u.s, times, a1grid[x,y], Pgrid[x,y], 0)

def plot_m_x(array):
    fig = plt.figure(figsize=(10,5))
    ax1 = fig.add_subplot(121)
    ax2 = fig.add_subplot(122)
    for ax in [ax1, ax2]:
        cs = ax.contour(agrid, m2grid, array)
        ax.clabel(cs)
        ax.set_xlabel('Semi-major axis [AU]')
        ax.set_ylabel('mass of secondary [$M_{\odot}$]')
        ax.set_title('Probability of detection for $e$ = 0')
    ax2.set_title('Same data, logarithmic axes')
    ax2.set_xscale('log')
    ax2.set_yscale('log')
    return fig

def plot_q_p(array):
    fig = plt.figure(figsize=(10,5))
    ax1 = fig.add_subplot(121)
    ax2 = fig.add_subplot(122)
    for ax in [ax1, ax2]:
        cs = ax.contour(Pgrid, m2grid/m1, array)
        ax.clabel(cs)
        ax.set_xlabel('Period [yr]')
        ax.set_ylabel('mass ratio [$M_{\odot}$]')
        ax.set_title('Probability of detection for $e$ = 0')
    ax2.set_title('Same data, logarithmic axes')
    ax2.set_xscale('log')
    ax2.set_yscale('log')
    return fig


fig = plot_m_x(prop_e0)

# Same plot as above, but with different axis

fig = plot_q_p(prop_e0)

"""
Explanation: Next, I show the probability of detection for circular orbits. The plots on the left and on the right show the same data, but one has a linear and one has logarithmic axis. Both of them are a little coarse, because the grid is not very dense. That can easily be changed if we want to make nicer looking plots.
On the right plot, there are islands where the chance to detect a binary is essentially 0. This happens when the sampling frequency is (or is close to) an integer multiple of the orbital period. Again, this will look better (and more regular) when I use a denser grid.
Those thoughts aside, we see that we expect to detect almost every solar-mass or heavier companion out to 20 AU, and have a chance of 50% out to about 40 AU.
For all plots I use 2 km/s as minimum $\Delta v$ that would be detected in the observations.
End of explanation
"""


times = np.arange(0.,10.,1.)
prop_e05 = np.zeros_like(Pgrid)
for x, y in np.ndindex(Pgrid.shape):
    prop_e05[x,y] = prob_to_detect(2.*u.km/u.s, times, a1grid[x,y], Pgrid[x,y], 0.5)

cs = plt.contour(agrid, m2grid, prop_e05)
plt.clabel(cs)
plt.xlabel('Semi-major axis [AU]')
plt.ylabel('mass of secondary [$M_{\odot}$]')
plt.title('Probability of detection for $e$ = 0.5')

"""
Explanation: Now, we repeat the excersice for $e=0.5$. As you can see the probability contours move inwards a little, but overall that is not a dramatic change.
End of explanation
"""


times = np.arange(0.,20.,1.)
prop_e05 = np.zeros_like(Pgrid)
for x, y in np.ndindex(Pgrid.shape):
    prop_e05[x,y] = prob_to_detect(2.*u.km/u.s, times, a1grid[x,y], Pgrid[x,y], 0)

cs = plt.contour(agrid, m2grid, prop_e05)
plt.clabel(cs)
plt.xlabel('Semi-major axis [AU]')
plt.ylabel('mass of secondary [$M_{\odot}$]')
plt.title('Probability of detection for $e$ = 0 with 20 year baseline' )

"""
Explanation: As a last point, I use circular orbits again to predict how much better this will be with a 20 year baseline. You can use that to predict the number of additional binary systems that will be identified when new data becomes available.
End of explanation
"""


actual_obs_years = table.Table.read('datafile4.txt', format='ascii')
actual_obs_years = actual_obs_years.filled()
actual_obs_years = actual_obs_years.group_by('ID')

"""
Explanation: Run simulations with the actual observing dates
End of explanation
"""


allsims00 = {}
for star, group in izip(actual_obs_years.groups.keys, actual_obs_years.groups):
    print datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S'), 'Working on: ', star[0]
    prop = np.zeros_like(Pgrid)
    for x, y in np.ndindex(Pgrid.shape):
        prop[x,y] = prob_to_detect(2.*u.km/u.s, np.array(group['Year']), a1grid[x,y], Pgrid[x,y], 0.0)
    allsims00[star[0]] = prop

allsims05 = {}
for star, group in izip(actual_obs_years.groups.keys, actual_obs_years.groups):
    print datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S'), 'Working on: ', star[0]
    prop = np.zeros_like(Pgrid)
    for x, y in np.ndindex(Pgrid.shape):
        prop[x, y] = prob_to_detect(2. * u.km/u.s, np.array(group['Year']), a1grid[x, y], Pgrid[x, y], 0.5)
    allsims05[star[0]] = prop

allnames = set(allsims00.keys()) - set(['CK  Cam'])
all00 = np.dstack([allsims00[n] for n in allnames]).mean(axis=2)
all05 = np.dstack([allsims05[n] for n in allnames]).mean(axis=2)

fig = plot_q_p(allsims00[list(allnames)[1]])

fig = plot_q_p(all00)

"""
Explanation: We just do this for eccentricity 0 and 0.5. Since there is little prior information and (see above) it also does not matter much in that range, we will just show that the contours are very similar and no more details are required.
End of explanation
"""


fig = plt.figure(figsize=(10,5))
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)

cs = ax1.contour(Pgrid, [0.1,0.2,1,5,25,50,100,300])
ax1.clabel(cs)
ax1.set_title('Period grid')
cs = ax2.contour(m2grid/m1)
ax2.clabel(cs)
ax2.set_title('mass ratio grid')

print np.min(Pgrid), np.max(Pgrid)

"""
Explanation: How dense are the grids in period and mass ratio space?
End of explanation
"""


fig = plt.figure(figsize=(4,3))
ax = fig.add_subplot(111)
cs = ax.contour(Pgrid, m2grid/m1, all00, [0.95, 0.90, 0.75], linewidths=4, linestyles=['solid', 'dashed','dotted'])
cs2 = ax.contour(Pgrid, m2grid/m1, all05, [0.95, 0.90, 0.75], linestyles=['solid', 'dashed','dotted'])
ax.set_xlim([0,60])
ax.clabel(cs, fmt='%1.2f', use_clabeltext=True, manual=[(40,.4), (30,.6), (20,.5)])

ax.set_xlabel('Period [yr]')
ax.set_ylabel('mass ratio')

#ax.set_xscale("log")
fig.subplots_adjust(left=0.16, bottom=0.16, top=0.97, right=0.97)
fig.savefig('detectionprobability.png')
fig.savefig('detectionprobability.pdf')
fig.savefig('detectionprobability.eps')

"""
Explanation: Get the figures for publication together
End of explanation
"""


