#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------

from setuptools import setup, find_packages
from os import path, getcwd
import platform
import sys
import datetime

nightly_build = False
package_name = 'onnxruntime'

if '--use_tensorrt' in sys.argv:
    package_name = 'onnxruntime-gpu-tensorrt'
    sys.argv.remove('--use_tensorrt')
    if '--nightly_build' in sys.argv:
        package_name = 'ort-trt-nightly'
        nightly_build = True
        sys.argv.remove('--nightly_build')
elif '--use_cuda' in sys.argv:
    package_name = 'onnxruntime-gpu'
    sys.argv.remove('--use_cuda')
    if '--nightly_build' in sys.argv:
        package_name = 'ort-gpu-nightly'
        nightly_build = True
        sys.argv.remove('--nightly_build')
elif '--use_ngraph' in sys.argv:
    package_name = 'onnxruntime-ngraph'
    sys.argv.remove('--use_ngraph')

elif '--use_openvino' in sys.argv:
    package_name = 'onnxruntime-openvino'

if '--nightly_build' in sys.argv:
    package_name = 'ort-nightly'
    nightly_build = True
    sys.argv.remove('--nightly_build')

try:
    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
    class bdist_wheel(_bdist_wheel):
        def finalize_options(self):
            _bdist_wheel.finalize_options(self)
            self.root_is_pure = False
except ImportError:
    bdist_wheel = None

# Additional binaries
if platform.system() == 'Linux':
  libs = ['onnxruntime_pybind11_state.so', 'libmkldnn.so.0', 'libmklml_intel.so', 'libiomp5.so']
  # nGraph Libs
  libs.extend(['libngraph.so', 'libcodegen.so', 'libcpu_backend.so', 'libmkldnn.so', 'libtbb_debug.so', 'libtbb_debug.so.2', 'libtbb.so', 'libtbb.so.2'])
elif platform.system() == "Darwin":
  libs = ['onnxruntime_pybind11_state.so', 'libmkldnn.0.dylib'] # TODO add libmklml and libiomp5 later.
else:
  libs = ['onnxruntime_pybind11_state.pyd', 'mkldnn.dll', 'mklml.dll', 'libiomp5md.dll']

data = [path.join('capi', x) for x in libs if path.isfile(path.join('onnxruntime', 'capi', x))]

python_modules_list = list()
if '--use_openvino' in sys.argv:
  #Adding python modules required for openvino ep
  python_modules_list.extend(['openvino_mo', 'openvino_emitter'])
  sys.argv.remove('--use_openvino')

# Additional examples
examples_names = ["mul_1.pb", "logreg_iris.onnx", "sigmoid.onnx"]
examples = [path.join('datasets', x) for x in examples_names]

# Extra files such as EULA and ThirdPartyNotices
extra = ["LICENSE", "ThirdPartyNotices.txt"]

# Description
README = path.join(getcwd(), "docs/python/README.rst")
if not path.exists(README):
    this = path.dirname(__file__)
    README = path.join(this, "docs/python/README.rst")
if not path.exists(README):
    raise FileNotFoundError("Unable to find 'README.rst'")
with open(README) as f:
    long_description = f.read()


version_number = ''
with open('VERSION_NUMBER') as f:
    version_number = f.readline().strip()
if nightly_build:
    date_suffix = str(datetime.datetime.now().date().strftime("%m%d"))
    version_number = version_number + ".dev" + date_suffix

# Setup
setup(
    name=package_name,
    version=version_number,
    description='ONNX Runtime Python bindings',
    long_description=long_description,
    author='Microsoft Corporation',
    author_email='onnx@microsoft.com',
    cmdclass={'bdist_wheel': bdist_wheel},
    license="MIT License",
    packages=['onnxruntime',
              'onnxruntime.backend',
              'onnxruntime.capi',
              'onnxruntime.datasets',
              'onnxruntime.tools',
              ],
    package_data={
        'onnxruntime': data + examples + extra,
    },
    py_modules=python_modules_list,
    extras_require={
        'backend': ['onnx>=1.2.3'],
        'numpy': ['numpy>=1.15.0']
    },
    entry_points= {
        'console_scripts': [
            'onnxruntime_test = onnxruntime.tools.onnxruntime_test:main',
        ]
    },
    classifiers=[
        'Development Status :: 4 - Beta',
        'Environment :: Console',
        'Intended Audience :: Developers',
        'Operating System :: POSIX :: Linux',
        'Programming Language :: Python',
        'Programming Language :: Python :: 3 :: Only',
        'Programming Language :: Python :: 3.5',
        'Programming Language :: Python :: 3.6',
        'Programming Language :: Python :: 3.7'],
    )