from setuptools import setup, Distribution, Extension
import subprocess
import shutil
import os
class BinaryDistribution(Distribution):
    def has_ext_modules(foo):
        return True

MAJOR   = 0
MINOR   = 14
PATCH   = 0
RC      = 0
ISTAGED = False



def git_commit():
    try:
        cmd = ['git', 'rev-parse', 'HEAD']
        git_commit = subprocess.Popen(cmd, stdout = subprocess.PIPE).communicate()[0].strip()
    except:
        git_commit = 'Unknown'
    return git_commit

def write_version_py(filename='paddle/version.py'):
    cnt = '''
# THIS FILE IS GENERATED FROM PADDLEPADDLE SETUP.PY
#
full_version    = '%(major)d.%(minor)d.%(patch)d'
major           = '%(major)d'
minor           = '%(minor)d'
patch           = '%(patch)d'
rc              = '%(rc)d'
istaged         = %(istaged)s
commit          = '%(commit)s'
with_mkl        = '%(with_mkl)s'

def show():
    if istaged:
        print 'full_version:', full_version
        print 'major:', major
        print 'minor:', minor
        print 'patch:', patch
        print 'rc:', rc
    else:
        print 'commit:', commit

def mkl():
    return with_mkl
'''
    commit = git_commit()
    with open(filename, 'w') as f:
        f.write(cnt % {
            'major': MAJOR,
            'minor': MINOR,
            'patch': PATCH,
            'rc': RC,
            'version': '${PADDLE_VERSION}',
            'commit': commit,
            'istaged': ISTAGED,
            'with_mkl': '@WITH_MKL@'})

write_version_py(filename='@PADDLE_BINARY_DIR@/python/paddle/version.py')


packages=['paddle',
          'paddle.libs',
          'paddle.utils',
          'paddle.dataset',
          'paddle.reader',
          'paddle.fluid',
          'paddle.fluid.proto',
          'paddle.fluid.proto.profiler',
          'paddle.fluid.layers',
          'paddle.fluid.transpiler',
          'paddle.fluid.transpiler.details']

if '${WITH_FLUID_ONLY}'== 'OFF':
    packages+=['paddle.proto',
               'paddle.trainer',
               'paddle.trainer_config_helpers',
               'paddle.v2',
               'paddle.v2.master',
               'paddle.v2.plot',
               'paddle.v2.reader',
               'paddle.v2.dataset',
               'py_paddle']

with open('@PADDLE_SOURCE_DIR@/python/requirements.txt') as f:
    setup_requires = f.read().splitlines()

if '${CMAKE_SYSTEM_PROCESSOR}' not in ['arm', 'armv7-a', 'aarch64']:
    setup_requires+=['opencv-python']

# the prefix is sys.prefix which should always be usr
paddle_bins = ''
if '${WITH_FLUID_ONLY}'== 'OFF':
    paddle_bin_dir = 'opt/paddle/bin'
    paddle_bins = ['${PADDLE_BINARY_DIR}/paddle/legacy/trainer/paddle_trainer',
                   '${PADDLE_BINARY_DIR}/paddle/legacy/trainer/paddle_merge_model',
                   '${PADDLE_BINARY_DIR}/paddle/legacy/pserver/paddle_pserver_main',
                   '${PADDLE_BINARY_DIR}/paddle/scripts/paddle']

package_data={'paddle.fluid': ['core.so']}
if '${WITH_FLUID_ONLY}'== 'OFF':
    package_data['paddle.v2.master']=['libpaddle_master.so']
    package_data['py_paddle']=['*.py','_swig_paddle.so']

package_dir={
    '': '${PADDLE_BINARY_DIR}/python',
    # The paddle.fluid.proto will be generated while compiling.
    # So that package points to other directory.
    'paddle.fluid.proto.profiler': '${PADDLE_BINARY_DIR}/paddle/fluid/platform',
    'paddle.fluid.proto': '${PADDLE_BINARY_DIR}/paddle/fluid/framework',
    'paddle.fluid': '${PADDLE_BINARY_DIR}/python/paddle/fluid',
}
if '${WITH_FLUID_ONLY}'== 'OFF':
    package_dir['py_paddle']='${PADDLE_BINARY_DIR}/python/py_paddle'

# put all thirdparty libraries in paddle.libs
package_data['paddle.libs']=['libwarpctc.so']
libs_path='${PADDLE_BINARY_DIR}/python/paddle/libs'
shutil.copy('${WARPCTC_LIBRARIES}', libs_path)
if '${WITH_MKL}' == 'ON':
    shutil.copy('${MKLML_LIB}', libs_path)
    shutil.copy('${MKLML_IOMP_LIB}', libs_path)
    package_data['paddle.libs']+=['libmklml_intel.so','libiomp5.so']
if '${WITH_MKLDNN}' == 'ON':
    # change rpath of libmkldnn.so.0, add $ORIGIN/ to it.
    # The reason is that all thirdparty libraries in the same directory,
    # thus, libmkldnn.so.0 will find libmklml_intel.so and libiomp5.so.
    command = "patchelf --set-rpath '$ORIGIN/' ${MKLDNN_SHARED_LIB}"
    if os.system(command) != 0:
        raise Exception("patchelf --set-rpath for libmkldnn.so.0 fails")
    package_data['paddle.libs']+=['libmkldnn.so.0']
    shutil.copy('${MKLDNN_SHARED_LIB}', libs_path)
# remove unused paddle/libs/__init__.py
os.remove(libs_path+'/__init__.py')
package_dir['paddle.libs']=libs_path

# change rpath of core.so, add $ORIGIN/../libs/ to it.
# The reason is that libwarpctc.so, libiomp5.so etc are in paddle.libs, and
# core.so is in paddle.fluid, thus paddle/fluid/../libs will pointer to above libraries.
# This operation will fix https://github.com/PaddlePaddle/Paddle/issues/3213
command = "patchelf --set-rpath '$ORIGIN/../libs/' ${PADDLE_BINARY_DIR}/python/paddle/fluid/core.so"
if os.system(command) != 0:
    raise Exception("patchelf --set-rpath for core.so fails")

setup(name='${PACKAGE_NAME}',
      version='${PADDLE_VERSION}',
      description='Parallel Distributed Deep Learning',
      install_requires=setup_requires,
      packages=packages,
      ext_modules=[Extension('_foo', ['stub.cc'])],
      package_data=package_data,
      package_dir=package_dir,
      scripts=paddle_bins
)