#!/usr/bin/env python
#
# Copyright (c) 2005-2006 The ABINIT Group (Yann Pouillon)
# All rights reserved.
#
# This file is part of the ABINIT software package. For license information,
# please see the COPYING file in the top-level directory of the ABINIT source
# distribution.
#

from time import gmtime,strftime

import commands
import os
import re
import sys

# ---------------------------------------------------------------------------- #

#
# Subprograms
#

# Forwarder for grouped data access
def code_data(action):

 return eval("code_data_%s()" % (action))



# Code for grouped data closing
def code_data_close():

 ret = "ncerr = nf90_close(ncid)\n" \
     + "call etsf_handle_ncerr('I/O','END','current file',my_name,ncerr)"

 return ret



# Code for grouped data initialization
def code_data_init():

 ret = """! Create the NetCDF file
ncerr = nf90_create(path=filename,cmode=NF90_NOCLOBBER,ncid=ncid)
call etsf_handle_ncerr('I/O','INI',filename,my_name,ncerr)

! Write global attributes
call etsf_put_globals(ncid,my_title,my_history)

! Define dimensions
call etsf_def_dims(ncid,dims)
"""

 ret += code_data_select("def")
 
 return ret



# Code for grouped data reading
def code_data_read():

 return code_data_select("get")



# Code for grouped data selection
def code_data_select(action):

 ret = "do i=1,etsf_ngroups\n select case ( groups )"

 for group in etsf_group_list:
  ret += "\n\n  case (etsf_grp_%s)\n   %s" % \
   (group,indent_code(code_optional_argument(group,action),3))

 ret += "\n\n end select\nend do"

 return ret



# Code for grouped data writing
def code_data_write():

 return code_data_select("put")
 


# Code for dimensions
def code_dims(action):

 ret = ""
 for dim in etsf_dimensions:
  if ( ret != "" ):
   ret += "\n\n"

  if ( action == "def" ):
   ret += "ncerr = nf90_def_dim(ncid,name='%s', &\n" % (dim) \
        + " & len=dims%%%s,dimid=ncdim)\n" % (dim) \
        + "call etsf_handle_ncerr('DEF','DIM','%s', &\n" % (dim) \
        + " & my_name,ncerr)"
  elif ( action == "get" ):
   ret += "ncerr = nf90_inq_dimid(ncid,name='%s', &\n" % (dim) \
        + " & dimid=ncdim)\n" \
        + "call etsf_handle_ncerr('GET','DID','%s', &\n" % (dim) \
        + " & my_name,ncerr)\n" \
        + "ncerr = nf90_inq_dimlen(ncid,dimid=ncdim, &\n" \
        + " & len=dims%%%s)\n" % (dim) \
        + "call etsf_handle_ncerr('GET','DIM','%s', &\n" % (dim) \
        + " & my_name,ncerr)"

 return ret

 

# Code for global attributes
def code_globals(action):

 ret = ""
 for att in etsf_attributes:
  if ( ret != "" ):
   ret += "\n\n"

  ret += "ncerr = nf90_%s_att(ncid,varid=NF90_GLOBAL,name='%s', &\n" % \
           (action,att[0]) \
       + " & values=etsf_%s)\n" % (att[0].lower()) \
       + "call etsf_handle_ncerr('%s','ATT','%s', &\n" % (action.upper(),att[0]) \
       + " & my_name,ncerr)"

 return ret



