#!/usr/bin/env python
"""
Generate simplex legendre 2D basis coffecients and exponents matrices and
save them to legendre2D_simplex_coefs.txt and legendre2D_simplex_expos.txt
"""
from argparse import ArgumentParser
from sympy import symbols
from sympy import jacobi_poly as jacobi_P
from sympy import expand_mul as expand
from sympy import cancel, simplify

import numpy as np

from sfepy.base.ioutils import InDir
from sfepy.discrete.dg.poly_spaces import iter_by_order, get_n_el_nod

helps = {
    'max_order' :
        'maximum order of polynomials [default: %(default)s]',
    'output_dir':
        'output directory',
}


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument('-m', '--max-order', metavar='order', type=int,
                        action='store', dest='max_order',
                        default=10, help=helps['max_order'])
    parser.add_argument('output_dir', help=helps['output_dir'])
    options = parser.parse_args()

    indir = InDir(options.output_dir)

    a, b = symbols("a, b")
    r, s = symbols("r, s")
    x, y = symbols("x, y")
    order = options.max_order
    dim = 2
    n_el_nod = get_n_el_nod(order, dim)  # number of DOFs per element

    simplexP = []
    exponentM = np.zeros((n_el_nod, 3), dtype=np.int32)
    coefM = np.zeros((n_el_nod, n_el_nod))
    exponentList = []
    for m, idx in enumerate(iter_by_order(order, dim)):
        # print(m, idx)
        exponentM[m, :dim] = idx
        pa = jacobi_P(idx[0], 0, 0, a)
        pb = jacobi_P(idx[1], 2 * idx[0] + 1, 0, b) * (1 - b) ** idx[0]
        # print("P_{} = {}".format(m, pa*pb))
        polrs = cancel((pa * pb).subs(b, s).subs(
                a, 2 * (1 + r) / (1 - s) - 1))
        # print("P_{} = {}".format(m, polrs))
        polxy = expand(polrs.subs(r, 2 * x - 1).subs(s, 2 * y - 1))
        # polxy = expand(polrs.subs(r, x).subs(s, y))

        simplexP.append(simplify(polxy))
        exponentList.append(x ** idx[0] * y ** idx[1])
        for j, exponent in enumerate(exponentList):
            coefM[m, j] = simplexP[m].as_coefficients_dict()[exponent]
        print("P_{}{} = {}".format(m, idx, simplexP[m]))
        print()

    np.savetxt(indir("legendre2D_simplex_expos.txt"), exponentM, fmt="%d")
    # are coefs always integers?
    np.savetxt(indir("legendre2D_simplex_coefs.txt"), coefM, fmt="%d")


if __name__ == '__main__':
    main()
