import numpy as np
from scipy.special import erfc

# 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))     

 

# Scale-heigth estimation (uppermost 100 km) 

R_EARTH = 6371e3  # meters, this should be replaced by earth curvature r_roc

def estimate_scale_height(r, Ne, top_km=100.0):
    """
    Estimate ionospheric scale height and electron density at p_max
    using log-linear regression over the uppermost 100 km.

    Parameters
    ----------
    r : ndarray, shape (n+1,)
        Radial grid boundaries (meters)
    Ne : ndarray, shape (n,)
        Retrieved electron density profile (m^-3)
    top_km : float
        Thickness of top region for regression (km)

    Returns
    -------
    H : float
        Scale height (meters)
    Ne_pmax : float
        Electron density at highest tangent point p_max (m^-3)
    """

    # Tangent point heights
    h = r[:-1] - R_EARTH
    pmax = r[-2]  # highest tangent point
    hmax = h[-1]

    # Select uppermost top_km
    mask = (h >= hmax - top_km * 1e3) & (Ne > 0.0) & np.isfinite(Ne)

    if np.count_nonzero(mask) < 3:
        raise RuntimeError("Not enough valid points for scale height estimation")

    h_fit = h[mask]
    lnNe = np.log(Ne[mask])

    # Log-linear regression: ln(Ne) = a + b*h
    b, a = np.polyfit(h_fit, lnNe, 1)

    # Scale height
    H = -1.0 / b

    # Electron density at p_max
    h_pmax = pmax - R_EARTH
    Ne_pmax = np.exp(a + b * h_pmax)

    return H, Ne_pmax

#Mode-0 analysis topside TEC correction 
def mode0_topside_tec(p, pmax, Ne_pmax, H):
    """
    Mode-0 analytic topside TEC correction (invTEC / gmrion)

    Parameters
    ----------
    p : ndarray
        Impact parameters of rays (meters)
    pmax : float
        Maximum tangent point (meters)
    Ne_pmax : float
        Electron density at pmax (m^-3)
    H : float
        Scale height (meters)

    Returns
    -------
    dTEC : ndarray
        Topside TEC correction for each ray (m^-2)
    """

    x = (pmax - p) / H
    x = np.maximum(x, 0.0)  # numerical safety

    prefactor = Ne_pmax * np.sqrt(0.5 * np.pi * H * pmax)

    dTEC = prefactor * (
        1.0 - np.exp(x) * erfc(np.sqrt(x))
    )

    return dTEC

# Full Mode-0 calibration loop (10 iterations)
# First iteration assumes a very extended topside
# Subsequent iterations t0 pull the scale height downward
# Final solution satisfies:
#   TEC consistency
#   Exponential topside
#   Observed Ne gradient

def mode0_calibration(tec0, r, n_iter=10):
    """
    Mode-0 TEC calibration with initial H = 1000 km.
 
    Parameters
    ----------
    tec0 : ndarray
        Initial (quasi-calibrated) TEC
    r : ndarray
        Radial grid (n+1)
    n_iter : int
        Number of iterations (default: 10)

    Returns
    -------
    Ne : ndarray
        Final electron density profile
    H : float
        Final scale height
    Ne_pmax : float
        Final Ne at p_max
    tec : ndarray
        Final calibrated TEC
    """

    p = r[:-1]
    pmax = p[-1]

    # --- Initialization ---
    H = 1000e3  # 1000 km initial scale height
    tec = tec0.copy()

    # Initial inversion to get Ne and Ne(pmax)
    Ne = onion_peeling_linear(tec, r)
    Ne_pmax = Ne[-1]

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

        # Update TEC using current H and Ne(pmax)
        tec = tec0 + mode0_topside_tec(p, pmax, Ne_pmax, H)

        # Reconstruct Ne
        Ne = onion_peeling_linear(tec, r)

        # Update H and Ne(pmax) from Ne profile
        H, Ne_pmax = estimate_scale_height(r, Ne)

        # Physical constraint (optional but recommended)
        H = np.clip(H, 20e3, 150e3)

    return Ne, H, Ne_pmax, tec