import sympy as sp
import numpy as np
from time import time

ti = time()

xi, eta = sp.symbols(['xi', 'eta'])
f1 = 1/2 - 3/4*xi + 1/4*xi**3
f2 = 1/8 - 1/8*xi - 1/8*xi**2 + 1/8*xi**3
f3 = 1/2 + 3/4*xi - 1/4*xi**3
f4 = -1/8 - 1/8*xi + 1/8*xi**2 + 1/8*xi**3

g1 = 1/2 - 3/4*eta + 1/4*eta**3
g2 = 1/8 - 1/8*eta - 1/8*eta**2 + 1/8*eta**3
g3 = 1/2 + 3/4*eta - 1/4*eta**3
g4 = -1/8 - 1/8*eta + 1/8*eta**2 + 1/8*eta**3

FXI = {0: f1, 1: f2, 2: f3, 3: f4} # first terms of set of functions
GETA = {0: g1, 1: g2, 2: g3, 3: g4}


def _orth_polynomials(r, sym):   
    '''
    r term from the set of functions used by Bardell (1991)

    Parameters
    ----------
    r : int
    sym : sympy.symbol
        symbol to be used in equation

    Returns
    -------
    f_sym_r : sympy equation
    r term from the set of functions.
    '''
    nmax = r//2 + 1 #denotes its own integer part (Bardel, 1991).
    f_sym_r = 0

    for n in range(0, nmax):
        f_sym_r += (
            (-1)**n * sp.factorial2(2*r-2*n-7)
            / (2**n * sp.factorial(n) * sp.factorial(r-2*n-1))
            * sym**(r-2*n-1)
            )
    f_sym_r = sp.sympify(f_sym_r)
    return f_sym_r


def convert_to_str(d, name):
    '''
    Parameters
    ----------
    d : dict 
        dictionary to be converted into string
    name : str
        name of the variable
    '''

    txt = name + ' = {\n'
    for k, v in d.items():
        txt += '    ' + str(k) + ': \'' + str(v) + '\',\n'
    txt += '}' + '\n' * 3
    return txt


if __name__ == '__main__':
    
    n_terms = 50

    # Complete set of bardell functions
    for i in range(4, n_terms):
        FXI[i] = _orth_polynomials(i+1, xi)
        GETA[i] = _orth_polynomials(i+1, eta)

    S = {} # shape function
    SXI, SETA, SXIXI, SETAETA, SXIETA = {}, {}, {}, {}, {} # derivatives


    for i, fi in FXI.items():
        for j, gj in GETA.items():
            S[(i, j)] = fi * gj 
                

    txt = f'# This file is generated by {__file__}, it contains shape function \n\n'
    txt += '__all__ = [\'FXI\', \'S\']\n\n\n'
    txt += convert_to_str(FXI, 'FXI')
    txt += convert_to_str(S, 'S')
   

    with open('_S.py', 'w') as f:
        f.write(txt)
        
    print(f'time is {time() - ti} s')

