# Copyright 2013-2024 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

import os
import unittest.mock
import xml.etree.ElementTree as ET

from spack import *
from spack.util.environment import EnvironmentModifications
import spack.build_environment

import importlib
build_brainscales = importlib.import_module("spack.pkg.ebrains-spack-builds.build_brainscales")


class Hxtorch(build_brainscales.BuildBrainscales):
    """hxtorch --- a PyTorch-based toplevel for the BrainScaleS-2 neuromorphic hardware systems"""

    homepage = "https://github.com/electronicvisions/hxtorch"
    # This repo provides a custom waf binary used for the build below
    git      = "https://github.com/electronicvisions/pynn-brainscales.git"

    maintainers = ['emuller']

    version('8.0-a5',         tag='hxtorch-8.0-a5')
    version('8.0-a4',         tag='hxtorch-8.0-a4')
    version('8.0-a3',         tag='hxtorch-8.0-a3')
    version('8.0-a2',         tag='hxtorch-8.0-a2')
    version('8.0-a1',         tag='hxtorch-8.0-a1')
    version('7.0-rc1-fixup3', tag='hxtorch-7.0-rc1-fixup3')
    version('7.0-rc1-fixup2', tag='hxtorch-7.0-rc1-fixup2')
    version('7.0-rc1-fixup1', branch='waf')

    deps_hxtorch_core = [
        # compiler for the BrainScaleS-2 embedded processor ("PPU"); needed for
        # building/linking, at runtime and for testing
        ('oppulance@8.0-a5', { "when":'@8.0-a5', "type": ('build', 'link', 'run', 'test') } ),
        ('oppulance@8.0-a4', { "when":'@8.0-a4', "type": ('build', 'link', 'run', 'test') } ),
        ('oppulance@8.0-a3', { "when":'@8.0-a3', "type": ('build', 'link', 'run', 'test') } ),
        ('oppulance@8.0-a2', { "when":'@8.0-a2', "type": ('build', 'link', 'run', 'test') } ),
        ('oppulance@8.0-a1', { "when":'@8.0-a1', "type": ('build', 'link', 'run', 'test') } ),
        ('oppulance@7.0-rc1-fixup3', { "when":'@7.0-rc1-fixup3', "type": ('build', 'link', 'run', 'test') } ),
        ('oppulance@7.0-rc1-fixup2', { "when":'@7.0-rc1-fixup2', "type": ('build', 'link', 'run', 'test') } ),
        ('oppulance@7.0-rc1-fixup1', { "when":'@7.0-rc1-fixup1', "type": ('build', 'link', 'run', 'test') } ),

        # host software dependencies
        ('bitsery', { "type": ('build', 'link', 'run', 'test') } ),
        ('binutils+gold+ld+plugins', { "type": ('build', 'link', 'run') } ), # specialize
        ('boost@1.69.0: +graph+icu+mpi+numpy+coroutine+context+filesystem+python+serialization+system+thread+program_options cxxstd=17', { "type": ('build', 'link', 'run', 'test') } ),
        ('cereal', { "type": ('build', 'link', 'run', 'test') } ),
        ('cppcheck', { "type": ('build', 'link', 'run') } ),
        ('genpybind@ebrains', { "type": ('build', 'link') } ),
        ('gflags', { "type": ('build', 'link', 'run') } ),
        ('googletest@1.11.0:+gmock', { "type": ('build', 'link', 'run') } ), # variadic templates needed
        ('inja', { "type": ('build', 'link', 'run', 'test')  } ),# template engine for PPU source jit generation
        ('intel-tbb', { "type": ('build', 'link', 'run') } ), # ppu gdbserver
        ('libelf', { "type": ('build', 'link', 'run') } ),
        ('liblockfile', { "type": ('build', 'link', 'run') } ),
        ('log4cxx@0.12.1:1.0', { "when": "@:8.0-a3", "type": ('build', 'link', 'run') } ),
        ('log4cxx@1.1: +events_at_exit', { "when": "@8.0-a4:", "type": ('build', 'link', 'run') } ),
        ('pkgconfig', { "type": ('build', 'link', 'run') } ),
        ('psmisc', { "type": ('run', 'test') } ),
        ('python@3.7.0:', { "type": ('build', 'link', 'run') } ), # BrainScaleS-2 only supports Python >= 3.7
        ('py-h5py', { "type": ('build', 'link', 'run') } ), # PyNN tests need it
        ('py-matplotlib', { "type": ('build', 'link', 'run') } ),
        ('py-networkx', { "type": ('build', 'link', 'run') } ),
        ('py-nose', { "type": ('build', 'link', 'run') } ),
        ('py-numpy', { "type": ('build', 'link', 'run') } ),
        ('py-pybind11', { "type": ('build', 'link', 'run') } ),
        ('py-pybind11-stubgen', { "type": ('build', 'link', 'run') } ),
        ('py-pycodestyle', { "type": ('build', 'link', 'run') } ),
        ('py-pyelftools', { "type": ('build', 'link', 'run') } ),
        ('py-pylint', { "type": ('build', 'link', 'run') } ),
        ('py-torch@1.9.1:', { "type": ('build', 'link', 'run', 'test') } ),
        ('py-torchvision', { "type": ('run') } ), # for demos
        ('py-pyyaml', { "type": ('build', 'link', 'run') } ),
        ('py-scipy', { "type": ('build', 'link', 'run') } ),
        ('py-sqlalchemy', { "type": ('build', 'link', 'run') } ),
        ('util-linux', { "type": ('build', 'link', 'run') } ),
        ('yaml-cpp+shared', { "type": ('build', 'link', 'run') } ),
    ]

    for dep, dep_kw in deps_hxtorch_core:
        depends_on(dep, **dep_kw)

    extends('python')

    patch("include-SparseTensorUtils.patch", when="@:8.0-a5")

    def install_test(self):
        with working_dir('spack-test', create=True):
            old_pythonpath = os.environ.get('PYTHONPATH', '')
            os.environ['PYTHONPATH'] = ':'.join([str(self.prefix.lib), old_pythonpath])
            bash = which("bash")
            # ignore segfaults for now (exit code 139)
            bash('-c', '(python -c "import hxtorch; print(hxtorch.__file__)" || ( test $? -eq 139 && echo "segfault")) || exit $?')