#!/usr/bin/env python
#
# Copyright (C) 2011-2014 ABINIT Group (Yann Pouillon)
#
# 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.
#

import os
import re
import sys

from ConfigParser import ConfigParser
from copy import deepcopy
from datetime import datetime
from optparse import OptionParser

class MyConfigParser(ConfigParser):

  def optionxform(self,option):
    return str(option)

# Parse command-line arguments
my_help = """\
Usage: %prog [options] [files ...]

  This script is expected to be run from the top directory of an Abinit
  source tree. If this is not the case, you must provide options to
  pinpoint the locations of the required config files.

  When called without arguments, this script will update the
  ~/.abinit/build/hostname.ac file. Otherwise it will update all files
  specified on the command line, one by one.

  For each processed file, it will create a <filename>.new upgraded
  file. The original data will never be touched. You might also wish to
  review the new file before using it.

  If you're running this script from the top directory of an Abinit
  source tree, specifying an unexisting file to the --template option
  will result in using the default template. It will produce an error
  otherwise.

"""
parser = OptionParser(usage=my_help, version="%prog 0.1.0")
parser.add_option("-e", "--env", dest="env_file", metavar="FILE",
  help="Use the specified config file to read environment variables")
parser.add_option("-f", "--force", action="store_true", dest="force_write",
 default=False, help="Force overwriting destination file")
parser.add_option("-o", "--options", dest="opt_file", metavar="FILE",
  help="Use the specified config file to read options")
parser.add_option("-q", "--quiet", action="store_false", dest="show_actions",
 default=True, help="Don't display actions")
parser.add_option("-t", "--template", dest="tpl_file", metavar="FILE",
  help="Use the specified template when creating the new files")
(opts, args) = parser.parse_args()

# Banner
show_actions = opts.show_actions
if ( show_actions ):
  parser.print_version(sys.stdout)
  sys.stdout.write("\nType %s -h for help.\n\n")

# Check argument consistency
if ( opts.env_file ):
  if ( os.path.exists(opts.env_file) ):
    env_file = opts.env_file
  else:
    parser.error("environment file not found")
else:
  env_file = "config/specs/environment.conf"
  if ( not os.path.exists(env_file) ):
    parser.error("you must specify an environment file when outside an Abinit source tree")

if ( opts.opt_file ):
  if ( os.path.exists(opts.opt_file) ):
    opt_file = opts.opt_file
  else:
    parser.error("option file not found")
else:
  opt_file = "config/specs/options.conf"
  if ( not os.path.exists(opt_file) ):
    parser.error("you must specify an option file when outside an Abinit source tree")

tpl_file = "doc/build/config-template.ac"
if ( opts.tpl_file ):
  if ( os.path.exists(opts.tpl_file) ):
    tpl_file = opts.tpl_file
    tpl_use = True
  elif ( os.path.exists(tpl_file) ):
    tpl_use = True
  else:
    parser.error("template file not found")
else:
  tpl_use = False

for cfg in args:
  if ( not os.path.exists(cfg) ):
    parser.error("version numbers may not be equal")

if ( len(args) == 0 ):
  args = ["%s/.abinit/build/%s.ac" % (os.environ.get("HOME"),os.uname()[1])]

# Display parameters
if ( show_actions ):
  sys.stdout.write("Environment : %s\n" % (env_file))
  sys.stdout.write("Options     : %s\n" % (opt_file))
  sys.stdout.write("Template    : %s\n" % (tpl_file))
  sys.stdout.write("\n")

# Init parsers
cfg_env = MyConfigParser()
cfg_env.read(env_file)

cfg_opt = MyConfigParser()
cfg_opt.read(opt_file)

if ( tpl_use ):
  tpl_data = file(tpl_file).read()

# Init renames
env_renames = dict()
for env in cfg_env.sections():
  tmp = cfg_env.get(env,"status")
  if ( re.match("renamed",tmp) ):
    env_renames[tmp.split()[1]] = env

opt_renames = dict()
for opt in cfg_opt.sections():
  tmp = cfg_opt.get(opt,"status")
  if ( re.match("renamed",tmp) ):
    opt_renames[tmp.split()[1]] = opt

