# 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

from spack.package import *
from spack_repo.builtin.packages.amrex.package import Amrex as SpackAmrex


class Amrex(SpackAmrex):

    version("18.07", sha256="2e8d3a91c5d972e3cebb88d10c4c92112459c88c2342c5a63337f3110bdbff13")

    # Config options
    variant(
        "dimensions",
        default="3",
        values=("1", "2", "3"),
        multi=False,
        description="Dimensionality",
        when="@18.07",
    )

    variant("fbaselib", default=False, description="Fbaselib", when="@18.07")
    variant("fpe", default=False, description="FPE", when="@18.07")

    patch("AMReX_ParallelDescriptor.patch", when="@18.07+mpi")

    @when("@18.07")
    def cmake_args(self):
        args = [
            self.define_from_variant("DIM", "dimensions"),
            self.define_from_variant("BUILD_SHARED_LIBS", "shared"),
            self.define_from_variant("ENABLE_MPI", "mpi"),
            self.define_from_variant("ENABLE_OMP", "openmp"),
            self.define_from_variant("ENABLE_FORTRAN_INTERFACES", "fortran"),
            self.define_from_variant("ENABLE_EB", "eb"),
            self.define_from_variant("ENABLE_LINEAR_SOLVERS", "linear_solvers"),
            self.define_from_variant("ENABLE_AMRDATA", "amrdata"),
            self.define_from_variant("ENABLE_PARTICLES", "particles"),
            self.define_from_variant("ENABLE_SUNDIALS", "sundials"),
            self.define_from_variant("ENABLE_HDF5", "hdf5"),
            self.define_from_variant("ENABLE_HYPRE", "hypre"),
            self.define_from_variant("ENABLE_PETSC", "petsc"),
            self.define_from_variant("ENABLE_CUDA", "cuda"),
            self.define_from_variant("ENABLE_PIC", "pic"),
            self.define_from_variant("ENABLE_FBASELIB", "fbaselib"),
            self.define_from_variant("ENABLE_FPE", "fpe"),
            "-DENABLE_DP=%s" % self.spec.variants["precision"].value.upper(),
            "-DENABLE_DP_PARTICLES=%s" % self.spec.variants["precision"].value.upper(),
        ]

        if "+linear_solvers" in self.spec:
            args.append("-DENABLE_LINEAR_SOLVERS_LEGACY=1")

        if self.spec.satisfies("%fj"):
            args.append("-DCMAKE_Fortran_MODDIR_FLAG=-M")

        if "+cuda" in self.spec:
            cuda_arch = self.spec.variants["cuda_arch"].value
            args.append("-DCUDA_ARCH=" + self.get_cuda_arch_string(cuda_arch))

        return args