import sympy as sp

def centered_difference_coefficients_1D(
        field_name = 'F',
        number_of_neighbours = 1):
    """
        return all centered differences that can be computed
        with a certain number of neighbours.
    """
    # variable
    x = sp.Symbol('x')
    h = sp.Symbol('h')
    #Generate "actual values" of function on grid points
    F = []
    F.append(sp.Symbol(field_name + '[i]'))
    for i in range(1, number_of_neighbours + 1):
        F.append(sp.Symbol(field_name + '[i{0:+}]'.format(i)))
    for i in range(-number_of_neighbours, 0):
        F.append(sp.Symbol(field_name + '[i{0:+}]'.format(i)))
    #Taylor expansion:
    f = [sp.Symbol('f^{(0)}')]
    tef = f[0]
    for i in range(1, number_of_neighbours*2 + 1):
        f.append(sp.Symbol('f^{{({0})}}'.format(i)))
        tef += f[-1] * x**i / sp.factorial(i)
    eq = [tef.subs(x, 0) - F[0]]
    for i in range(1, number_of_neighbours + 1):
        eq += [tef.subs(x, -i) - F[-i],
               tef.subs(x,  i) - F[ i]]
    sol = sp.solve(eq, f)
    cd = []
    coeff = [] 
    for i in range(len(f)):
        cd.append(sp.factor(f[i].subs(sol)))
        coeff.append([])
        for j in range(-number_of_neighbours, number_of_neighbours+1):
            coeff[i].append((f[i].subs(sol)).diff(F[j]))
    return F, cd, coeff

def centered_difference_coefficient_ND(
        degrees = [1, 1],
        grid_point = [0, 0],
        coeff_matrix = None):
    number_of_neighbours = (coeff_matrix.shape[0]-1)/2
    coeff = 1.0
    for i in range(len(degrees)):
        coeff *= coeff_matrix[degrees[i]][grid_point[i]+number_of_neighbours]
    return coeff

def main():
    out_file = open('centered_differences.tex', 'w')
    out_file.write('\\documentclass{article}\n'
                 + '\\usepackage{amsmath}\n'
                 + '\\begin{document}\n')
    out_file.write('Generally any centered difference of order $k$ (computed with neighbours of degree $n$) is a linear combination of values of the field $F$:\n'
                 + '\\begin{equation}\n'
                 + 'F^{(k)}_n [i] \\approx \\frac{1}{h^k}\\sum_{j = -n}^n c^n_{kj} F[i+j]\n'
                 + '\\end{equation}\n')
    out_file.write('A tensor product has to be used for the 3D case:\n'
                 + '\\begin{align}\n'
                 + 'F^{(k_x, k_y, k_z)}_n [i_x, i_y, i_z] &\\approx \\\\ '
                    + '\\frac{1}{h_x^{k_x}h_y^{k_y}h_z^{k_z}} '
                    + '&\\sum_{j_x = -n}^n \\sum_{j_y = -n}^n \\sum_{j_z = -n}^n '
                    + 'c^n_{k_xj_x} c^n_{k_yj_y} c^n_{k_zj_z} F[i_x+j_x, i_y + j_y, i_z + j_z]\n'
                 + '\\end{align}\n'
                 + '(well, the same is true for any number of dimensions).')
    for neighbours in [1, 2]:
        out_file.write('\\paragraph{{{0} neighbours}}'.format(neighbours))
        F, cd, coeff = centered_difference_coefficients_1D(number_of_neighbours = neighbours)
        for n in range(len(cd)):
            out_file.write(
                    '\\begin{equation}\n'
                  + 'F^{{({0})}}[i] \\approx \\frac{{1}}{{h^{0}}} \\left('.format(n) + sp.latex(cd[n]) + '\\right)\n'
                  + '\\end{equation}\n')
        out_file.write('and now just coefficient matrix:')
        out_file.write('\\begin{equation}\n')
        out_file.write('c^{{{0}}} = \\left(\\begin{{matrix}}\n'.format(neighbours))
        for k in range(len(coeff)):
            for j in range(len(coeff[k])-1):
                out_file.write(sp.latex(coeff[k][j]) + ' & ')
            out_file.write(sp.latex(coeff[k][-1]))
            if k < len(coeff) - 1:
                out_file.write(' \\\\[7pt] ')
            out_file.write('\n')
        out_file.write('\\end{matrix}\\right)\n')
        out_file.write('\\end{equation}\n')
    for neighbours in [3, 4]:
        out_file.write('\\paragraph{{{0} neighbours}}'.format(neighbours))
        F, cd, coeff = centered_difference_coefficients_1D(number_of_neighbours = neighbours)
        out_file.write('\\begin{equation}\n')
        out_file.write('c^{{{0}}} = \\left(\\begin{{matrix}}\n'.format(neighbours))
        for k in range(len(coeff)):
            for j in range(len(coeff[k])-1):
                out_file.write(sp.latex(coeff[k][j]) + ' & ')
            out_file.write(sp.latex(coeff[k][-1]))
            if k < len(coeff) - 1:
                out_file.write(' \\\\[7pt] ')
            out_file.write('\n')
        out_file.write('\\end{matrix}\\right)\n')
        out_file.write('\\end{equation}\n')
    out_file.write('\\end{document}\n')
    out_file.close()
    return None

if __name__ == '__main__':
    main()

