import jplephem
from jplephem.spk import SPK
kernel = SPK.open('/home/daniel/data/de430.bsp')
import astropy
from astropy.coordinates import SkyCoord, ICRS, FK5, EarthLocation, AltAz
from astropy.time import Time
import astropy.units as u

import math
import numpy as np
@u.quantity_input(canon_velocity=(u.meter / u.second))
def v_sun(source, apex = "18h03m50.29s +30d00m16.8s", 
                  canon_velocity = 20.0 * 1000 * u.meter / u.second):
    """
    Calculate the velocity of the sun relative to a point in the sky.
    
    Parameters
    ----------
    source : `astropy.coordinate.SkyCoord` object
        The point on the sky which the Earth's projected velocity should be 
        calculated for.
        
    apex : str, optional
        The point in the sky which the Sun is travelling towards. This should 
        be given as a string in the format "ra dec". The default value is
        18h03m50.29s +30d00m16.8s.
        
    canon_velocity : float
        The velocity at which the sun is travelling towards `apex`. This 
        defaults to 20 000 m/s. The value must be specified using astropy 
        units.
        
    Returns
    -------
    v_sun : `astropy.quantity`
        The velocity of the sun, relative to the source, with correct units.
    
    Notes
    -----
    The Sun is moving towards a point in the sky close to Vega at a velocity
    of around 20 km/s. In order to calculate the velocity of the Earth towards 
    any other point in the sky we need to project the velocity vector onto
    the position vector of that point in the sky, using a dot product.
    """
    if isinstance(source.obstime.value, str):
        equinox = source.obstime
    else:
        equinox = source.obstime[0]
    c = SkyCoord(apex, equinox="J2000", frame=FK5)
    c = c.transform_to(FK5(equinox=equinox, representation='cartesian'))
    c = c.cartesian
    source = source.cartesian
    v = canon_velocity * c.xyz
    return canon_velocity*np.dot(c.xyz,source.xyz)
def v_obs(source, observer):
    """
    Calculate the velocity of the observatory due to its rotation about the 
    Earth's barycentre in the direction of a point on the sky.
    
    Parameters
    ----------
    source : `astropy.coordinate.SkyCoord` object
        The point on the sky which the Earth's projected velocity should be 
        calculated for.
    observer : `astropy.coordinates.EarthLocation` object
        The location on the Earth of the observatory.
    
    Returns
    -------
    v_obs : `astropy.quantity` object
        The velocity of the observer relative to the source.
    """
    pos = [observer.x.value, observer.y.value, observer.z.value] * u.meter
    vel = [0, 0,  (2*np.pi)/(24.*3600.)]  / u.second # 459
    source = source.transform_to(AltAz(location=observer))
    return np.dot(np.cross(pos, vel), source.cartesian.xyz)*u.meter/u.second
def v_earth(source):
    """
    Calculate the velocity of the Earth relative to some point on the sky.
    
    Parameters
    ----------
    source : `astropy.coordinate.SkyCoord`
        The point on the sky which the Earth's projected velocity should be 
        calculated for.
        
    Returns
    -------
    v_earth : `astropy.quantity` object
        The velocity of the Earth relative to the source.
    """
    time = source.obstime.jd
    position, velocity = kernel[0,3].compute_and_differentiate(time)
    position2, velocity2 = kernel[3,399].compute_and_differentiate(time)
    velocity = velocity - velocity2
    source = source.cartesian
    return np.dot(((velocity * 1000 * u.meter / u.day).to(u.meter / u.second)).T, source.xyz) 
def v_lsr(source, observer):
    """
    Calculate the velocity of the local standard of rest for an observatory in
    the direction of a source on the sky.
    
    Parameters
    ----------
    source : `astropy.coordinate.SkyCoord` object
        The point on the sky which the Earth's projected velocity should be 
        calculated for.
    observer : `astropy.coordinates.EarthLocation` object
        The location on the Earth of the observatory.
    
    Returns
    -------
    v_lsr : `astropy.quantity` object
        The velocity of the observer relative to the local standard of rest.
    """
    v =  v_sun(source) + (v_earth(source) + v_obs(source, observer))
    return v
start_time = Time('2016-01-28T12:15:00', format='isot', scale='utc')
time_range = np.linspace(0, 100, 100)*u.second
times = start_time + time_range
source = SkyCoord(306.51141024*u.degree, 40.84802122*u.degree, frame=FK5, obstime=times)
acre_road = EarthLocation(lat=55.9024278*u.deg, lon=-4.307582*u.deg, height=61*u.m)
v_lsr(source, acre_road)

$[11878.335,~11878.353,~11878.37,~11878.388,~11878.406,~11878.423,~11878.441,~11878.459,~11878.476,~11878.494,~11878.512,~11878.529,~11878.547,~11878.564,~11878.582,~11878.6,~11878.617,~11878.635,~11878.653,~11878.67,~11878.688,~11878.706,~11878.723,~11878.741,~11878.758,~11878.776,~11878.794,~11878.811,~11878.829,~11878.847,~11878.864,~11878.882,~11878.9,~11878.917,~11878.935,~11878.952,~11878.97,~11878.988,~11879.005,~11879.023,~11879.041,~11879.058,~11879.076,~11879.094,~11879.111,~11879.129,~11879.146,~11879.164,~11879.182,~11879.199,~11879.217,~11879.235,~11879.252,~11879.27,~11879.288,~11879.305,~11879.323,~11879.34,~11879.358,~11879.376,~11879.393,~11879.411,~11879.429,~11879.446,~11879.464,~11879.482,~11879.499,~11879.517,~11879.534,~11879.552,~11879.57,~11879.587,~11879.605,~11879.623,~11879.64,~11879.658,~11879.675,~11879.693,~11879.711,~11879.728,~11879.746,~11879.764,~11879.781,~11879.799,~11879.817,~11879.834,~11879.852,~11879.869,~11879.887,~11879.905,~11879.922,~11879.94,~11879.958,~11879.975,~11879.993,~11880.011,~11880.028,~11880.046,~11880.063,~11880.081] \; \mathrm{\frac{m}{s}}$