import numpy as np
import random
from planet_overhead import is_planet_overhead
from datetime import datetime, timezone
from tqdm import tqdm

names = ["sun", "moon", "mercury", "venus", "mars",
         "jupiter BARYCENTER", "saturn BARYCENTER", "uranus BARYCENTER",
         "neptune BARYCENTER", "pluto BARYCENTER"]

test_planet = "mars"
USER_LAT = 0
USER_LON = 0
trials = 500

"""
FIRST RUN build_ephemeris_tables.py

This is a code script to test your tables for accuracy 

what this code does:
    we have saved sparse tables of inertial planet (and moon) state vectors (position and velocity).
    given a random user time, lat, and lon, we grab the closest planet state vector, sun vector,
    and unix timestamp.

    using these vectors we approximate the planet’s position, usually using a simple linear
    approximation and the sun. the sun has the biggest effect on planet motion

    we then see where these approximated planet locations (xyz) are relative to Earth, and then
    spin them around based on Earth’s rotation. (ECI to ECEF)

    once we have our planet locations spun to match the rotating Earth, confirming whether it’s
    overhead is easy. lat / lon are also stuck on the spinning Earth.

"""

SECONDS_PER_DAY = 86400.0  # sidereal
JD_UNIX_EPOCH = 2440587.5  # julian date at unix epoch
G = 2.9591220828559093e-4  # units of au^3 / (day^2)
MASS_EARTH = 3.003e-06  # solar mass
MASS_MOON_RATIO = 0.0123000
MASS_SUN = 1.0
SOLAR_MASS_FRACTION = {
    "sun": 1.0,
    "mercury": 1.66e-7,
    "venus": 2.45e-6,
    "earth": 3.03e-6,
    "mars": 3.23e-7,
    "jupiter barycenter": 9.55e-4,
    "saturn barycenter": 2.86e-5,
    "uranus barycenter": 4.37e-6,
    "neptune barycenter": 5.15e-5,
    "pluto barycenter": 7.340e-9
}

MU_SUN = G * MASS_SUN


def accel_sun(r):
    """f = ma => a = F / m => a = r * G * M_sun / |r|^3"""
    # thank you, newton
    d = np.linalg.norm(r)
    return -MU_SUN * r / (d ** 3 + 1e-12)


def euler_step(r, v, dt_days):
    """euler integration — linear."""
    a = accel_sun(r)
    r_new = r + v * dt_days
    v_new = v + a * dt_days
    return r_new, v_new


def propagate_planet_grcs(idx, unix_time, target_unix, planet_table):
    """
    Take one big linear step to move the planet to where it should ~be.

    We need the planet vector to have Earth at its origin, so we also move Earth with a big step,
    then subtract Earth’s new position from the propagated planet position. This effectively
    slides the frame back to where it should be.

    The convenient thing about the GCRS frame is that Earth, by definition, is always at
    xyz = [0, 0, 0] and d(xyz)/dt = [0, 0, 0] (at the time the table was made),
    so we don’t really need Earth’s state vector to move it — we just use the sun table.

    """

    dt_sec = target_unix - unix_time
    dt_days = dt_sec / SECONDS_PER_DAY

    au_xyz = planet_table[idx, 0:3]
    au_per_d = planet_table[idx, 3:6]

    if abs(dt_days) < 1e-9:
        return au_xyz.copy()

    r_p = au_xyz.copy() + sun_data[idx, 0:3]
    v_p = au_per_d.copy() + sun_data[idx, 3:6]

    # earth helio = sun wrt earth
    r_e = sun_data[idx, 0:3].copy()
    v_e = sun_data[idx, 3:6].copy()

    r_p, v_p = euler_step(r_p, v_p, dt_days)
    r_e, v_e = euler_step(r_e, v_e, dt_days)

    # r_p, v_p = rk4_step_sun_planet_system(r_p, v_p, dt_days)
    # r_e, v_e = rk4_step_sun_planet_system(r_e, v_e, dt_days)

    r_geo = r_p - r_e
    return r_geo


### moon stuff
MU_EARTH = 398600.4418


def accel_earth_moon(r_km):
    """f = ma => a = F / m => a = r * G * M_sun / |r|^3"""
    d = np.linalg.norm(r_km)
    return -MU_EARTH * r_km / (d ** 3 + 1e-12)


AU_KM = 149597870.700
DAY_SEC = 86400.0
MOON_DT = 300.0  # the moon needs special care — smaller time steps and higher-order integration


def rk4_step_moon_earth_system(r, v, dt):
    """
    rk4 integrator for approximating 2nd-order differential equations.
    """

    def f(r, v):
        return v, accel_earth_moon(r)

    k1_r, k1_v = f(r, v)
    k2_r, k2_v = f(
        r + 0.5 * dt * k1_r,
        v + 0.5 * dt * k1_v
    )
    k3_r, k3_v = f(
        r + 0.5 * dt * k2_r,
        v + 0.5 * dt * k2_v
    )
    k4_r, k4_v = f(
        r + dt * k3_r,
        v + dt * k3_v
    )
    r_new = r + (dt / 6.0) * (k1_r + 2 * k2_r + 2 * k3_r + k4_r)
    v_new = v + (dt / 6.0) * (k1_v + 2 * k2_v + 2 * k3_v + k4_v)
    return r_new, v_new


