import numpy as np
'''
Functions for eta interpolation
    1. build_xy_lut_from_hist(H)
    2. bilinear_xy_lookup(etaX, etaY, X_tab, Y_tab, x_centers, y_centers)
    3. map_xy_no_interp(etaX, etaY, X_tab, Y_tab, x_edges, y_edges)
'''
def build_xy_lut_Rosenblatt(H):
    """
    Generate two 2D tables X_tab, Y_tab from a 2D histogram H
    such that (X_tab[i,j], Y_tab[i,j]) is the (x,y) value at the center of bin (i,j).
    Input:
        H: 2D histogram of charge centroid (eta_x, eta_y) distribution, shape (nBinX, nBinY), ranging from 0 to 1.
    Output:
        X_tab, Y_tab: two 2D tables of interpolated (x,y) values in position space, shape (nBinX, nBinY), ranging from 0 to 1.
    """
    H = np.asarray(H, dtype=float)
    nBinX, nBinY = H.shape
    etaX_edges = np.linspace(0, 1, nBinX+1)
    etaY_edges = np.linspace(0, 1, nBinY+1)

    # slight smoothing to avoid zero-count rows/columns causing flat CDF
    eps=1e-9
    alpha = max(eps, 1e-12 * max(1.0, H.sum()))
    Hs = H + alpha

    # marginal 1D CDF F_EtaX
    pEtaX = Hs.sum(axis=1)                                  # Marginal PDF (nBinX,)
    F_EtaX_edges = np.concatenate([[0.0], np.cumsum(pEtaX)])# CDF at bin edges (nBinX+1,)
    F_EtaX_edges = F_EtaX_edges / F_EtaX_edges[-1]                   

    # conditional CDF F_{Y|X} (a CDF in y direction for each etaX bin)
    F_EtaY_cond_edges = np.zeros((nBinX, nBinY + 1), float)
    for i in range(nBinX):
        cy = np.concatenate([[0.0], np.cumsum(Hs[i, :])])
        F_EtaY_cond_edges[i, :] = cy / cy[-1]

    # calculate (u,v) at the center of each original bin
    Xc_1d = 0.5 * (F_EtaX_edges[:-1] + F_EtaX_edges[1:])
    Yc_2d = 0.5 * (F_EtaY_cond_edges[:, :-1] + F_EtaY_cond_edges[:, 1:])

    X_tab = np.repeat(Xc_1d[:, None], nBinY, axis=1)        # (nBinX, nBinY)
    Y_tab = Yc_2d                                           # (nBinX, nBinY)

    return (X_tab, Y_tab)

def build_xy_lut_DoubleCDF(H):
    """
    Generate two 2D tables X_tab, Y_tab from a 2D histogram H
    such that (X_tab[i,j], Y_tab[i,j]) is the (x,y) value at the center of bin (i,j).
    Input:
        H: 2D histogram of charge centroid (eta_x, eta_y) distribution, shape (nBinX, nBinY), ranging from 0 to 1.
    Output:
        X_tab, Y_tab: two 2D tables of interpolated (x,y) values in position space, shape (nBinX, nBinY), ranging from 0 to 1.
    """
    H = np.asarray(H, dtype=float)
    nBinX, nBinY = H.shape
    etaX_edges = np.linspace(0, 1, nBinX+1)
    etaY_edges = np.linspace(0, 1, nBinY+1)

    # slight smoothing to avoid zero-count rows/columns causing flat CDF
    eps=1e-9
    alpha = max(eps, 1e-12 * max(1.0, H.sum()))
    Hs = H + alpha

    # marginal 1D CDF F_EtaX
    pEtaX = Hs.sum(axis=1)                                  # Marginal PDF (nBinX,)
    F_EtaX_edges = np.concatenate([[0.0], np.cumsum(pEtaX)])# CDF at bin edges (nBinX+1,)
    F_EtaX_edges = F_EtaX_edges / F_EtaX_edges[-1]                   

    # marginal 1D CDF F_EtaY
    pEtaY = Hs.sum(axis=0)                                  # Marginal PDF (nBinY,)
    F_EtaY_edges = np.concatenate([[0.0], np.cumsum(pEtaY)])# CDF at bin edges (nBinY+1,)
    F_EtaY_edges = F_EtaY_edges / F_EtaY_edges[-1]                   

    # calculate (u,v) at the center of each original bin
    Xc_1d = 0.5 * (F_EtaX_edges[:-1] + F_EtaX_edges[1:])
    Yc_1d = 0.5 * (F_EtaY_edges[:-1] + F_EtaY_edges[1:])

    X_tab = np.repeat(Xc_1d[:, None], nBinY, axis=1)        # (nBinX, nBinY)
    Y_tab = np.repeat(Yc_1d[None, :], nBinX, axis=0)        # (nBinX, nBinY)

    return (X_tab, Y_tab)

def bilinear_xy_lookup(etaX, etaY, X_tab, Y_tab):
    """
    bilinear interpolation, from (eta_x, eta_y) to (u,v):
    """
    x = np.asarray(etaX); y = np.asarray(etaY)
    nBinX, nBinY = X_tab.shape

    ### bin centers
    x_centers = (np.arange(nBinX) + 0.5) / nBinX
    y_centers = (np.arange(nBinY) + 0.5) / nBinY

    # search bin
    ix = np.clip(np.searchsorted(x_centers, x) - 1, 0, nBinX-2)
    iy = np.clip(np.searchsorted(y_centers, y) - 1, 0, nBinY-2)

    # four corners
    x0 = x_centers[ix];     x1 = x_centers[ix+1]
    y0 = y_centers[iy];     y1 = y_centers[iy+1]
    tx = np.where(x1>x0, (x - x0)/(x1 - x0), 0.0)
    ty = np.where(y1>y0, (y - y0)/(y1 - y0), 0.0)

    # values at corners
    X00 = X_tab[ix,   iy  ]
    X10 = X_tab[ix+1, iy  ]
    X01 = X_tab[ix,   iy+1]
    X11 = X_tab[ix+1, iy+1]
    Y00 = Y_tab[ix,   iy  ]
    Y10 = Y_tab[ix+1, iy  ]
    Y01 = Y_tab[ix,   iy+1]
    Y11 = Y_tab[ix+1, iy+1]

    # bilinear interpolation
    X0 = (1-tx)*X00 + tx*X10
    X1 = (1-tx)*X01 + tx*X11
    Y0 = (1-tx)*Y00 + tx*Y10
    Y1 = (1-tx)*Y01 + tx*Y11
    X = (1-ty)*X0 + ty*X1
    Y = (1-ty)*Y0 + ty*Y1
    return X, Y

def xy_lookup(etaX, etaY, X_tab, Y_tab):
    """
    mapping without interpolation.
    """
    x = np.asarray(etaX)
    y = np.asarray(etaY)
    nBinX, nBinY = X_tab.shape

    # Compute nearest bin index directly
    ix = np.clip(np.rint(x * nBinX - 0.5).astype(int), 0, nBinX - 1)
    iy = np.clip(np.rint(y * nBinY - 0.5).astype(int), 0, nBinY - 1)

    return X_tab[ix, iy], Y_tab[ix, iy]