import numpy as np
import scipy.sparse as scs
import scipy.sparse.linalg as scsl
import matplotlib.pyplot as plt


# Defines the range for t
TSTART = 0
TEND = 80

# Number of points for plotting
NDP = 1000

# Valid range for N
NMIN = 10
NMAX = 400

# The mesh (initialized later)
N = 0
mesh = 0


# Create mesh

def create_mesh(n):
    return np.linspace(TSTART, TEND, n)


# Compute basis function fk in point t
def fk(t, k):

    if k == 0:
        if t < mesh[1]:
            return (mesh[1] - t) / (mesh[1] - mesh[0])
        else:
            return 0

    if k == N - 1:
        if t > mesh[N - 2]:
            return (t - mesh[N - 2]) / (mesh[N - 1] - mesh[N - 2])
        else:
            return 0

    if t < mesh[k - 1] or t > mesh[k + 1]:
        return 0

    if mesh[k - 1] <= t <= mesh[k]:
        return (t - mesh[k - 1]) / (mesh[k] - mesh[k - 1])

    return (mesh[k + 1] - t) / (mesh[k + 1] - mesh[k])


# Compute the right-hand side (vector b) analytically

def compute_integral_l(a, b):
    return (1. / (a - b)) * ((b - a) * np.cos(b) + np.sin(a) - np.sin(b))


def compute_integral_r(a, b):
    return (1. / (a - b)) * ((a - b) * np.cos(a) - np.sin(a) + np.sin(b))


def bk(k):
    first = 0
    second = 0

    tk = mesh[k]

    if (k > 0):
        tkm1 = mesh[k - 1]
        first = compute_integral_l(tkm1, tk)

    if (k < N - 1):
        tkp1 = mesh[k + 1]
        second = compute_integral_r(tk, tkp1)

    return first + second


def compute_b():
    b = np.ndarray((N, 1))
    for i in range(N):
        b[i] = bk(i)
    return b


# Compute the matrix A analytically

def aki(k, i):
    if (k == i - 1):
        return 0.5
    if (k == i + 1):
        return -0.5
    if (k == i and i == 0):
        return -0.5
    if (k == i and i == N - 1):
        return 0.5
    return 0.0


def compute_a():
    a = scs.lil_matrix((N, N))
    for i in range(N):
        for j in range(N):
            a[i, j] = aki(i, j)
    return a


# Apply boundary condition

def apply_bc(a, b, y0):
    b[0] = y0
    a[0, 0] = 1.0
    a[0, 1:] = np.zeros((1, N - 1))


# Compute true solution

def true_solution(t):
    return -np.cos(t) + 2


# Compute the approximate solution after solving the system

def approx_solution(t, alpha):
    res = 0
    for i in range(N):
        res += alpha[i] * fk(t, i)
    return res


# Plot function

def pplot(X, Y, title):
    X = np.asarray(X)
    Y = np.asarray(Y)
    plt.plot(X, Y, label=title)
    plt.xlabel('t')
    plt.ylabel('f(t)')


#	====
#	MAIN
#	====


print('==============================================================')
print('- This program finds the solution to the differential equation')
print('- It looks for an approximate (piecewise-linear) solution')
print('- You need to determine the optimal number of linear segments')
print('- If this number is too small, precision is not enough')
print('- If this number is too big, the computation time is long')


# Repeat forever until user exits
while(1 > 0):

    N = NMIN

    # Read input and process it
    tmp = input(
        ' ====> Enter number of segments in a range from ' +
        str(NMIN) +
        ' to ' +
        str(NMAX) +
        ' (0 to exit): ')
    try:
        N = int(tmp)
    except:
        print('       The input is incorrect')
        continue

    if (N == 0):
        break

    if (N < NMIN or N > NMAX):
        print('       From ' + str(NMIN) + ' to ' + str(NMAX))
        continue

        # Define mesh and matrices
    mesh = create_mesh(N)
    a = compute_a()
    b = compute_b()

    # Apply boundary condition f(0) = 1 ( = y0 )
    y0 = 1
    apply_bc(a, b, y0)

    # Solve the system of linear equations
    a = a.tocsr()
    alpha = scsl.spsolve(a, b)

    xx = np.linspace(TSTART, TEND, NDP)
    yy = np.ndarray((NDP, 1))
    zz = np.ndarray((NDP, 1))
    for i in range(NDP):
        zz[i] = true_solution(xx[i])

    for i in range(NDP):
        yy[i] = approx_solution(xx[i], alpha)

    pplot(xx, zz, 'True solution')
    pplot(xx, yy, 'Approximate solution with ' + str(N) + ' segments')

    plt.title('Solutions to the equation f\'(t) = sin(t), f(0) = 1')

    plt.xlim(np.min(xx), np.max(xx))
    plt.ylim(np.min([np.min(yy), np.min(zz)]) - 2,
             np.max([np.max(yy), np.max(zz)]) + 2)
    plt.legend()
    plt.show()

# All done
print('Bye!')