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

SECONDS_PER_DAY = 86400.0
JD_UNIX_EPOCH = 2440587.5

G = 2.9591220828559093e-4  # AU^3 / day^2
MASS_SUN = 1.0
MU_SUN = G * MASS_SUN

MU_EARTH = 398600.4418  # km^3 / s^2
AU_KM = 149597870.700
DAY_SEC = 86400.0
MOON_DT = 300.0


def accel_sun(r):
    d = np.linalg.norm(r)
    return -MU_SUN * r / (d ** 3 + 1e-12)


def euler_step(r, v, dt_days):
    a = accel_sun(r)
    return r + v * dt_days, v + a * dt_days


def accel_earth_moon(r_km):
    d = np.linalg.norm(r_km)
    return -MU_EARTH * r_km / (d ** 3 + 1e-12)


def rk4_step(r, v, dt, accel_fn):
    def f(r, v):
        return v, accel_fn(r)

    k1r, k1v = f(r, v)
    k2r, k2v = f(r + 0.5 * dt * k1r, v + 0.5 * dt * k1v)
    k3r, k3v = f(r + 0.5 * dt * k2r, v + 0.5 * dt * k2v)
    k4r, k4v = f(r + dt * k3r, v + dt * k3v)

    r_new = r + (dt / 6) * (k1r + 2 * k2r + 2 * k3r + k4r)
    v_new = v + (dt / 6) * (k1v + 2 * k2v + 2 * k3v + k4v)
    return r_new, v_new


def propagate_planet_grcs(idx, unix_time, target_unix, planet_table, sun_data):
    dt_days = (target_unix - unix_time) / SECONDS_PER_DAY

    r_p = planet_table[idx, :3] + sun_data[idx, :3]
    v_p = planet_table[idx, 3:6] + sun_data[idx, 3:6]

    r_e = sun_data[idx, :3]
    v_e = sun_data[idx, 3:6]

    if abs(dt_days) < 1e-9:
        return planet_table[idx, :3].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)

    return r_p - r_e


def propagate_moon_grcs_fine(idx, unix_time, target_unix, moon_table):
    dt_total = target_unix - unix_time
    if abs(dt_total) < 1e-6:
        return moon_table[idx, :3].copy()

    r = moon_table[idx, :3] * AU_KM
    v = moon_table[idx, 3:] * AU_KM / DAY_SEC

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

    while remaining > 0:
        h = min(MOON_DT, remaining) * sign
        r, v = rk4_step(r, v, h, accel_earth_moon)
        remaining -= abs(h)

    return r / AU_KM


def earth_rotation_angle_from_unix(t_unix):
    JD = JD_UNIX_EPOCH + t_unix / SECONDS_PER_DAY
    d = JD - 2451545.0
    theta = 2 * np.pi * (0.7790572732640 + 1.00273781191135448 * d)
    return theta % (2 * np.pi)


def eci_to_ecef(r_eci, t_unix):
    theta = earth_rotation_angle_from_unix(t_unix)
    c, s = np.cos(theta), np.sin(theta)
    R = np.array([[c, s, 0], [-s, c, 0], [0, 0, 1]])
    return R @ r_eci


def latlon_to_ecef(lat, lon, radius=6371.0):
    lat, lon = np.radians(lat), np.radians(lon)
    return np.array([
        radius * np.cos(lat) * np.cos(lon),
        radius * np.cos(lat) * np.sin(lon),
        radius * np.sin(lat),
    ])


def is_visible_est(planet_ecef, user_lat, user_lon):
    user = latlon_to_ecef(user_lat, user_lon)
    return np.dot(planet_ecef, user) > 0


def closest_time_index(times, target):
    return np.searchsorted(times, target, side="left").clip(0, len(times) - 1)


def check_planet_visibility(
        planet_name: str,
        utc_timestamp,
        user_lat: float,
        user_lon: float,
        data_dir: str = "data",
):
    """
    Args:
        planet_name: "mars", "moon", "venus", etc
        utc_timestamp:
            - unix seconds (float or int), OR
            - timezone-aware datetime in UTC
        user_lat, user_lon: degrees

    Returns:
        dict with:
            est_overhead   -> bool (your estimator)
            true_overhead  -> bool (skyfield)
            match          -> bool
    """

    if isinstance(utc_timestamp, datetime):
        if utc_timestamp.tzinfo is None:
            raise ValueError("datetime must be timezone-aware UTC")
        unix_time = utc_timestamp.timestamp()
        dt_utc = utc_timestamp.astimezone(timezone.utc)
    else:
        unix_time = float(utc_timestamp)
        dt_utc = datetime.fromtimestamp(unix_time, tz=timezone.utc)

    planet_table = np.load(f"{data_dir}/{planet_name}_sv_table.npy")
    time_table = np.load(f"{data_dir}/time_table.npy")
    sun_data = np.load(f"{data_dir}/sun_sv_table.npy")

    idx = closest_time_index(time_table, unix_time)

    if planet_name.lower() == "moon":
        r_eci = propagate_moon_grcs_fine(idx, time_table[idx], unix_time, planet_table)
    else:
        r_eci = propagate_planet_grcs(idx, time_table[idx], unix_time, planet_table, sun_data)

    planet_ecef = eci_to_ecef(r_eci, unix_time)
    est_overhead = is_visible_est(planet_ecef, user_lat, user_lon)

    true_overhead = is_planet_overhead(
        planet_name.lower(),
        dt_utc,
        user_lat,
        user_lon,
    )

    return {
        "planet": planet_name.lower(),
        "est_overhead": bool(est_overhead),
        "true_overhead": bool(true_overhead),
        "match": bool(est_overhead == true_overhead),
    }