import numpy as np
import scipy.integrate as integ

# Data generation routines
def calculate_strain_stress(
    input_type, time_array, input_expr, E_mods, viscs, D_input_lambda=None
):
    # In this incarnation, the kwarg is non-optional.

    #     if D_input_lambda:
    input_lambda = input_expr
    #     else:
    #         t = sym.symbols('t', real=True)
    #         D_input_expr = input_expr.diff(t)

    #         input_lambda = sym.lambdify(t, input_expr)
    #         D_input_lambda = sym.lambdify(t, D_input_expr)

    # The following function interprets the provided model parameters differently depending on the input_type.
    # If the input_type is 'Strain' then the parameters are assumed to refer to a Maxwell model, whereas
    # if the input_type is 'Stress' then the parameters are assumed to refer to a Kelvin model.
    relax_creep_lambda = relax_creep(E_mods, viscs, input_type)

    if relax_creep_lambda == False:
        return False, False

    start_time_point = time_array[0]

    integrand_lambda = lambda x, t: relax_creep_lambda(t - x) * D_input_lambda(x)
    integral_lambda = lambda t: integ.quad(
        integrand_lambda, start_time_point, t, args=(t)
    )[0]

    output_array = np.array([])
    input_array = np.array([])
    for time_point in time_array:
        first_term = input_lambda(start_time_point) * relax_creep_lambda(
            time_point - start_time_point
        )
        second_term = integral_lambda(time_point)
        output_array = np.append(output_array, first_term + second_term)
        input_array = np.append(input_array, input_lambda(time_point))

    if input_type == "Strain":
        strain_array = input_array
        stress_array = output_array
    else:
        strain_array = output_array
        stress_array = input_array

    strain_array = strain_array.reshape(time_array.shape)
    stress_array = stress_array.reshape(time_array.shape)

    return strain_array, stress_array


def relax_creep(E_mods, viscs, input_type):

    # The following function interprets the provided model parameters differently depending on the input_type.
    # If the input_type is 'Strain' then the parameters are assumed to refer to a Maxwell model, whereas
    # if the input_type is 'Stress' then the parameters are assumed to refer to a Kelvin model.
    # The equations used thus allow the data to be generated according to the model now designated.

    E_mods_1plus_array = np.array(E_mods[1:]).reshape(-1, 1)
    viscs_array = np.array(viscs).reshape(-1, 1)

    taus = viscs_array / E_mods_1plus_array

    if input_type == "Strain":
        relax_creep_lambda = lambda t: E_mods[0] + np.sum(
            np.exp(-t / taus) * E_mods_1plus_array
        )
    elif input_type == "Stress":
        relax_creep_lambda = lambda t: 1 / E_mods[0] + np.sum(
            (1 - np.exp(-t / taus)) / E_mods_1plus_array
        )
    else:
        print("Incorrect input_type")
        relax_creep_lambda = False

    return relax_creep_lambda
