"""
This file provides functions for the log-periodic power law function that the JLS model uses, along
with a symbolic gradient for use in curve-fitting and nonlinear optimization.
"""
import numpy as np


def lppl_func_from_params(log_p_c, b_0, b_1, kappa, beta, omega, phi, t_c):
    """Given parameters, outputs a numpy vectorized function that computes the log-periodic power
    law. Note that putting in t_c to this function will return NaN!"""
    def f(x):
        delta_t = np.subtract(t_c, x)
        exp_delta_t = np.power(delta_t, beta)
        t1 = np.multiply(b_0, exp_delta_t)
        cos_arg = np.cos(np.add(np.multiply(omega, np.log(delta_t)), phi))
        t2 = np.multiply(b_1, np.multiply(exp_delta_t, cos_arg))
        return np.subtract(log_p_c, np.multiply(np.divide(kappa, beta), np.add(t1, t2)))
    return f


def transformed_lppl_from_params(A, B, C_1, C_2, t_c, m, w):
    """Returns a numpy vectorized function for the transformed version of the LPPL."""
    # for reference, we're computing the following:
    # f(t) = A + B(t_c-t)^m + C_1(t_c-t)^m cos(w ln(t_c-t)) + C_2(t_c-t)^m sin(w ln(t_c-t))
    # don't say I didn't warn you

    # this is different enough to the above that it makes sense to just do it again
    def f(x):
        delta_t = np.subtract(t_c, x)
        exp_delta_t = np.power(delta_t, m)
        angle_term = np.multiply(w, np.log(delta_t))
        trig_term1 = np.multiply(C_1, np.multiply(exp_delta_t, np.cos(angle_term)))
        trig_term2 = np.multiply(C_1, np.multiply(exp_delta_t, np.sin(angle_term)))
        return np.add(A, np.add(np.multiply(B, exp_delta_t), trig_term1, trig_term2))
    return f


def lppl(t, log_p_c, b_0, b_1, kappa, beta, omega, phi, t_c):
    """Just calculates the LPPL for the given parameters."""
    return lppl_func_from_params(log_p_c, b_0, b_1, kappa, beta, omega, phi, t_c)(t)


def transformed_lppl(t, A, B, C_1, C_2, t_c, m, w):
    """Calculates the LPPL for the given parameters, using a slightly different transformed version
    of the LPPL."""
    return transformed_lppl_from_params(A, B, C_1, C_2, t_c, m, w)(t)
