"""Script to initialize kernels via SymPy.

Running this script generates Fortran code for kernels
that is written to `kernels_base.f90`. This code has to
be compiled via `make` subsequently.
"""
#%%
from sympy import symbols, sqrt, exp, diff
from sympy.utilities.codegen import codegen
import shutil

#%% # Kernel functions (kernel, 1st and second derivatives)
r, ra, rb = symbols(r'r r_a r_b', real=True)
l = symbols(r'l', positive=True)

kern = {}
kern['sqexp'] = exp(-r**2/2)
kern['matern32'] = (1 + sqrt(3)*r)*exp(-sqrt(3)*r)
kern['matern52'] = (1 + sqrt(5)*r + 5*r/3)*exp(-sqrt(5)*r)
kern['wend4'] = (1 - r)**4*(1 + 4*r)

dkern = {}
for k, v in kern.items():
    dkern[k] = diff(v, r).simplify()

d2kern = {}
for k, v in dkern.items():
    d2kern[k] = diff(v, r).simplify()

funlist = []
for k in kern:
    funlist = funlist + [
        ('kern_{}'.format(k), kern[k]),
        ('dkern_{}'.format(k), dkern[k]),
        ('d2kern_{}'.format(k), d2kern[k])
    ]

[(name, code), (h_name, header)] = codegen(
    funlist, 'F95', 'kernels_base', header=False, empty=False)

with open(name, 'w') as f:
    f.write(code)