# Init replacements
opt_replacements = {
  "enable_atompaw":["with_dft_flavor","atompaw"],
  "enable_bigdft":["with_dft_flavor","bigdft"],
  "enable_etsf_io":["with_trio_flavor","etsf_io"],
  "enable_fox":["with_trio_flavor","fox"],
  "enable_libxc":["with_dft_flavor","libxc"],
  "enable_netcdf":["with_trio_flavor","netcdf"],
  "enable_wannier90":["with_dft_flavor","wannier90"]}

# Init connector defaults
opt_conn_set = dict()
opt_conn_set["with_dft_flavor"] = \
  cfg_opt.get("with_dft_flavor","default").split("+")
opt_conn_set["with_fft_flavor"] = \
  [cfg_opt.get("with_fft_flavor","default")]
opt_conn_set["with_linalg_flavor"] = \
  [cfg_opt.get("with_linalg_flavor","default")]
opt_conn_set["with_math_flavor"] = \
  [cfg_opt.get("with_math_flavor","default")]
opt_conn_set["with_timer_flavor"] = \
  [cfg_opt.get("with_timer_flavor","default")]
opt_conn_set["with_trio_flavor"] = \
  cfg_opt.get("with_trio_flavor","default").split("+")
opt_conn_set["with_dft_flavor"].sort()
opt_conn_set["with_trio_flavor"].sort()
opt_conn_def = deepcopy(opt_conn_set)

