# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
#
# 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.
import os
from setuptools import setup, Extension, find_packages
import subprocess

import torch
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, ROCM_HOME

from typing import Optional, Union
import subprocess
from pathlib import Path


version_dependent_macros = [
    '-DVERSION_GE_1_1',
    '-DVERSION_GE_1_3',
    '-DVERSION_GE_1_5',
]

extra_cuda_flags = [
    '-std=c++14',
    #'-maxrregcount=50',
    '-U__CUDA_NO_HALF_OPERATORS__',
    '-U__CUDA_NO_HALF_CONVERSIONS__'
]

cc_flag = ['-DAMDGPU_TARGETS', 'gfx900;gfx906;gfx926']

extra_cuda_flags += cc_flag


def get_sha(root: Union[str, Path]) -> str:
    try:
        return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=root).decode('ascii').strip()
    except Exception:
        return 'Unknown'


def get_abi():
    try:
        command = "echo '#include <string>' | gcc -x c++ -E -dM - | fgrep _GLIBCXX_USE_CXX11_ABI" 
        result = subprocess.run(command, shell=True, capture_output=True, text=True) 
        output = result.stdout.strip() 
        abi = "abi" + output.split(" ")[-1]
        return abi
    except Exception:
        return 'abiUnknown'
    
    
def get_version_add(sha: Optional[str] = None) -> str:
    openfold_root = os.path.dirname(os.path.abspath(__file__))
    add_version_path = os.path.join(os.path.join(openfold_root, "openfold"), "version.py")
    if sha != 'Unknown':
        if sha is None:
            sha = get_sha(openfold_root)
        version = 'git' + sha[:7]

    # abi version
    version += "." + get_abi()

    # dtk version
    if os.getenv("ROCM_PATH"):
        rocm_path = os.getenv('ROCM_PATH', "")
        rocm_version_path = os.path.join(rocm_path, '.info', "rocm_version")
        with open(rocm_version_path, 'r',encoding='utf-8') as file:
            lines = file.readlines()
        rocm_version=lines[0][:-2].replace(".", "")
        version += ".dtk" + rocm_version

    # torch version
    version += ".torch" + torch.__version__[:4]

    with open(add_version_path, encoding="utf-8", mode="w") as file:
        file.write("__version__='1.0.1'\n")
        file.write("__dcu_version__='1.0.1+{}'\n".format(version))
    file.close()


def get_version():
    get_version_add()
    version_file = 'version.py'
    with open(version_file, encoding='utf-8') as f:
        exec(compile(f.read(), version_file, 'exec'))
    return locals()['__dcu_version__']

setup(
    name='openfold',
    #version='1.0.1',
    version=get_version(),
    description='A PyTorch reimplementation of DeepMind\'s AlphaFold 2',
    author='Gustaf Ahdritz & DeepMind',
    author_email='gahdritz@gmail.com',
    license='Apache License, Version 2.0',
    url='https://github.com/aqlaboratory/openfold',
    packages=find_packages(exclude=["tests", "scripts"]),
    include_package_data=False,
    package_data={
        "openfold": ['utils/kernel/csrc/*'],
        "": ["resources/stereo_chemical_props.txt"]
    },
    ext_modules=[CUDAExtension(
        name="attn_core_inplace_cuda",
        sources=[
            "openfold/utils/kernel/csrc/softmax_cuda.cpp",
            "openfold/utils/kernel/csrc/softmax_cuda_kernel.cu",
        ],
        include_dirs=[
            os.path.join(
                os.path.dirname(os.path.abspath(__file__)),
                'openfold/utils/kernel/csrc/'
            )
        ],
        extra_compile_args={
            'cxx': ['-O3'] + version_dependent_macros,
            'hipcc': (
                ['-O3'] + version_dependent_macros + extra_cuda_flags
            ),
        }
    )],
    cmdclass={'build_ext': BuildExtension},
    classifiers=[
        'License :: OSI Approved :: Apache Software License',
        'Operating System :: POSIX :: Linux',
        'Programming Language :: Python :: 3.7,'
        'Topic :: Scientific/Engineering :: Artificial Intelligence',
    ],
)