# Generic code for a group
def code_group_generic(group,action):

 ret = ""

 # Look for peculiarities
 if ( group in etsf_properties ):
  specs = etsf_properties[group]
 else:
  specs = ETSF_PRO_NONE

 # Process each variable in the group
 for var in etsf_groups[group]:
  var_desc = etsf_variables[var]

  if ( ret != "" ):
   ret += "\n\n"

  # Define: fetch dimensions first
  if ( action == "def" ):
   target = "name='%s',xtype=%s" % (var,nc_types[var_desc[0]])

   if ( len(var_desc) > 1 ):
    ret += "allocate(ncdims(%d))\n" % (len(var_desc)-1)
    index  = len(var_desc)-1
    for dim in var_desc[1:]:
     ret += "ncerr = nf90_inq_dimid(ncid,'%s',ncdims(%d))\n" % (dim,index) \
         +  "call etsf_handle_ncerr('INQ','DID','%s', &\n" % (dim) \
         +  " & my_name,ncerr)\n"
     index -= 1
    target += ",dimids=ncdims,varid=ncvar"
  else:
   target = "varid=ncvar,"
   if ( (specs & ETSF_PRO_GRP_MAIN) == 0 ):
    target += "values=folder%%%s" % (var)
   else:
    target += "values=%s" % (var)

   if ( action != "def" ):
    ret += "ncerr = nf90_inq_varid(ncid,'%s',ncvar)\n" % (var) \
         + "call etsf_handle_ncerr('INQ','VID','%s', &\n" % (var) \
         + " & my_name,ncerr)\n"

  # Perform the action
  ret += "ncerr = nf90_%s_var(ncid, &\n" % (action) \
       + " & %s)\n" % (target) \
       + "call etsf_handle_ncerr('%s','VAR','%s', &\n" % (action.upper(),var) \
       + " & my_name,ncerr)"

  # Free memory
  if ( (action == "def") and (len(var_desc) > 1) ):
   ret += "\ndeallocate(ncdims)"

 return ret



# Code for main group checking
def code_main(action):

 ret = ""

 if ( action == "check" ):
  ret = "etsf_main_check = .true.\nexclude_others  = .false."

  for group in etsf_groups:
   if ( group in etsf_properties ):
    specs = etsf_properties[group]
   else:
    specs = ETSF_PRO_NONE

   if ( (specs & ETSF_PRO_GRP_MAIN) != 0 ):
    ret += "\n\nif ( iand(groups,etsf_grp_%s) ) then\n" % (group) \
         + " if ( exclude_others ) etsf_main_check = .false.\n" \
         + " exclude_others = .true.\nend if"

 return ret



# Transfer data to and from an optional argument
def code_optional_argument(group,action):

 ret = "if ( present(%s) ) then\n" % (group)

 if ( action == "get" ):
  ret += " call etsf_get_dims(ncid,dims)\n\n"

  for var in etsf_groups[group]:
   var_desc = etsf_variables[var]
   if ( len(var_desc) > 1 ):
    ret += " allocate(my_%s%%%s(%s))\n" % \
     (group,var,"dims%"+",dims%".join(var_desc[1:]))
  ret += " call etsf_get_%s(ncid,my_%s)\n" % (group,group)

 if ( action == "def" ):
  ret += " call etsf_%s_%s(ncid,my_%s)\n" % (action,group,group)

 for var in etsf_groups[group]:
  var_dims = len(etsf_variables[var])-1
  if ( action == "put" ):
   if ( var_dims > 0 ):
    ret += " my_%s%%%s => %s%%%s\n" % (group,var,group,var)
   else:
    ret += " my_%s%%%s = %s%%%s\n" % (group,var,group,var)
  elif ( action == "get" ):
    ret += " %s%%%s = my_%s%%%s\n" % (group,var,group,var)

 if ( action == "put" ):
  ret += " call etsf_put_%s(ncid,my_%s)\n" % (group,group)
 elif ( action == "get" ):
  for var in group[1:]:
   ret += "\n deallocate(my_%s%%%s)" % (group,var)

 ret += "endif"

 return ret



# Indent a piece of code
def indent_code(code,offset):

 tmp = ""
 for i in range(offset):
  tmp += " "

 return tmp+re.sub("\n","\n"+tmp,code)



