import numpy as np
from scipy.interpolate import UnivariateSpline

# Onion-peeling inversion (linear Abel) 
def onion_peeling_linear(tec, r):
    """
    Reconstruct electron density profile from calibrated TEC
    using linear-layer onion peeling (invtec / gmrion formulation).

    Parameters
    ----------
    tec : ndarray, shape (n,)
        Calibrated TEC for rays with impact parameters p_i = r_i
    r : ndarray, shape (n+1,)
        Radial grid boundaries (meters), monotonically increasing

    Returns
    -------
    Ne : ndarray, shape (n,)
        Electron density profile (electrons / m^3)
    """

    n = len(tec)
    Ne = np.zeros(n)

    for i in range(n-1, -1, -1):  # peel from top down

        ri = r[i]
        dr = r[i+1] - r[i]

        # Contribution from already-recovered outer shells
        outer_sum = 0.0
        for k in range(i+1, n):
            outer_sum += Ne[k] * (
                np.sqrt(r[k+1]**2 - ri**2)
                - np.sqrt(r[k]**2 - ri**2)
            )

        # Local inversion (linear-layer normalization)
        Ne[i] = (3.0 / 4.0) * (
            tec[i] - outer_sum
        ) / np.sqrt(2.0 * ri * dr)

    return Ne


   #def Cki(r, i, k):                                              
   #    return (3.0 / 4.0) * (                                     
   #        np.sqrt(r[k+1]**2 - r[i]**2)                           
   #        - np.sqrt(r[k]**2 - r[i]**2)                           
   #    ) / np.sqrt(2.0 * r[i] * (r[i+1] - r[i]))                  
   #Ne[i] = (3.0 / 4.0) * tec[i] / np.sqrt(2*r[i]*(r[i+1]-r[i])) \ 
   #        - sum(Cki(r, i, k) * Ne[k] for k in range(i+1, n))     

# TEC square spline interpolation to regular p-grid
def spline_tec_squared(p_obs, tec_obs, p_grid, s=0.0):
    """
    Spline interpolation of TEC^2 onto regular impact parameter grid.
    """

    tec2 = tec_obs**2
    spline = UnivariateSpline(p_obs, tec2, s=s, k=3)
    tec2_grid = spline(p_grid)

    # Numerical safety
    tec2_grid = np.maximum(tec2_grid, 0.0)

    return np.sqrt(tec2_grid)

# Linear regression for Ne(pmax) (Mode-1)
# Uses uppermost few kilometers (top_km can be tuned).

def estimate_Ne_pmax_mode1(p, tec, top_km=5.0):
    """
    Estimate Ne(pmax) from linear regression of TEC^2.
    """

    pmax = p[-1]

    # Select uppermost few km
    mask = (pmax - p) <= top_km * 1e3
    mask &= (p < pmax)

    dp = pmax - p[mask]
    tec2 = tec[mask]**2

    if np.count_nonzero(mask) < 3:
        raise RuntimeError("Insufficient points for Mode-1 regression")

    # Linear regression: TEC^2 = A * (pmax - p)
    A, _ = np.polyfit(dp, tec2, 1)

    Ne_pmax = np.sqrt(A / (8.0 * pmax))

    return Ne_pmax
    
# Full Mode-1 calibration procedure
def mode1_calibration(
    p_obs,
    tec_obs,
    r_grid,
    n_iter=10,
    top_km=5.0
):
    """
    Full Mode-1 calibration following gmrion description.

    Parameters
    ----------
    p_obs : ndarray
        Observed impact parameters (occultation side)
    tec_obs : ndarray
        Observed quasi-calibrated TEC
    r_grid : ndarray
        Regular radial grid (n+1)
    n_iter : int
        Number of iterations
    top_km : float
        Height range for regression near orbit altitude

    Returns
    -------
    Ne : ndarray
        Electron density profile
    Ne_pmax : float
        Electron density at orbit altitude
    tec : ndarray
        Final calibrated TEC
    """

    p_grid = r_grid[:-1]
    pmax = p_grid[-1]

    # --- Step 1: spline TEC^2 to regular grid ---
    tec = spline_tec_squared(p_obs, tec_obs, p_grid)

    # --- Iterative refinement ---
    for it in range(n_iter):

        # Step 2: estimate Ne(pmax)
        Ne_pmax = estimate_Ne_pmax_mode1(p_grid, tec, top_km)

        # Step 3: invert TEC ? Ne profile
        Ne = onion_peeling_linear(tec, r_grid)

        # Step 4: enforce Ne(pmax)
        Ne[-1] = Ne_pmax

    return Ne, Ne_pmax, tec