"""Small independent orbit/projection functions; Python standard library only.

Distances AU, time TDB days, velocities AU/day. No author code is copied.
"""
import math

AU_METRES = 149597870700.0
DAY_SECONDS = 86400.0
GM_SI = 1.3271244004193938e20
MU = GM_SI * DAY_SECONDS**2 / AU_METRES**3
C_AU_DAY = 299792458.0 * DAY_SECONDS / AU_METRES

def dot(first, second):
    return sum(a*b for a, b in zip(first, second))

def magnitude(vector):
    return math.sqrt(dot(vector, vector))

def add(first, second):
    return [a+b for a, b in zip(first, second)]

def scale(vector, factor):
    return [value*factor for value in vector]

def subtract(first, second):
    return add(first, scale(second, -1))

def stumpff(argument):
    if abs(argument) < 1e-5:
        return (0.5-argument/24+argument**2/720-argument**3/40320+argument**4/3628800,
                1/6-argument/120+argument**2/5040-argument**3/362880+argument**4/39916800)
    if argument > 0:
        root = math.sqrt(argument)
        return (1-math.cos(root))/argument, (root-math.sin(root))/root**3
    root = math.sqrt(-argument)
    return (math.cosh(root)-1)/(-argument), (math.sinh(root)-root)/root**3

def kepler(position, velocity, elapsed_days, gravity=MU):
    radius = magnitude(position)
    root_gravity = math.sqrt(gravity)
    radial_product = dot(position, velocity)
    alpha = 2/radius-dot(velocity, velocity)/gravity
    anomaly = root_gravity*elapsed_days/radius
    if alpha > 1e-6:
        anomaly = root_gravity*alpha*elapsed_days
    for _ in range(80):
        z = alpha*anomaly**2
        cosine, sine = stumpff(z)
        value = radial_product/root_gravity*anomaly**2*cosine + (1-alpha*radius)*anomaly**3*sine + radius*anomaly-root_gravity*elapsed_days
        derivative = radial_product/root_gravity*anomaly*(1-z*sine)+(1-alpha*radius)*anomaly**2*cosine+radius
        correction = value/derivative
        limit = max(1, abs(anomaly))
        correction = max(-limit, min(limit, correction))
        anomaly -= correction
        if abs(correction) < 1e-12:
            break
    else:
        raise ValueError('Universal conic solve did not converge')
    cosine, sine = stumpff(alpha*anomaly**2)
    f = 1-anomaly**2*cosine/radius
    g = elapsed_days-anomaly**3*sine/root_gravity
    final_position = add(scale(position, f), scale(velocity, g))
    final_radius = magnitude(final_position)
    f_dot = root_gravity/(radius*final_radius)*(alpha*anomaly**3*sine-anomaly)
    g_dot = 1-anomaly**2*cosine/final_radius
    return final_position, add(scale(position, f_dot), scale(velocity, g_dot))

def rk4(position, velocity, elapsed_days, gravity=MU, maximum_step=0.125, beta_pr=0):
    """Separate fine-step integration for independent conic and energy-loss checks."""
    count = max(1, math.ceil(abs(elapsed_days)/maximum_step))
    step = elapsed_days/count
    state = list(position)+list(velocity)
    def derivative(values):
        radial = values[:3]
        speed = values[3:]
        radius = magnitude(radial)
        acceleration = scale(radial, -gravity/radius**3)
        if beta_pr:
            direction = scale(radial, 1/radius)
            radial_speed = dot(speed, direction)
            acceleration = add(acceleration, scale(add(scale(direction, radial_speed), speed), -beta_pr*MU/(C_AU_DAY*radius**2)))
        return speed+acceleration
    for _ in range(count):
        first = derivative(state)
        second = derivative(add(state, scale(first, step/2)))
        third = derivative(add(state, scale(second, step/2)))
        fourth = derivative(add(state, scale(third, step)))
        state = [value+step*(a+2*b+2*c+d)/6 for value,a,b,c,d in zip(state,first,second,third,fourth)]
    return state[:3], state[3:]

def sky_basis(right_ascension_degrees, declination_degrees):
    ra, dec = map(math.radians, [right_ascension_degrees, declination_degrees])
    return ([math.cos(dec)*math.cos(ra), math.cos(dec)*math.sin(ra), math.sin(dec)],
            [-math.sin(ra), math.cos(ra), 0],
            [-math.sin(dec)*math.cos(ra), -math.sin(dec)*math.sin(ra), math.cos(dec)])

def position_angle(vector, ra, dec):
    _, east, north = sky_basis(ra, dec)
    return math.degrees(math.atan2(dot(vector, east), dot(vector, north))) % 360

def tangent_plane(vector, ra, dec):
    centre, east, north = sky_basis(ra, dec)
    denominator = dot(vector, centre)
    if denominator <= 0:
        raise ValueError('Point lies outside this tangent hemisphere')
    return [math.degrees(dot(vector, east)/denominator), math.degrees(dot(vector, north)/denominator)]

def inverse_cd(wcs, tangent_degrees):
    a,b,c,d = [wcs[key] for key in ['CD1_1','CD1_2','CD2_1','CD2_2']]
    determinant = a*d-b*c
    x,y = tangent_degrees
    return [(d*x-b*y)/determinant, (-c*x+a*y)/determinant]

def sky_to_native_pixel(vector, wcs):
    offsets = inverse_cd(wcs, tangent_plane(vector, wcs['CRVAL1'], wcs['CRVAL2']))
    return [offsets[0]+wcs['CRPIX1']-1, offsets[1]+wcs['CRPIX2']-1]