# Initialize a routine from a template
def init_routine(name,template,info,script,args,type="subroutine"):

 # Init
 ret = file("config/etsf/template.%s" % (template),"r").read()

 if ( type == "subroutine" ):
  ret = re.sub("@SUBPROGRAM@",type,ret)
 else:
  ret = re.sub("@SUBPROGRAM@","function",ret)
  ret = re.sub("@FUNCTYPE@",type,ret)

 ret = re.sub("@NAME@",name,ret)
 ret = re.sub("@INFO@",info,ret)
 ret = re.sub("@SCRIPT@",script,ret)

 arg_list = ""
 arg_move = len(name)+16
 arg_stop = 0

 # Process arguments
 if ( args != None ):
  for arg_str in args:
   arg_info = arg_str.split()

   if ( arg_list != "" ):
    if ( (len(arg_list)+arg_move)/72 > arg_stop ):
     arg_list += ", &\n  & "
     arg_stop += 1
    else:
     arg_list += ","
   arg_list += arg_info[0]
  ret = re.sub("@ARG_LIST@",arg_list,ret)

  arg_desc = ""
  loc_vars = ""
  for arg_str in args:
   arg_info = arg_str.split()
   arg = arg_info[0]

   # Optional arguments
   if ( (len(arg_info) > 3) and (arg_info[3] == "optional") ):
    opt = ",optional"
   else:
    opt = ""

   # Arrays
   if ( len(arg_info) > 4 ):
    dim = ":"
    for i in range(int(arg_info[4])-1):
     dim += ",:"
    arg_desc += " %s%s,intent(%s) :: %s(%s)\n" % \
     (arg_info[1],opt,arg_info[2],arg,dim)
    if ( arg_info[3] == 1 ):
     loc_vars += " %s,allocatable :: my_%s(%s)\n" % (arg_info[1],arg,dim)
   else:
    arg_desc += " %s%s,intent(%s) :: %s\n" % (arg_info[1],opt,arg_info[2],arg)
    if ( (len(arg_info) > 3) and (arg_info[3] == "optional") ):
     loc_vars += " %s :: my_%s\n" % (arg_info[1],arg)

  if ( type != "subroutine" ):
   loc_vars += " %s :: etsf_%s\n" % (type,name)

  ret = re.sub("@ARG_DESC@",arg_desc,ret)
  ret = re.sub("@LOCAL_VARS@",loc_vars,ret)
 else:
  ret = re.sub(",@ARG_LIST@","",ret)
  ret = re.sub("@ARG_DESC@","",ret)
  ret = re.sub("@LOCAL_VARS@","",ret)

 return ret



# ---------------------------------------------------------------------------- #

#
# Main program
#

# Initial setup
my_name    = "etsf-make-routines"
my_configs = ["config/etsf/specs.cf",
              "config/etsf/library.cf",
              "config/etsf/abinit.cf"]

# Check if we are in the top of the ABINIT source tree
if ( not os.path.exists("configure.ac") or
     not os.path.exists("src/main/abinit.F90") ):
 print "%s: You must be in the top of an ABINIT source tree." % my_name
 print "%s: Aborting now." % my_name
 sys.exit(1)

# Read config file(s)
for cnf in my_configs:
 if ( os.path.exists(cnf) ):
  execfile(cnf)
 else:
  print "%s: Could not find config file (%s)." % (my_name,cnf)
  print "%s: Aborting now." % my_name
  sys.exit(2)

# Create routines
for sub in etsf_subprograms.keys():
 dsc = etsf_subprograms[sub]

 # Look for peculiarities
 if ( sub in etsf_subs_properties ):
  specs = etsf_subs_properties[sub]
 else:
  specs = ETSF_SUBS_PRO_NONE

 # Check type
 if ( len(dsc) > 3 ):
  sub_type = dsc[3]
 else:
  sub_type = "subroutine"

 # Write action routines
 for action in dsc[2].split():
  if ( sub == "@GROUP@" ):
   sub_list = etsf_group_list
   sub_code = "group_generic"
   sub_cprm = "sub_name,action"
  else:
   sub_list = [sub]
   sub_code = sub
   sub_cprm = "action"

  for sub_name in sub_list:
   # Look for arguments
   try:
    if ( sub == "@GROUP@" ):
     sub_args = eval("etsf_subs_%s_args[sub]" % (action))
    else:
     sub_args = eval("etsf_subs_%s_args[sub_name]" % (action))
   except KeyError:
    sub_args = None

   # Load template
   src = init_routine("%s_%s" % (sub_name,action),dsc[0],dsc[1],my_name,
    sub_args,type=sub_type)

   # Substitute patterns
   src = re.sub("@ACTION_TEXT@",etsf_subs_actions[action].capitalize()+"s",src)
   src = re.sub("@GROUP@",sub_name,src)
   src = re.sub("@GROUP_TYPE@","type(etsf_%s)" % (sub_name),src)
   src = re.sub("@CODE@",
          indent_code(eval("code_%s(%s)" % (sub_code,sub_cprm)),1),src)

   # Write routine
   out = file("%s/etsf_%s_%s.F90" % (etsf_file_srcdir,sub_name,action),"w")
   out.write(src)
   out.close()