"""
Generalized functions used for the formulation of the problem. The notation generally follows that used in the
assignment: https://byu.instructure.com/courses/14234/assignments/545719
"""
import numpy as np

fx = 1
x_min, x_max = 0, 1
x_range = x_max - x_min


def formulate(num_points, gamma):
    """Formulate the A matrix and b vector"""
    A = A_matrix(num_points, gamma)
    b = b_vector(num_points)
    return A, b


def A_matrix(num_points, gamma):
    """Generate the A matrix"""
    dx = x_range / num_points
    A = np.zeros((num_points, num_points))
    for i, row in enumerate(A):
        # Create the diagonal terms
        row[i] = 2 / dx
        if i != 0:
            # Create the lower-diagonal terms
            row[i - 1] = -(1 / dx) - (gamma / 2)
        if i != num_points - 1:
            # Create the upper-diagonal terms
            row[i + 1] = -(1 / dx) + (gamma / 2)
            pass
    return A


def A1_inv(num_points):
    """Generate the A matrix"""
    dx = x_range / num_points
    A = np.zeros((num_points, num_points))
    for i, row in enumerate(A):
        # Create the diagonal terms
        row[i] = 2 / dx
        if i != 0:
            # Create the lower-diagonal terms
            row[i - 1] = -(1 / dx)
        if i != num_points - 1:
            # Create the upper-diagonal terms
            row[i + 1] = -(1 / dx)
            pass
    return np.linalg.inv(A)


def b_vector(num_points):
    """Generate the b vector"""
    dx = x_range / num_points
    b = np.ones((num_points, 1)) * dx
    return b
