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

F = {0: f1, 1: f2, 2: f3, 3: f4} # first terms of set of functions
FXI = {
    0: sp.diff(f1, xi),
    1: sp.diff(f2, xi),
    2: sp.diff(f3, xi),
    3: sp.diff(f4, xi)   
}
FXIXI = {
    0: sp.diff(f1, xi, xi),
    1: sp.diff(f2, xi, xi),
    2: sp.diff(f3, xi, xi),
    3: sp.diff(f4, xi, xi)   
}


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):
        F[i] = _orth_polynomials(i+1, xi)
        FXI[i] = sp.diff(F[i], xi)
        FXIXI[i] = sp.diff(F[i], xi, xi)


    FF, FXI_F, FXI_FXI, FXIXI_F, FXIXI_FXI, FXIXI_FXIXI = {}, {}, {}, {}, {}, {} # shape function
    ii_FF, ii_FXI_F, ii_FXI_FXI, ii_FXIXI_F, ii_FXIXI_FXI, ii_FXIXI_FXIXI = {}, {}, {}, {}, {}, {} # shape function

    for i in range(n_terms):
        for k in range(n_terms):
            FF[(i, k)] = F[i] * F[k]
            FXI_F[(i, k)] = FXI[i] * F[k]
            FXI_FXI[(i, k)] = FXI[i] * FXI[k]
            FXIXI_F[(i, k)] = FXIXI[i] * F[k]
            FXIXI_FXI[(i, k)] = FXIXI[i] * FXI[k]
            FXIXI_FXIXI[(i, k)] = FXIXI[i] * FXIXI[k]

            ii_FF[(i, k)] = sp.integrate(FF[(i, k)], (xi, -1, 1))
            ii_FXI_F[(i, k)] = sp.integrate(FXI_F[(i, k)], (xi, -1, 1))
            ii_FXI_FXI[(i, k)] = sp.integrate(FXI_FXI[(i, k)], (xi, -1, 1))
            ii_FXIXI_F[(i, k)] = sp.integrate(FXIXI_F[(i, k)], (xi, -1, 1))
            ii_FXIXI_FXI[(i, k)] = sp.integrate(FXIXI_FXI[(i, k)], (xi, -1, 1))
            ii_FXIXI_FXIXI[(i, k)] = sp.integrate(FXIXI_FXIXI[(i, k)], (xi, -1, 1))
     

    txt = f'# This file is generated by {__file__}, it contains the preintegrated terms\n\n'
    txt += '__all__ = [\'ii_FF\', \'ii_FXI_F\', \'ii_FXI_FXI\', \'ii_FXIXI_F\', \'ii_FXIXI_FXI\', \'ii_FXIXI_FXIXI\']\n\n\n'
    txt += convert_to_str(ii_FF, 'ii_FF')
    txt += convert_to_str(ii_FXI_F, 'ii_FXI_F')
    txt += convert_to_str(ii_FXI_FXI, 'ii_FXI_FXI')
    txt += convert_to_str(ii_FXIXI_F, 'ii_FXIXI_F')
    txt += convert_to_str(ii_FXIXI_FXI, 'ii_FXIXI_FXI')
    txt += convert_to_str(ii_FXIXI_FXIXI, 'ii_FXIXI_FXIXI')


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

