#!/usr/bin/env python
"""
Produce the routines used to write the content of a Fortran array to a netcdf file
These tools are mainly used for debugging or for a quick analysis of data.
"""

def dimension_shape_type_kind(ndim, type, kind):
    """Return a string with the fortran dimension."""
    if ndim:
        dimension = ("(" + ndim*":,")[:-1] + ")"
    else:
        dimension = ""
                                                    
    shape = "(%d)" % ndim if ndim > 0 else ""
                                                    
    if kind:
        type_kind = "%s(%s)" % (type, kind)
    else:
        type_kind = type
        if type == "character":
            type_kind = "%s(len=%s)" % (type, kind)

    return dimension, shape, type_kind

def ncwrite(type, kind, ndim):
    """
    Generate the routine for the allocation of pointers/arrays

    Args:
        type:
            Fortran intrinsic type (real, integer, complex, logical, character)
        kind: 
            Kind of the type. if type is 'character' ...
        ndim:
            integer giving the number of dimensions.
    """
    # Construct the name of the function
    function = "farr_" + type + "_" + kind + str(ndim)

    dimension, shape, type_kind = dimension_shape_type_kind(ndim, type, kind)
    
    # Fortran type --> ETSF-IO flags.
    xtype = {
        ("integer", ""): "int",
        ("real", "sp"): "sp",
        ("real", "dp"): "dp",
        ("complex", "spc"): "sp",
        ("complex", "dpc"): "dp",
        ("character", ""): "char",
    }[(type, kind)]

    # Use c2r to convert from complex to real
    # Multiply dims by 2 to take into account both real and imag part.
    # This trick can cause a sigfault for large arrays since
    # the compiler may store the result of c2r on the stack.
    convert_farray = "farray"; convert_fact = 1
    if type == "complex": 
        convert_farray = "c2r(farray)"
        convert_fact = 2

    # Template string 
    template = """\
subroutine %(function)s(farray, varname, ncid)

 implicit none

!Arguments ------------------------------------
 integer,intent(in) :: ncid
 character(len=*),intent(in) :: varname
 %(type_kind)s,intent(in) :: farray%(dimension)s

!Local variables-------------------------------
!scalars
#ifdef HAVE_NETCDF
 integer :: dimval,ii,varid
 character(len=nctk_slen) :: dimname,shape_str   

! *********************************************************************

! Define dimensions.
! Use *private* names for dimensions to avoid possible name collisions.
! if varname is "foo", we use "__foo_dim1__", __foo_dim2__" for name of the dimensions
 NCF_CHECK(nctk_set_datamode(ncid))
 
 shape_str = ""
 do ii=1,size(shape(farray))
   dimval = size(farray, dim=ii)
   if (ii == 1) dimval = size(farray, dim=ii) * %(convert_fact)s
   write(dimname,"(3a,i0,a)")"__",trim(varname),"_dim",ii,"__"

   NCF_CHECK(nctk_def_dims(ncid, nctkdim_t(dimname, dimval)))
   shape_str = trim(shape_str) // ", " // dimname 
 end do

 ! Define the variable.
 NCF_CHECK(nctk_def_arrays(ncid, nctkarr_t(varname, '%(xtype)s', shape_str)))
 varid = nctk_idname(ncid, varname)

 ! Write data
 NCF_CHECK(nctk_set_datamode(ncid))
 NCF_CHECK(nf90_put_var(ncid, varid, %(convert_farray)s))

#else 
 MSG_ERROR("netcdf support is not activated.")
#endif

end subroutine %(function)s
"""
    return template % locals()


def main():
    all_ndims = range(1,2)

    # Table with types and kinds
    type_and_kinds = [
        #("integer", ("",)),
        ("real", ("dp",)),
        ("complex", ("dpc",)),
        #("logical", ("",)),
        #("character", ("",)),
    ]

    # The subroutines we are gonna create
    subs = [ 
        ncwrite,
    ]

    text = []
    for sub in subs:
        for ndim in all_ndims:
            for type, kinds in type_and_kinds:
                for kind in kinds:
                    text.append(sub(type=type, kind=kind, ndim=ndim))

    header = "! NOTE: This code has been generated by genarray.py\n"
    text = ("\n!" + 80*"-" +"\n").join(text)
    print(header + text)

    return 0


if __name__ == "__main__":
    import sys
    sys.exit(main())
