""" 
    Function supporting monotonic cubic polynomials: 
        second parametrization
    We define the polynomial:
          poly(x) = sum_i beta_i * x^i
    where the coefficients beta_i are parameterized by the vector 'pars'.
"""
import numpy as np

# Forward map: Parameters to polynomial coefficients
def forward_map(pars):
    """
    Converts parameters (pars = [C, epsilon, k1, k2]) into polynomial coefficients (beta = [d, c, b, a]).
    
    Polynomial Definition:
    - d = C: Constant term of the polynomial.
    - c = k2^2 + epsilon^2: Coefficient of the linear term.
    - b = k1 * k2: Coefficient of the quadratic term.
    - a = k1**2 / 3: Coefficient of the cubic term.

    Parameters:
    - pars: NumPy array of parameters [C, epsilon, k1, k2].

    Returns:
    - beta: NumPy array of polynomial coefficients [d, c, b, a].
    """
    C, epsilon, k1, k2 = pars  # Unpack the parameter vector
    
    # Compute coefficients
    d = C
    c = k2**2 + epsilon**2
    b = k1 * k2
    a = k1**2 / 3
    
    # Return coefficients as a NumPy array
    beta = np.array([d, c, b, a])
    return beta


# Jacobian of the forward map: First-order derivatives
def forward_map_jacobian(pars):
    """
    Computes the Jacobian matrix of the forward map analytically.
    
    Parameters:
    - pars: NumPy array of parameters [C, epsilon, k1, k2].
    
    Returns:
    - J: NumPy 4x4 Jacobian matrix, where J[i, j] = d(beta[i])/d(pars[j]).
    """
    C, epsilon, k1, k2 = pars  # Unpack parameters
    
    # Initialize Jacobian matrix
    J = np.zeros((4, 4))  # 4x4 matrix
    
    # Partial derivatives for d = C
    J[0, 0] = 1  # d(d)/dC
    J[0, 1] = 0  # d(d)/d(epsilon)
    J[0, 2] = 0  # d(d)/d(k1)
    J[0, 3] = 0  # d(d)/d(k2)

    # Partial derivatives for c = k2^2 + epsilon^2
    J[1, 0] = 0  # d(c)/dC
    J[1, 1] = 2 * epsilon  # d(c)/d(epsilon)
    J[1, 2] = 0  # d(c)/d(k1)
    J[1, 3] = 2 * k2  # d(c)/d(k2)

    # Partial derivatives for b = k1 * k2
    J[2, 0] = 0  # d(b)/dC
    J[2, 1] = 0  # d(b)/d(epsilon)
    J[2, 2] = k2  # d(b)/d(k1)
    J[2, 3] = k1  # d(b)/d(k2)

    # Partial derivatives for a = k1^2 / 3
    J[3, 0] = 0  # d(a)/dC
    J[3, 1] = 0  # d(a)/d(epsilon)
    J[3, 2] = 2 * k1 / 3  # d(a)/d(k1)
    J[3, 3] = 0  # d(a)/d(k2)
    
    return J


# Hessian of the forward map: Second-order derivatives
def forward_map_hessian(pars):
    """
    Computes the Hessian tensor of the forward map analytically.
    
    Parameters:
    - pars: NumPy array of parameters [C, epsilon, k1, k2].
    
    Returns:
    - H: NumPy 4x4x4 Hessian tensor, where H[i, j, k] = d^2(beta[i])/d(pars[j])d(pars[k]).
    """
    C, epsilon, k1, k2 = pars  # Unpack parameters
    
    # Initialize Hessian tensor (4 x 4 x 4)
    H = np.zeros((4, 4, 4))
    
    # Hessian for d = C: All second derivatives are zero
    # Already H[0, :, :] is initialized to zero
    
    # Hessian for c = k2^2 + epsilon^2
    H[1, 1, 1] = 2  # d^2(c)/d(epsilon^2)
    H[1, 3, 3] = 2  # d^2(c)/d(k2^2)
    
    # Hessian for b = k1 * k2: All second derivatives are zero
    # Already H[2, :, :] is initialized to zero
    
    # Hessian for a = k1^2 / 3
    H[3, 2, 2] = 2 / 3  # d^2(a)/d(k1^2)
    
    return H


# Backward map: Polynomial coefficients to parameters
def backward_map(beta, only_one = True):
    """
    Computes the parameters (pars = [C, epsilon, k1, k2]) from the polynomial coefficients (beta = [d, c, b, a]).
    
    Polynomial Definition:
    - d = C: Constant term of the polynomial.
    - c = k2^2 + epsilon^2: Used to recover k2 and epsilon.
    - b = k1 * k2: Used to recover k1 and k2.
    - a = k1^2 / 3: Used to recover k1.

    Parameters:
    - beta: NumPy array of polynomial coefficients [d, c, b, a].

    Returns:
    - List of possible parameter sets [(C, epsilon, k1, k2)].
    """
    d, c, b, a = beta  # Unpack the coefficients
    
    # Recover k1 from a (two possible values due to ± sqrt)
    if a < 0:
        raise ValueError("Coefficient 'a' must be non-negative for monotonic polynomials.")
    
    k1_options = np.unique([np.sqrt(3 * a), -np.sqrt(3 * a)])  # Two possible k1 values
    
    possible_parameters = []
    
    # For each possible k1, compute k2 and epsilon
    for k1 in k1_options:
        k2 = None
        if k1 != 0:  # Ensure k1 is non-zero (avoids division by zero)
            k2 = b / k1  # Compute k2 from b and k1
        elif b == 0:
            k2 = 0

        if k2 is None: continue    
       
        # Check if c >= k2^2 for valid epsilon computation
        if c >= k2**2:
            epsilon_options = np.unique([np.sqrt(c - k2**2), -np.sqrt(c - k2**2)])  # Two possible epsilon values
            
            for epsilon in epsilon_options:
                # Constant term d maps directly to C
                C = d
                sol = np.array([C, epsilon, k1, k2])

                if only_one: return sol
                possible_parameters.append(sol)


    return possible_parameters


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

    print("\nLinear func:")
    lin_fun_beta = [1,0.2,0,0]
    only_one = False

    lin_fun_pars = backward_map(lin_fun_beta, only_one=only_one)
    lin_fun_beta_recover = lin_fun_pars if only_one else np.unique([forward_map(pars) for pars in lin_fun_pars], axis=0)
    
    print(f"  {only_one = }")
    print("  lin_fun_beta:", lin_fun_beta)
    print("  backwards:", lin_fun_pars)
    print("  forwards:", lin_fun_beta_recover)