# coding=utf-8
# Copyright 2020 The TensorFlow GAN Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""TF-GAN: A Generative Adversarial Networks library for TensorFlow.

TF-GAN is a lightweight library for training and evaluating Generative
Adversarial Networks (GANs).

See the README on GitHub for further documentation.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import io
import sys
import unittest

from setuptools import find_packages
from setuptools import setup
from setuptools.command.install import install as InstallCommandBase
from setuptools.command.test import test as TestCommandBase
from setuptools.dist import Distribution


project_name = 'tensorflow-gan'

# Get version from version module.
with open('tensorflow_gan/python/version.py') as fp:
  globals_dict = {}
  exec(fp.read(), globals_dict)  # pylint: disable=exec-used
__version__ = globals_dict['__version__']
version = __version__


class StderrWrapper(io.IOBase):

  def write(self, *args, **kwargs):
    return sys.stderr.write(*args, **kwargs)

  def writeln(self, *args, **kwargs):
    if args or kwargs:
      sys.stderr.write(*args, **kwargs)
    sys.stderr.write('\n')


class Test(TestCommandBase):

  def run_tests(self):
    # Import absl inside run, where dependencies have been loaded already.
    from absl import app  # pylint: disable=g-import-not-at-top

    def main(_):
      test_loader = unittest.TestLoader()
      test_suite = test_loader.discover('tensorflow_gan', pattern='*_test.py')
      stderr = StderrWrapper()
      result = unittest.TextTestResult(stderr, descriptions=True, verbosity=2)
      test_suite.run(result)

      result.printErrors()

      final_output = ('Tests run: {}. Errors: {}  Failures: {}.'.format(
          result.testsRun, len(result.errors), len(result.failures)))

      header = '=' * len(final_output)
      stderr.writeln(header)
      stderr.writeln(final_output)
      stderr.writeln(header)

      if result.wasSuccessful():
        return 0
      else:
        return 1

    # Run inside absl.app.run to ensure flags parsing is done.
    return app.run(main)


class BinaryDistribution(Distribution):
  """This class is needed in order to create OS specific wheels."""

  def has_ext_modules(self):
    return False

# TODO(joelshor): Maybe someday, when TF-GAN grows up, we can have our
# description be a `README.md` like `tensorflow_probability`.
DOCLINES = __doc__.split('\n')

setup(
    name=project_name,
    version=version,
    description=DOCLINES[0],
    long_description='\n'.join(DOCLINES[2:]),
    author='Google Inc.',
    author_email='packages@tensorflow.org',
    url='http://github.com/tensorflow/gan',
    license='Apache 2.0',
    packages=find_packages(),
    zip_safe=False,
    distclass=BinaryDistribution,
    cmdclass={
        'test': Test,
        'pip_pkg': InstallCommandBase,
    },
    install_requires=[
        'tensorflow_hub>=0.2',
        'tensorflow_probability>=0.7',
    ],
    extras_require={
        'tf': ['tensorflow>=1.12'],
        'tensorflow-datasets': ['tensorflow-datasets>=0.5.0'],
    },
    classifiers=[
        'Development Status :: 3 - Alpha',
        'Intended Audience :: Developers',
        'Intended Audience :: Education',
        'Intended Audience :: Science/Research',
        'License :: OSI Approved :: Apache Software License',
        'Programming Language :: Python :: 3',
        'Programming Language :: Python :: 3.4',
        'Programming Language :: Python :: 3.5',
        'Programming Language :: Python :: 3.6',
        'Topic :: Scientific/Engineering',
        'Topic :: Scientific/Engineering :: Artificial Intelligence',
        'Topic :: Software Development',
        'Topic :: Software Development :: Libraries',
        'Topic :: Software Development :: Libraries :: Python Modules',
    ],
    keywords='tensorflow GAN generative model machine learning',
)