"""
    Function supporting monotonic cubic polynomials: 
        first parametrization
    We define the polynomial:
          poly(x) = sum_i beta_i * x^i
    where the coefficients beta_i are parameterized by the vector 'pars':
        beta_i is a function pars
"""

import numpy as np

def forward_map(pars, small=1e-8):
    """
    Map parameter vector pars to coefficients beta.

    beta = [p0, (p2^2 + p1^2) / a, p2, p3^2],
    where a = 3*p3^2 + small.
    """
    pars = np.asarray(pars, dtype=float)
    p0, p1, p2, p3 = pars
    a = 3.0 * p3**2 + small
    beta = np.array([
        p0,
        (p2**2 + p1**2) / a,
        p2,
        p3**2
    ], dtype=float)
    return beta

def backward_map(beta, small=1e-8):
    """
    Map coefficients beta back to parameter vector pars.
    Nonnegative roots are used for p1 and p3.
    """
    beta = np.asarray(beta, dtype=float)
    b0, b1, b2, b3 = beta
    p0 = b0
    p3 = np.sqrt(b3)
    a = 3.0 * b3 + small
    p2 = b2
    p1_sq = b1 * a - p2**2
    if p1_sq < 0:
        raise ValueError(f"Negative square root encountered for p1^2 = {p1_sq}")
    p1 = np.sqrt(p1_sq)
    return np.array([p0, p1, p2, p3], dtype=float)

def forward_map_jacobian(pars, small=1e-8):
    """
    Compute Jacobian of forward_map at given pars.
    Returns J: d(beta)/d(pars), shape (4, 4).
    """
    p0, p1, p2, p3 = np.asarray(pars, dtype=float)
    a = 3.0 * p3**2 + small
    J = np.zeros((4, 4), dtype=float)

    # beta0 row
    J[0, 0] = 1.0

    # beta1 row
    J[1, 1] = 2.0 * p1 / a
    J[1, 2] = 2.0 * p2 / a
    J[1, 3] = -6.0 * p3 * (p2**2 + p1**2) / a**2

    # beta2 row
    J[2, 2] = 1.0

    # beta3 row
    J[3, 3] = 2.0 * p3

    return J

def forward_map_hessian(pars, small=1e-8):
    """
    Compute Hessians of forward_map at given pars.

    Returns H: shape (4, 4, 4),
    where H[i] is the 4x4 Hessian of beta[i] w.r.t. pars.
    """
    p0, p1, p2, p3 = np.asarray(pars, dtype=float)
    a = 3.0 * p3**2 + small
    N = p1**2 + p2**2
    H = np.zeros((4, 4, 4), dtype=float)

    # beta0: all zeros

    # beta1 Hessian
    H1 = np.zeros((4, 4), dtype=float)
    H1[1, 1] = 2.0 / a
    H1[2, 2] = 2.0 / a
    H1[1, 3] = -12.0 * p1 * p3 / a**2
    H1[3, 1] = H1[1, 3]
    H1[2, 3] = -12.0 * p2 * p3 / a**2
    H1[3, 2] = H1[2, 3]
    H1[3, 3] = 6.0 * N * (12.0 * p3**2 - a) / a**3
    H[1] = H1

    # beta2: all zeros

    # beta3 Hessian
    H3 = np.zeros((4, 4), dtype=float)
    H3[3, 3] = 2.0
    H[3] = H3

    return H


# -------------------------
# Round-trip test
# -------------------------
if __name__ == "__main__":
    pars_original = np.array([1.0, 2.0, 3.0, 4.0])
    beta = forward_map(pars_original)
    pars_recovered = backward_map(beta)

    print("Original pars: ", pars_original)
    print("Beta:          ", beta)
    print("Recovered pars:", pars_recovered)
    print("Difference:    ", pars_recovered - pars_original)

    print("\nJacobian at pars:")
    print(forward_map_jacobian(pars_original))

    print("\nHessian for beta1:")
    print(forward_map_hessian(pars_original)[1])