import sympy as sym
import numpy as np
import mpmath as mp
import matplotlib.pyplot as plt

x = sym.Symbol('x')

def least_squares(f, psi, Omega, symbolic=True):
    """
    Given a function f(x) on an interval Omega (2-list)
    return the best approximation to f(x) in the space V
    spanned by the functions in the list psi.
    """
    N = len(psi) - 1
    A = sym.zeros(N+1, N+1)
    b = sym.zeros(N+1, 1)
    x = sym.Symbol('x')
    print ('...evaluating matrix...',)
    for i in range(N+1):
        for j in range(i, N+1):
            print ('(%d,%d)' % (i, j))

            integrand = psi[i]*psi[j]
            if symbolic:
                I = sym.integrate(integrand, (x, Omega[0], Omega[1]))
            if not symbolic or isinstance(I, sym.Integral):
                # Could not integrate symbolically, use numerical int.
                print ('numerical integration of', integrand)
                integrand = sym.lambdify([x], integrand)
                I = mp.quad(integrand, [Omega[0], Omega[1]])
            A[i,j] = A[j,i] = I
        integrand = psi[i]*f
        if symbolic:
            I = sym.integrate(integrand, (x, Omega[0], Omega[1]))
        if not symbolic or isinstance(I, sym.Integral):
            # Could not integrate symbolically, use numerical int.
            print ('numerical integration of', integrand)
            integrand = sym.lambdify([x], integrand)
            I = mp.quad(integrand, [Omega[0], Omega[1]])
        b[i,0] = I
    print ('A:\n', A, '\nb:\n', b)
    if symbolic:
        c = A.LUsolve(b)  # symbolic solve
        # c is a sympy Matrix object, numbers are in c[i,0]
        c = [sym.simplify(c[i,0]) for i in range(c.shape[0])]
    else:
        c = mp.lu_solve(A, b)  # numerical solve
    print ('coeff:', c)

    u = sum(c[i]*psi[i] for i in range(len(psi)))
    print ('approximation:', u)
    return u, c

def comparison_plot(f, u, Omega):
    # Turn f and u to ordinary Python functions
    f = sym.lambdify([x], f, modules="numpy")
    u = sym.lambdify([x], u, modules="numpy")
    resolution = 401  # no of points in plot
    xcoor  = np.linspace(Omega[0], Omega[1], resolution)
    exact  = f(xcoor)
    approx = u(xcoor)
    plt.plot(xcoor, approx)
    plt.plot(xcoor, exact)
    plt.legend(['approximation', 'exact'])

def run_parabola_by_linear_leastsq():
    f = 10*(x-1)**2 - 1
    psi = [1, x]
    Omega = [1, 2]
    u, c = least_squares(f, psi, Omega)
    comparison_plot(f, u, Omega)