# Process each file
for src in args:

  if ( os.path.exists(src+".new") and (not opts.force_write) ):
    sys.stderr.write("WARNING: not overwriting %s\n" % (src+".new"))
    continue

  if ( show_actions ):
    sys.stdout.write("Processing %s\n" % (src))

  src_data = file(src).readlines()
  src_sets = dict()
  dst_data = list()
  index = 0

  opt_removed = dict()
  fcflags = dict()

  for line in src_data:
    re_cmt = re.search("#",line)

    if ( (not re_cmt) or (re_cmt.start() > 0) ):

      # Process environment variables
      for var in cfg_env.sections()+env_renames.keys():
        re_env = re.search("^[ ]*%s=[\"']?[^#]+[\"']?" % (var),line)
        if ( re_env ):
          if ( (not re_cmt) or (re_cmt and (re_env.start() < re_cmt.start())) ):
            val = re_env.group().partition("=")[2]
            val = val.strip()
            val = re.sub("^[\"']","",val)
            val = re.sub("[\"']$","",val)

            try:
              qry = cfg_env.get(var,"status")
            except:
              qry = "dummy"
            if ( var in env_renames ):
              key = env_renames[var]
              line = re.sub(var,key,line)
              src_sets[key] = val
              if ( show_actions ):
                sys.stdout.write("---> renamed '%s' to '%s'\n" % (var,key))
            elif ( qry == "removed" ):
              line = re.sub(var,"# Removed: "+var,line)
              if ( show_actions ):
                sys.stdout.write("---> commented '%s' assignment\n" % (var))
            elif ( qry == "dropped" ):
              line = re.sub(var,"# Dropped: "+var,line)
              if ( show_actions ):
                sys.stdout.write("---> commented '%s' assignment\n" % (var))
            else:
              src_sets[var] = val

      # Process options
      for var in cfg_opt.sections()+opt_renames.keys():
        re_opt = re.search("[ ]*%s=[\"']?[^#]+[\"']?" % (var),line)

        if ( re_opt ):
          if ( (not re_cmt) or (re_cmt and (re_opt.start() < re_cmt.start())) ):
            val = re_opt.group().partition("=")[2]
            val = val.strip()
            val = re.sub("^[\"']","",val)
            val = re.sub("[\"']$","",val)

            try:
              qry = cfg_opt.get(var,"status")
            except:
              qry = "dummy"
            if ( var in opt_renames ):
              key = opt_renames[var]
              line = re.sub(var,key,line)
              src_sets[key] = val
              if ( show_actions ):
                sys.stdout.write("---> renamed '%s' to '%s'\n" % (var,key))
            elif ( qry == "removed" ):
              line = re.sub(var,"# Removed: "+var,line)
              opt_removed[var] = val
              if ( show_actions ):
                sys.stdout.write("---> commented '%s' assignment\n" % (var))
            elif ( qry == "dropped" ):
              line = re.sub(var,"# Dropped: "+var,line)
              opt_removed[var] = val
              if ( show_actions ):
                sys.stdout.write("---> commented '%s' assignment\n" % (var))
            else:
              src_sets[var] = val

      # Collect Fortran optflags
      re_fcf = re.search("[ ]*fcflags_opt_[0-9a-z_]+=[\"']?[^#]+[\"']?",line)

      if ( re_fcf ):
        if ( (not re_cmt) or (re_cmt and (re_fcf.start() < re_cmt.start())) ):
          (var,sep,val) = re_fcf.group().partition("=")
          val = val.strip()
          val = re.sub("^[\"']","",val)
          val = re.sub("[\"']$","",val)
          fcflags[var] = val

    dst_data.append(line)
    index += 1

  # Replacements: banner
  if ( len(opt_removed) > 0 ):
    dst_data += "\n# Added by %s on %s\n" % (sys.argv[0],datetime.now()) \
             +  "# Replacements for removed options\n"

  # Generate replacements for FFT
  if ( "enable_fftw" in opt_removed ):
    if ( opt_removed["enable_fftw"] == "yes" ):
      opt_conn_set["with_fft_flavor"] = ["fftw3"]
    else:
      opt_conn_set["with_fft_flavor"] = ["none"]
  if ( "enable_fftw_threads" in opt_removed ):
    if ( opt_removed["enable_fftw_threads"] == "yes" ):
      opt_conn_set["with_fft_flavor"] = ["fftw3-threads"]
  else:
    if ( not "enable_fftw" in opt_removed ):
      opt_conn_set["with_fft_flavor"] = ["none"]
  if ( re.match("fftw3",opt_conn_set["with_fft_flavor"][0]) ):
    if ( "with_fftw_includes" in opt_removed ):
      dst_data += "with_fft_incs=\"" + \
        opt_removed["with_fftw_includes"]
      if ( "with_fft_incs" in src_sets ):
        dst_data += " " + src_sets["with_fft_incs"]
        src_sets["with_fft_incs"] += opt_removed["with_fftw_includes"]
      else:
        src_sets["with_fft_incs"] = opt_removed["with_fftw_includes"]
      dst_data += "\"\n"
    if ( "with_fftw_libs" in opt_removed ):
      dst_data += "with_fft_libs=\"" + \
        opt_removed["with_fftw_libs"]
      if ( "with_fft_libs" in src_sets ):
        dst_data += " " + src_sets["with_fft_libs"]
        src_sets["with_fft_libs"] = opt_removed["with_fftw_libs"] + \
          " " + src_sets["with_fft_libs"]
      else:
        src_sets["with_fft_libs"] = opt_removed["with_fftw_libs"]
      dst_data += "\"\n"

  # Generate replacements for MATH
  if ( "enable_gsl" in opt_removed ):
    if ( opt_removed["enable_gsl"] == "yes" ):
      opt_conn_set["with_math_flavor"] = ["gsl"]
    else:
      opt_conn_set["with_math_flavor"] = ["none"]
  if ( opt_conn_set["with_math_flavor"][0] == "gsl" ):
    if ( "with_gsl_includes" in opt_removed ):
      dst_data += "with_math_incs=\"" + \
        opt_removed["with_gsl_includes"]
      if ( "with_math_incs" in src_sets ):
        dst_data += " " + src_sets["with_math_incs"]
        src_sets["with_math_incs"] += opt_removed["with_gsl_includes"]
      else:
        src_sets["with_math_incs"] = opt_removed["with_gsl_includes"]
      dst_data += "\"\n"
    if ( "with_gsl_libs" in opt_removed ):
      dst_data += "with_math_libs=\"" + \
        opt_removed["with_gsl_libs"]
      if ( "with_math_libs" in src_sets ):
        dst_data += " " + src_sets["with_math_libs"]
        src_sets["with_math_libs"] = opt_removed["with_gsl_libs"] + \
          " " + src_sets["with_math_libs"]
      else:
        src_sets["with_math_libs"] = opt_removed["with_gsl_libs"]
      dst_data += "\"\n"

  # Generate replacements for TIMER
  if ( "enable_papi" in opt_removed ):
    if ( opt_removed["enable_papi"] == "yes" ):
      opt_conn_set["with_timer_flavor"] = ["papi"]
    else:
      opt_conn_set["with_timer_flavor"] = ["none"]
  if ( opt_conn_set["with_timer_flavor"][0] == "papi" ):
    if ( "with_papi_includes" in opt_removed ):
      dst_data += "with_timer_incs=\"" + \
        opt_removed["with_papi_includes"]
      if ( "with_timer_incs" in src_sets ):
        dst_data += " " + src_sets["with_timer_incs"]
        src_sets["with_timer_incs"] += opt_removed["with_papi_includes"]
      else:
        src_sets["with_timer_incs"] = opt_removed["with_papi_includes"]
      dst_data += "\"\n"
    if ( "with_papi_libs" in opt_removed ):
      dst_data += "with_timer_libs=\"" + \
        opt_removed["with_papi_libs"]
      if ( "with_timer_libs" in src_sets ):
        dst_data += " " + src_sets["with_timer_libs"]
        src_sets["with_timer_libs"] = opt_removed["with_papi_libs"] + \
          " " + src_sets["with_timer_libs"]
      else:
        src_sets["with_timer_libs"] = opt_removed["with_papi_libs"]
      dst_data += "\"\n"

  # Generate replacements for plugins
  if ( "enable_all_plugins" in opt_removed ):
    if ( opt_removed["enable_all_plugins"] == "yes" ):
      opt_conn_set["with_dft_flavor"] = \
        ["atompaw","bigdft","libxc","wannier90"]
      opt_conn_set["with_trio_flavor"] =\
        ["etsf_io","fox","netcdf"]
    else:
      opt_conn_set["with_dft_flavor"] = ["none"]
      opt_conn_set["with_trio_flavor"] = ["none"]
  else:
    for var in opt_replacements:
      if ( var in opt_removed ):
        (opt,val) = opt_replacements[var]
        if ( opt_removed[var] == "yes" ):
          if ( not val in opt_conn_set[opt] ):
            opt_conn_set[opt].append(val)
        else:
          if ( val in opt_conn_set[opt] ):
            opt_conn_set[opt].remove(val)

  # Generate replacements for ScaLAPACK
  if ( "enable_scalapack" in opt_removed ):
    if ( not "with_linalg_flavor" in src_sets ):
      if ( (opt_conn_set["with_linalg_flavor"][0] == "netlib") or \
           (opt_conn_set["with_linalg_flavor"][0] == "goto") ):
        opt_conn_set["with_linalg_flavor"][0] += "-mpi"
    if ( "with_scalapack_libs" in opt_removed ):
      dst_data += "with_linalg_libs=\"" + \
        opt_removed["with_scalapack_libs"]
      if ( "with_linalg_libs" in src_sets ):
        dst_data += " " + src_sets["with_linalg_libs"]
        src_sets["with_linalg_libs"] = opt_removed["with_scalapack_libs"] + \
          " " + src_sets["with_linalg_libs"]
      else:
        src_sets["with_linalg_libs"] = opt_removed["with_scalapack_libs"]
      dst_data += "\"\n"

  okeys = opt_conn_set.keys()
  okeys.sort()
  for key in okeys:
    if ( len(opt_conn_set[key]) > 0 ):
      opt_conn_set[key].sort()
      if ( key in opt_conn_def and
         (opt_conn_set[key] != opt_conn_def[key]) ):
        dst_data += "%s=\"%s\"\n" % (key,"+".join(opt_conn_set[key]))
        src_sets[key] = "+".join(opt_conn_set[key])
    else:
      dst_data += "%s=\"none\"\n" % (key)
      src_sets[key] = "none"

  # Output final result
  if ( tpl_use ):
    dst_data = ""
    re_fcflags = re.compile("^#fcflags_opt_95_drive=")
    for line in file(tpl_file,"r").readlines():
      if ( re_fcflags.match(line) ):
        fckeys = fcflags.keys()
        fckeys.sort()
        for var in fckeys:
          dst_data += "%s=\"%s\"\n" % (var,fcflags[var])
      else:
        for var in src_sets.keys():
          line = re.sub("^#"+var+"=.*","%s=\"%s\"" % (var,src_sets[var]),line)
        dst_data += line

  file(src+".new","w").write("".join(dst_data))

  if ( show_actions ):
    sys.stdout.write("\n")

