# Copyright 2013-2021 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)

# ----------------------------------------------------------------------------
# If you submit this package back to Spack as a pull request,
# please first remove this boilerplate and all FIXME comments.
#
# This is a template package file for Spack.  We've put "FIXME"
# next to all the things you'll want to change. Once you've handled
# them, you can save this file and test your package like this:
#
#     spack install icon
#
# You can edit this file again by typing:
#
#     spack edit icon
#
# See the Spack documentation for more information on packaging.
# ----------------------------------------------------------------------------

from spack import *
import os

class Icon(AutotoolsPackage):
    """ISC22 ICON package."""

    homepage = "https://hpcadvisorycouncil.atlassian.net/wiki/spaces/HPCWORKS/pages/2792161313/Getting+started+with+ICON+for+ISC22+SCC"
    url      = "http://localhost:31415/icon-2.6.4.tar.gz"

    maintainers = ['spedoske','victoryang00']

    version('2.6.4', sha256='1860028836d0894224ce301c3d0cb27a899548823267b08bf5c97ae696c3758d')

    depends_on('libxml2')
    depends_on('eccodes')
    depends_on('intel-oneapi-mkl')
    #if self.compiler.name == 'nvhpc':
    #depends_on('intel-oneapi-mpi')
    depends_on('netcdf-fortran')
    depends_on('netcdf-c')
    depends_on('hdf5+mpi+fortran+hl+szip')
    depends_on('zlib')
    depends_on('libszip')
    depends_on('mpi')
    depends_on('cuda')
    #depends_on('nvhpc')
    depends_on('serialbox')

    def configure_args(self):
        args = [#'--enable-art',
                '--enable-coupling',
                '--enable-serialization',
                '--enable-emvorado',
                '--enable-grib2',
                '--disable-yaxt',
                '--enable-sct',
                '--enable-ecrad',
                '--enable-rte-rrtmgp',
                '--enable-mpi',
                '--enable-gpu',
                '--disable-explicit-fpp'
                ]

        # if self.compiler.name == 'nvhpc':
        #    args.append('--enable-gpu')
        # else:
        #    args.append('--disable-gpu')

        #from print import pprint
        #pprint(vars(self.spec['nvhpc']),depth=4)

        #print(os.environ)

        os.environ['CC'] =  os.environ['MPICC']
        os.environ['CXX'] = os.environ['MPICXX']
        os.environ['F77'] = os.environ['MPIF77']
        os.environ['F90'] = os.environ['MPIF90']
        os.environ['FC'] =  os.environ['MPIF90']


        os.environ['MPICH_CC'] =  os.environ['MPICC']
        os.environ['MPICH_CXX'] = os.environ['MPICXX']
        os.environ['MPICH_F77'] = os.environ['MPIF77']
        os.environ['MPICH_F90'] = os.environ['MPIF90']
        os.environ['MPICH_FC'] =  os.environ['MPIF90']

        #nvhpc_bin = '/home/coffee/spack/opt/spack/linux-ubuntu20.04-zen2/oneapi-2021.4.0/nvhpc-21.9-ymnvxvr4e2q7475djr4lig7cbucg54uv/Linux_x86_64/21.9/compilers/bin'

        #os.environ['OMPI_CC'] =  os.path.join(nvhpc_bin,'nvc')
        #os.environ['OMPI_CXX'] = os.path.join(nvhpc_bin,'nvc++')
        #os.environ['OMPI_F77'] = os.path.join(nvhpc_bin,'nvfortran')
        #os.environ['OMPI_FC'] =  os.path.join(nvhpc_bin,'nvfortran')


        CPPFLAGS=["CPPFLAGS="]
        FCFLAGS=["FCFLAGS="]
        LDFLAGS=["LDFLAGS="]
        LIBS=["LIBS="]

        # zlib
        LDFLAGS.append('-L{}'.format(os.path.join(self.spec['zlib'].prefix,'lib')))
        LIBS.append('-lz')

        # hdf5
        CPPFLAGS.append('-I{}'.format(os.path.join(self.spec['hdf5'].prefix,'include')))
        FCFLAGS.append('-I{}'.format(os.path.join(self.spec['hdf5'].prefix,'include')))
        LDFLAGS.append('-L{}'.format(os.path.join(self.spec['hdf5'].prefix,'lib')))
        LIBS.append('-lhdf5hl_fortran -lhdf5_fortran -lhdf5')


        # netcdf-fortran
        FCFLAGS.append('-I{}'.format(os.path.join(self.spec['netcdf-fortran'].prefix,'include')))
        LDFLAGS.append('-L{}'.format(os.path.join(self.spec['netcdf-fortran'].prefix,'lib')))
        LIBS.append('-lnetcdff')

        # netcdf-c
        CPPFLAGS.append('-I{}'.format(os.path.join(self.spec['netcdf-c'].prefix,'include')))
        LDFLAGS.append('-L{}'.format(os.path.join(self.spec['netcdf-c'].prefix,'lib')))
        LIBS.append('-lnetcdf')

        # blas and lapack
        CPPFLAGS.append('-I{}'.format(os.path.join(self.spec['intel-oneapi-mkl'].prefix,'include')))
        LIBS.append('-lmkl_intel_lp64 -lmkl_intel_thread -lmkl_core -liomp5 -lpthread -lm -ldl -lmpi -lxml2')

        # xml2
        CPPFLAGS.append('-I{}'.format(os.path.join(self.spec['libxml2'].prefix,'include/libxml2/')))

        # eccodes c/fc
        CPPFLAGS.append('-I{}'.format(os.path.join(self.spec['eccodes'].prefix,'include')))
        LDFLAGS.append('-L{}'.format(os.path.join(self.spec['eccodes'].prefix,'lib')))
        FCFLAGS.append('-I{}'.format(os.path.join(self.spec['eccodes'].prefix,'include')))
        LIBS.append('-leccodes -leccodes_f90')

        # serialbox
        LDFLAGS.append('-L{}'.format(os.path.join(self.spec['serialbox'].prefix,'lib')))
        FCFLAGS.append('-I{}'.format(os.path.join(self.spec['serialbox'].prefix,'include')))
        LIBS.append('-lSerialboxFortran')


        # LDFLAGS.append('-L{}'.format(os.path.join(self.spec['cuda'].prefix,'lib64')))
        # LIBS.append('-lcudart')

        # gpu
        #LDFLAGS.append('-L/usr/lib')
        #LIBS.append('-lstdc++ -lstdc++ -L/usr/lib')

        # gpu
        if self.compiler.name == 'nvhpc':
           LDFLAGS.append('-L{}'.format(os.path.join(self.spec['cuda'].prefix,'lib64')))
           LDFLAGS.append('-L/lib/x86_64-linux-gnu/ -lstdc++')
           LDFLAGS.append('-I{}'.format(os.path.join(self.spec['openmpi'].prefix,'lib')))
           FCFLAGS.append('-I{}'.format(os.path.join(self.spec['openmpi'].prefix,'lib')))
           LIBS.append('-lcudart')

        args.append(' '.join(CPPFLAGS))
        args.append(' '.join(FCFLAGS))
        args.append(' '.join(LDFLAGS))
        args.append(' '.join(LIBS))
        args.append('SB2PP='+self.spec['serialbox'].prefix+'/python/pp_ser/pp_ser.py')
        if self.compiler.name == 'nvhpc':
            depends_on('nvhpc')
            #args.append('MPI_LAUNCH='+self.spec['openmpi'].prefix+'/bin/mpirun')
            args.append("MPI_LAUNCH=/usr/bin/mpirun")
            args.append('NVCC=nvcc -std=c++11  -ccbin {} -allow-unsupported-compiler'.format(self.compiler.cxx))
        else:
            depends_on('intel-oneapi-mpi')
            args.append('MPI_LAUNCH='+self.spec['intel-oneapi-mpi'].prefix+'/mpi/2021.4.0/bin/mpirun')
            args.append('NVCC=nvcc -std=c++11  -ccbin {} -allow-unsupported-compiler'.format(self.compiler.cc))
        return args

    def setup_run_environment(self, env):
        env.set('PYTHONPATH', join_path(self.spec['serialbox'].prefix,'python'))