def rk4_step_sun_planet_system(r, v, dt):
    """
    rk4 integrator for approximating 2nd-order differential equations.
    """

    def f(r, v):
        return v, accel_sun(r)

    k1_r, k1_v = f(r, v)
    k2_r, k2_v = f(
        r + 0.5 * dt * k1_r,
        v + 0.5 * dt * k1_v
    )
    k3_r, k3_v = f(
        r + 0.5 * dt * k2_r,
        v + 0.5 * dt * k2_v
    )
    k4_r, k4_v = f(
        r + dt * k3_r,
        v + dt * k3_v
    )
    r_new = r + (dt / 6.0) * (k1_r + 2 * k2_r + 2 * k3_r + k4_r)
    v_new = v + (dt / 6.0) * (k1_v + 2 * k2_v + 2 * k3_v + k4_v)
    return r_new, v_new


def propagate_moon_grcs_fine(idx, unix_time, target_unix, moon_table):
    """
    propagate the Moon around Earth using rk4.
    Earth's gravity completely dominates for the moon compared to everything else.
    """

    au_xyz = moon_table[idx][0:3]
    au_per_d = moon_table[idx][3:]

    dt_total = target_unix - unix_time
    if abs(dt_total) < 1e-6:
        return au_xyz.copy()

    r = au_xyz * AU_KM
    v = au_per_d * AU_KM / DAY_SEC

    sign = np.sign(dt_total)
    dt_remaining = abs(dt_total)

    dt = MOON_DT

    while dt_remaining > 0.0:
        h = min(dt, dt_remaining) * sign
        r, v = rk4_step_moon_earth_system(r, v, h)
        dt_remaining -= abs(h)

    return r / AU_KM


###########################################################

## earth rotate

def eci_to_ecef_target_timestamp(r_eci, t_unix_start, t_unix_target):
    """
    Our GCRS vectors don’t care about the spinning Earth (if they did,
    classical physics approximations wouldnt work).

    GCRS shares the same origin (0, 0, 0) with ECEF, so we just rotate the vectors to match
    Earth’s current spin.

    IMHO I can’t believe how accurate earth_rotation_angle_from_unix is. This was the part I was most worried about,
    but grabbing some numbers online actually worked surprisingly well.
    """
    theta = earth_rotation_angle_from_unix(t_unix_start + (t_unix_target - t_unix_start))
    c = np.cos(theta)
    s = np.sin(theta)
    R = np.array([
        [c, s, 0.0],
        [-s, c, 0.0],
        [0.0, 0.0, 1.0]
    ])
    return R @ r_eci


def earth_rotation_angle_from_unix(t_unix):
    # unix seconds to Julian Date
    JD = JD_UNIX_EPOCH + t_unix / SECONDS_PER_DAY

    d = JD - 2451545.0

    theta = 2.0 * np.pi * (
            0.7790572732640 +
            1.00273781191135448 * d
    )

    return theta % (2.0 * np.pi)


def latlon_to_ecef(lat, lon, radius=6371.0):
    # classic equation - z axis goes thru the north pole
    lat_rad = np.radians(lat)
    lon_rad = np.radians(lon)
    x = radius * np.cos(lat_rad) * np.cos(lon_rad)
    y = radius * np.cos(lat_rad) * np.sin(lon_rad)
    z = radius * np.sin(lat_rad)
    return np.array([x, y, z])


def is_visible_est(planet_ecef, user_lat, user_lon):
    # once agin the dot product is king
    user_ecef = latlon_to_ecef(user_lat, user_lon)
    return np.dot(planet_ecef, user_ecef) > 0


#########################

def closest_time_index(unix_time, n, target):
    # given our sparse lookup tables (ephemeris), find the closest saved timestamp
    if target <= unix_time[0]:
        return 0
    if target >= unix_time[n - 1]:
        return n - 1
    left = 0
    right = n - 1
    while left <= right:
        mid = (left + right) // 2
        if unix_time[mid] == target:
            return mid
        elif unix_time[mid] < target:
            left = mid + 1
        else:
            right = mid - 1
    if (target - unix_time[right]) <= (unix_time[left] - target):
        return right
    else:
        return left


filename = f"data/{test_planet}_sv_table.npy"
test_planet_table = np.load(filename)
time_table = np.load("data/time_table.npy")
sun_data = np.load("data/sun_sv_table.npy")

accuracy = 0

for i in tqdm(range(trials)):

    USER_TIMESTAMP = random.uniform(time_table[0], time_table[-1])
    USER_TIMESTAMP = int(datetime.now(timezone.utc).timestamp()) + 60*60*20*262
    # USER_TIMESTAMP = datetime.now(timezone.utc).timestamp() + 60*60*24*162

    idx = closest_time_index(time_table, len(time_table), USER_TIMESTAMP)

    if test_planet == "moon":
        r_prop = propagate_moon_grcs_fine(idx, time_table[idx], USER_TIMESTAMP, test_planet_table)
    else:
        r_prop = propagate_planet_grcs(idx, time_table[idx], USER_TIMESTAMP, test_planet_table)

    itrs_test_position_est = eci_to_ecef_target_timestamp(r_prop, time_table[idx], USER_TIMESTAMP)
    person_lat_lon_est = latlon_to_ecef(USER_LAT, USER_LON)
    est_overhead = is_visible_est(itrs_test_position_est, USER_LAT, USER_LON)  ## bool

    ## skyfield truth
    dt_utc = datetime.fromtimestamp(USER_TIMESTAMP, tz=timezone.utc)
    planet_name_sf = test_planet.lower()

    true_overhead = is_planet_overhead(
        planet_name_sf,
        dt_utc,
        USER_LAT,
        USER_LON,
    )

    accuracy += (est_overhead == true_overhead)

print(f"accuracy {100* accuracy / trials}%")