#!/usr/bin/env python3

import os, glob, jinja2
from shutil import copyfile


def render(tpl_path, context):
    path, filename = os.path.split(tpl_path)
    return jinja2.Environment(undefined=jinja2.StrictUndefined,
        loader=jinja2.FileSystemLoader(path or './')
    ).get_template(filename).render(context)


rankmap = {1:'vector', 2:'matrix', 3:'tensor', 4:'four_dimensional'}

def gen_interface_list(types):
    intrfc_list = []
    for t in types:
        for rank in (1,2,3,4):
            intrfc_list.append({
                      'vtype'     : t,
                      'rank'      : rank,
                      'vrankname' : rankmap[rank],
                      'subname'   : t[0]+rankmap[rank][0],
                      'vsize'     : rankmap[rank][0]+'size',
                      'vsizedims' : "({})".format(rank),
                      'ranks'     : ",".join([':']*rank)
                      })
    return intrfc_list


def gen_module(context):
    fname = 'devxlib_buffers.jf90'
    bname, _ = os.path.splitext(fname)
    with open('devxlib_buffer.f90','w') as f:
        f.write(render(fname, context))

def gen_module2(context):
    fname = 'devxlib_buffers.jf90'
    bname, _ = os.path.splitext(fname)
    with open('devxlib_pinned.f90','w') as f:
        f.write(render(fname, context))

def main():
    # clean up
    #open('devxlib_buffer.f90','w').close()

    # Types used in the interface
    interfaces = gen_interface_list(['integer', 'real(real64)', 'complex(real64)'])
    context = {'interfaces'  : interfaces,
               'cudamod'     : 'use cudafor',
               'ompmod'      : 'use omp_lib',
               'attributes'  : 'device',
               'pointer_type': 'c_devptr',
               'modulename'  : 'devxlib_buffer',
               'typename'    : 'devxlib_buffer_t'}
    gen_module(context)

#    context = {'interfaces'  : interfaces,
#               'cudamod'     : 'use cudafor',
#               'ompmod'      : 'use omp_lib',
#               'attributes'  : '',
#               'pointer_type': 'c_ptr',
#               'modulename'  : 'devxlib_pinned',
#               'typename'    : 'devxlib_pinned_t'}
#    gen_module(context)

def main2():
    # clean up
    #open('devxlib_pinned.f90','w').close()

    # Types used in the interface
    interfaces = gen_interface_list(['integer', 'real(real64)', 'complex(real64)'])
    context = {'interfaces'  : interfaces,
               'cudamod'     : 'use cudafor',
               'ompmod'      : 'use omp_lib',
               'attributes'  : '',
               'pointer_type': 'c_ptr',
               'modulename'  : 'devxlib_pinned',
               'typename'    : 'devxlib_pinned_t'}
    gen_module2(context)

    ## set multiple extensions
    #copyfile("device_fbuff.f90","device_fbuff.F")


if __name__ == "__main__":
    main()
    main2()
