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

from pathlib import Path

from spack.package import *


class IconNwp(Package):
    """
    Recipe to build ICON-NWP using the LMU gitlab repository.
    """

    git = "ssh://git@gitlab.dkrz.de/icon/icon-nwp.git"
    homepage = "https://code.mpimet.mpg.de/projects/iconpublic"
    # maintainers("oriol.tinto")

    # Version w2w-B6 points to a different repository and branch)
    version("w2w-B6", git="ssh://git@gitlab.physik.uni-muenchen.de/w2w/icon-w2w.git", branch="icon-w2w/icon-nwp-B6")
    # FIXME: The ugly configuration system in older icon versions prevents us from using this package with it.
    # version("2.5.0-nwp3", branch="icon-nwp/op-release-2.5.0-nwp3")
    version("2.6.4-nwp3", branch="icon-nwp/op-release-2.6.4-nwp3")
    version("2.6.5-nwp0", branch="icon-nwp/op-release-2.6.5-nwp0")
    version("master", branch="master")
    version("psp", branch="icon-nwp/icon-nwp-psp")

    variant("lmu", default=False, description="if git.url and submodules should be patched to use the LMU mirrors")
    variant("plexrt", default=False, description="add 3D radiation support with TenStream")
    variant("petsc", default=False, description="add PETSc support")

    # Dependencies
    depends_on("mpi")
    depends_on("netcdf-c")
    depends_on("netcdf-fortran")
    depends_on("eccodes+fortran")
    depends_on("libxml2")

    depends_on("petsc", when="+petsc")
    depends_on("petsc", when="+plexrt")
    depends_on("tenstream", when="+plexrt")

    # Openblas? best way of doing it?
    depends_on("openblas", when="%gcc")
    depends_on("intel-mkl", when="%intel")

    # Extra dependencies for B6, including yaml, and the hdf5 filters for compression.
    depends_on("libyaml", when="@w2w-B6")
    depends_on("hdf5-blosc", when="@w2w-B6")
    depends_on("h5z-zfp", when="@w2w-B6")
    depends_on("sz~hdf5", when="@w2w-B6")
    depends_on("sz3~hdf5", when="@w2w-B6")

    phases = ["configure", "build", "install"]

    def do_fetch(self, mirror_only=False):
        if "+lmu" in self.spec:
            self.fetcher[0].url = self.git = "ssh://git@gitlab.physik.uni-muenchen.de/w2w/icon.git"
        super(IconNwp, self).do_fetch(mirror_only)

    def patch(self):
        # Run git submodule update
        git = which("git")
        if "+lmu" in self.spec:
            self.patch_submodules()
        git("submodule", "update", "--init", "--recursive")

    def patch_submodules(self):
        """
        Because of the lack of access rights to the original submodule repositories,
        we patch the gitmodules file to point to a different mirror.
        """
        git_submodules_file = Path().cwd() / ".gitmodules"

        git_mirror = "git@gitlab.lrz.de:dkrz-mirror"

        git_modules_patch = f"""
        [submodule "externals/mtime"]
                path = externals/mtime
                url = {git_mirror}/libmtime.git
        [submodule "externals/jsbach"]
                path = externals/jsbach
                url = {git_mirror}/jsbach.git
        [submodule "externals/yac"]
                path = externals/yac
                url = {git_mirror}/YAC.git
        [submodule "externals/self"]
                path = externals/self
                url = {git_mirror}/libself.git
        [submodule "externals/tixi"]
                path = externals/tixi
                url = {git_mirror}/libtixi.git
        [submodule "externals/yaxt"]
                path = externals/yaxt
                url = {git_mirror}/yaxt.git
        [submodule "externals/rte-rrtmgp"]
                path = externals/rte-rrtmgp
                url = https://github.com/earth-system-radiation/rte-rrtmgp.git
        [submodule "externals/cub"]
                path = externals/cub
                url = https://github.com/NVlabs/cub.git
        [submodule "externals/omni-xmod-pool"]
                path = externals/omni-xmod-pool
                url = https://github.com/claw-project/omni-xmod-pool.git
        [submodule "externals/cdi"]
                path = externals/cdi
                url = {git_mirror}/libcdi.git
        [submodule "externals/sct"]
                path = externals/sct
                url = {git_mirror}/sct.git
        [submodule "externals/ecrad"]
                path = externals/ecrad
                url = {git_mirror}/libecrad.git
        [submodule "externals/dace_icon"]
                path = externals/dace_icon
                url = {git_mirror}/dace-icon-interface.git
        [submodule "externals/emvorado"]
                path = externals/emvorado
                url = {git_mirror}/emvorado-for-icon.git
        [submodule "utils/mkexp"]
                path = utils/mkexp
                url = https://git.mpimet.mpg.de/public/mkexp
        [submodule "externals/art"]
                path = externals/art
                url = {git_mirror}/art.git
        [submodule "externals/ppm"]
                path = externals/ppm
                url = https://gitlab.dkrz.de/jahns/ppm.git
        [submodule "externals/probtest"]
                path = externals/probtest
                url = {git_mirror}/cscs-sw_probtest.git
        """

        # Replace the content of the original file with the patch
        with git_submodules_file.open("w") as out_f:
            out_f.write(git_modules_patch)

    def setup_build_environment(self, env):
        spec = self.spec

        # Some environment variables to set
        env_variables_to_set = {
            "CC": spec["mpi"].mpicc,
            "FC": spec["mpi"].mpifc,
            "F77": spec["mpi"].mpif77,
        }
        for variable, value in env_variables_to_set.items():
            env.set(variable, value)

    def configure(self, spec, prefix):
        spec = self.spec
        print(spec["mpi"].mpifc)

        libs = [
            f"-L{spec['netcdf-c'].prefix.lib} -lnetcdf ",  # netcdf-c libs
            f"-L{spec['netcdf-fortran'].prefix.lib} -lnetcdff",  # netcdf-fortran libs
            f"-L{spec['eccodes'].prefix.lib} -leccodes_f90 -leccodes",
            f"-L{spec['libxml2'].prefix.lib} -lxml2",  # XML2 libs
        ]

        if self.spec.satisfies("%gcc"):
            libs.append("-lopenblas")
        elif self.spec.satisfies("%intel"):
            libs.append("-qmkl=sequential")

        if self.spec.version == Version("w2w-B6"):
            libs.append("-lyaml")

        mtune = "generic"
        INCLUDES = [f"-I{spec['libxml2'].prefix}/include/libxml2", ]
        CPPFLAGS = []

        if "+plexrt" in self.spec:
            CPPFLAGS.append("-DHAVE_PLEXRT")
            INCLUDES.append(f"-I{spec['tenstream'].prefix}/include")
            libs.append(f"-L{spec['tenstream'].prefix.lib} -ltenstream")

        if "+petsc" in self.spec:
            CPPFLAGS.append("-DHAVE_PETSC")
            INCLUDES.append(f"-I{spec['petsc'].prefix}/include")
            libs.append(f"-L{spec['petsc'].prefix.lib} -lpetsc")


        options = [
            f"CC={spec['mpi'].mpicc}",
            f"FC={spec['mpi'].mpifc}",
            f"CFLAGS=-g -mpc64 {' '.join(INCLUDES)}",
            f"ICON_CFLAGS=-O3 -g -mtune={mtune}",
            f"ICON_BUNDLED_CFLAGS=-O3 -mtune={mtune}",
            f"CPPFLAGS={' '.join(CPPFLAGS)}",
            "FCFLAGS=-std=legacy -fmodule-private -fimplicit-none -fmax-identifier-length=63 -Wall -Wcharacter-truncation -Wconversion -Wunderflow -Wunused-parameter -Wno-surprising -fall-intrinsics -g -mpc64 -w",
            "ICON_FCFLAGS=-fbacktrace -fbounds-check -fstack-protector-all -finit-real=nan -finit-integer=-2147483648 -finit-character=127 -w -O2",
            f"ICON_OCEAN_FCFLAGS=-O3 -mtune={mtune}",
            f"LDFLAGS={' '.join(libs)}",
            f"LIBS={' '.join(libs)}",
            f"--prefix={prefix}",
            f"--enable-grib2",
        ]

        # For some reason there's a problem with OpenMPI with gcc@11.3.0 which makes the configuration fail.
        if self.spec.compiler.name == "gcc" and self.spec.compiler.version == Version("11.3.0"):
            options.append("--enable-mpi-checks=no")
        configure(*options)

    def build(self, spec, prefix):
        make()

    def install(self, spec, prefix):
        make("install")

        # Create extra folders with the data and the docs
        self.make_extra_folders(prefix)

    def setup_run_environment(self, env):
        env.set("ICON_BASE_PATH", self.spec.prefix)
        env.set("ICON_DATA_PATH", self.spec.prefix.join("data"))
        env.set("ICON_DOCS_PATH", self.spec.prefix.join("doc"))

    def make_extra_folders(self, prefix):
        mkdir = which("mkdir")
        rsync = which("rsync")
        curl = which("curl")

        # copy executables and documentation
        mkdir("-p", f"{prefix}/data")
        mkdir("-p", f"{prefix}/doc")
        mkdir("-p", f"{prefix}/run")
        rsync("-av", "data/", f"{prefix}/data")
        rsync("-av", "run/", f"{prefix}/run")
        rsync("-av", "--include='*.pdf'", "--exclude='*.*'", "doc/", f"{prefix}/doc/")
        curl(
            "https://code.mpimet.mpg.de/attachments/download/19568/ICON_tutorial_2019.pdf",
            "--output",
            f"{prefix}/doc/ICON_tutorial_2019.pdf",
        )
        curl(
            "http://www.cosmo-model.org/content/model/documentation/core/emvorado_userguide.pdf",
            "--output",
            f"{prefix}/doc/emvorado_userguide.pdf",
        )