import numpy as np

# ------------------------------------------------------
#  Utility functions
# ------------------------------------------------------

def los_from_az_el(az, el):
    """Convert az/el (radians) to topocentric LOS unit vector."""
    # Local topocentric frame: x = east, y = north, z = up
    x = np.cos(el) * np.sin(az)
    y = np.cos(el) * np.cos(az)
    z = np.sin(el)
    return np.vstack((x, y, z)).T


def quadratic_fit(t, values):
    """Fit a quadratic polynomial to component values vs time."""
    A = np.vstack([t**2, t, np.ones_like(t)]).T
    coeff, *_ = np.linalg.lstsq(A, values, rcond=None)
    return coeff  # a, b, c


def eval_quadratic(coeff, t):
    a, b, c = coeff
    return a * t**2 + b * t + c


def pca_plane_projection(R):
    """
    Fit a plane to the set of 3D points R (Nx3).
    Return projected points and plane normal.
    """
    mean_r = R.mean(axis=0)
    X = R - mean_r
    C = X.T @ X / len(R)
    eigvals, eigvecs = np.linalg.eigh(C)
    # Smallest eigenvalue eigenvector = plane normal
    n = eigvecs[:, np.argmin(eigvals)]
    # project
    proj = R - np.outer((R - mean_r) @ n, n)
    return proj, n


def reject_outliers(residuals, k=3.0):
    """Return boolean mask: True = keep point."""
    med = np.median(residuals)
    mad = np.median(np.abs(residuals - med))
    if mad == 0:
        return np.ones_like(residuals, dtype=bool)
    z = 0.6745 * (residuals - med) / mad
    return np.abs(z) < k


# ------------------------------------------------------
#  Main preprocessing pipeline
# ------------------------------------------------------

def preprocess_for_iod(t, az, el, rho, r_site_eci, el_min_deg=20.0):
    """
    Implements steps 2 → 6 from the earlier list.
    Returns: r1, r2, r3 (3 clean ECI position vectors)
    """

    # -------------------------
    # Step 2: Drop low elevation
    # -------------------------
    mask = (el * 180/np.pi) > el_min_deg
    t = t[mask]
    az = az[mask]
    el = el[mask]
    rho = rho[mask]
    r_site_eci = r_site_eci[mask]

    # -------------------------
    # Build raw ECI positions
    # -------------------------
    los = los_from_az_el(az, el)
    R = r_site_eci + rho[:, None] * los   # Nx3

    # -------------------------
    # Step 3: Quadratic fit on each component
    # -------------------------
    t0 = t - t.mean()  # center time for numerical stability

    coeff_x = quadratic_fit(t0, R[:, 0])
    coeff_y = quadratic_fit(t0, R[:, 1])
    coeff_z = quadratic_fit(t0, R[:, 2])

    R_fit = np.column_stack([
        eval_quadratic(coeff_x, t0),
        eval_quadratic(coeff_y, t0),
        eval_quadratic(coeff_z, t0)
    ])

    # -------------------------
    # Step 4: PCA plane projection
    # -------------------------
    R_proj, normal = pca_plane_projection(R_fit)

    # -------------------------
    # Step 5: Outlier rejection
    # -------------------------
    res = np.linalg.norm(R_proj - R_fit, axis=1)
    keep = reject_outliers(res, k=3.0)

    t = t[keep]
    R_proj = R_proj[keep]

    # -------------------------
    # Step 6: Pick 3 clean points around max elevation
    # -------------------------
    # Use the time closest to mid-pass (middle sample)
    mid_idx = len(t) // 2
    t2 = t[mid_idx]

    # Choose t1 and t3 symmetrically by time spacing
    # You may tune dt_window depending on pass length
    dt_window = (t[-1] - t[0]) * 0.25  # quarter-span on each side

    t1 = t2 - dt_window
    t3 = t2 + dt_window

    # interpolate smoothed/proj positions at t1, t2, t3
    def interp_R(t_query):
        return np.column_stack([
            np.interp(t_query, t, R_proj[:, 0]),
            np.interp(t_query, t, R_proj[:, 1]),
            np.interp(t_query, t, R_proj[:, 2])
        ])

    r1 = interp_R(t1)[0]
    r2 = interp_R(t2)[0]
    r3 = interp_R(t3)[0]

    return r1, r2, r3