#!/usr/bin/env python

"""
Installer script to install the LAMMPS python package and the corresponding
shared library into either the system-wide site-packages tree, or - failing
that - into the corresponding user tree. Called from the 'install-python'
build target in the conventional and CMake based build systems
"""

# copy LAMMPS shared library and lammps package to system dirs

from __future__ import print_function
import sys,os,re,shutil
from argparse import ArgumentParser

parser = ArgumentParser(prog='install.py',
                        description='LAMMPS python package installer script')

parser.add_argument("-p", "--package", required=True,
                    help="path to the LAMMPS Python package")
parser.add_argument("-l", "--lib", required=True,
                    help="path to the compiled LAMMPS shared library")
parser.add_argument("-v", "--version", required=True,
                    help="path to the LAMMPS version.h header file")

parser.add_argument("-d","--dir",
                    help="Legacy custom installation folder selection for package and library")

args = parser.parse_args()

# validate arguments and make paths absolute

if args.package:
  if not os.path.exists(args.package):
    print( "ERROR: LAMMPS package %s does not exist" % args.package)
    parser.print_help()
    sys.exit(1)
  else:
    args.package = os.path.abspath(args.package)

if args.lib:
  if not os.path.exists(args.lib):
    print( "ERROR: LAMMPS shared library %s does not exist" % args.lib)
    parser.print_help()
    sys.exit(1)
  else:
    args.lib = os.path.abspath(args.lib)

if args.version:
  if not os.path.exists(args.version):
    print( "ERROR: LAMMPS version header file %s does not exist" % args.version)
    parser.print_help()
    sys.exit(1)
  else:
    args.version = os.path.abspath(args.version)

if args.dir:
  if not os.path.isdir(args.dir):
    print( "ERROR: Installation folder %s does not exist" % args.dir)
    parser.print_help()
    sys.exit(1)
  else:
    args.dir = os.path.abspath(args.dir)

# if a custom directory is given, we copy the files directly
# without any special processing or additional steps to that folder

if args.dir:
  print("Copying LAMMPS Python package to custom folder %s" % args.dir)
  try:
    shutil.copytree(args.package, os.path.join(args.dir,'lammps'))
  except shutil.Error:
    pass # fail silently

  print("Copying LAMMPS shared library to custom folder %s" % args.dir)
  try:
    shutil.copyfile(args.lib, os.path.join(args.dir,os.path.basename(args.lib)))
  except shutil.Error:
    pass # fail silently

  sys.exit()

# extract version string from header
def get_lammps_version(header):
    with open(header, 'r') as f:
        line = f.readline()
        start_pos = line.find('"')+1
        end_pos = line.find('"', start_pos)
        return "".join(line[start_pos:end_pos].split())

verstr = get_lammps_version(args.version)

# convert string version to numeric version
vernum = 0
vregex = re.compile(r"([0-9]+)([A-Za-z]+)(2[0-9]+)")
m = vregex.match(verstr)
if (m):
    month2num = { 'Jan' : 1, 'Feb' : 2, 'Mar' : 3, 'Apr' : 4, 'May' : 5, 'Jun' : 6,
            'Jul' : 7, 'Aug' : 8, 'Sep' : 9, 'Oct' : 10, 'Nov' : 11, 'Dec' : 12 }
    try:
      vernum = int(m.group(3))*10000
      vernum += month2num[m.group(2)]*100
      vernum += int(m.group(1))
    except:
      exit('Failure to parse version string: %s' % verstr)

print("Installing LAMMPS Python package version %s into site-packages folder" % verstr)

# we need to switch to the folder of the python package
os.chdir(os.path.dirname(args.package))



# update version number in lammps module
vregex = re.compile(r".*(__version__ += +)[0-9]+")
with open(os.path.join('lammps','__init__.py'), "r+") as f:
    content = f.read()
    f.seek(0)
    f.write(re.sub(vregex, lambda match: '{}{}'.format(match.group(1), vernum), content))
    f.truncate()

from distutils.core import setup
from distutils.sysconfig import get_python_lib
import site

#Arguments common to global or user install -- everything but data_files
setup_kwargs= dict(name="lammps",
        version=verstr,
        author="Steve Plimpton",
        author_email="sjplimp@sandia.gov",
        url="https://lammps.sandia.gov",
        description="LAMMPS Molecular Dynamics Python package",
        license="GPL",
        packages=["lammps","lammps.mliap"],
        )

tryuser=False
try:
  sys.argv = ["setup.py","install"]    # as if had run "python setup.py install"
  setup_kwargs['data_files']=[(os.path.join(get_python_lib(), 'lammps'), [args.lib])]
  setup(**setup_kwargs)
except:
  tryuser=True
  print ("Installation into global site-packages folder failed.\nTrying user folder %s now." % site.USER_SITE)

if tryuser:
  try:
    sys.argv = ["setup.py","install","--user"]    # as if had run "python setup.py install --user"
    setup_kwargs['data_files']=[(os.path.join(site.USER_SITE, 'lammps'), [args.lib])]
    setup(**setup_kwargs)
  except:
    print("Installation into user site package folder failed.")

# restore __version__ == 0 for in place usage
with open(os.path.join('lammps','__init__.py'), "r+") as f:
    content = f.read()
    f.seek(0)
    f.write(re.sub(vregex, lambda match: '{}{}'.format(match.group(1), 0), content))
    f.truncate()