"""
`module  JAC.Hfs`
    ... a submodel of JAC that contains all methods for computing HFS A and B coefficients, hyperfine sublevel representations, etc.
"""
module Hfs

    using Printf, ..AngularMomentum, ..Basics,  ..Defaults, ..InteractionStrength, ..ManyElectron, ..Radial, ..Nuclear, ..TableStrings

    """
    `struct  Hfs.InteractionMatrix`  ... defines a type for storing the T^1 and T^2 interaction matrices for a given basis.

        + calcT1   ::Bool               ... true, if the matrixT1 has been calculated and false otherwise.
        + calcT2   ::Bool               ... true, if the matrixT2 has been calculated and false otherwise.
        + matrixT1 ::Array{Float64,2}   ... T1 interaction matrix
        + matrixT2 ::Array{Float64,2}   ... T2 interaction matrix

    """
    struct InteractionMatrix
        calcT1     ::Bool
        calcT2     ::Bool
        matrixT1   ::Array{Float64,2}
        matrixT2   ::Array
    end

    """
    `Hfs.InteractionMatrix()`  ... constructor for an `empty` instance of InteractionMatrix.
    """
    function InteractionMatrix()
        InteractionMatrix(false, false, zeros(2,2), zeros(2,2))
    end

    # `Base.show(io::IO, im::Hfs.InteractionMatrix)`  ... prepares a proper printout of the variable InteractionMatrix.
    function Base.show(io::IO, im::Hfs.InteractionMatrix)
        println(io, "calcT1:           $(im.calcT1)  ")
        println(io, "calcT1:           $(im.calcT1)  ")
        println(io, "matrixT1:         $(im.matrixT1)  ")
        println(io, "matrixT2:         $(im.matrixT2)  ")
    end

    """
    `struct  Hfs.IJF_Vector`  ... defines a type for a IJF-coupled basis vector, here based on an ASF.

        + I        ::AngularJ64   ... Nuclear spin I.
        + F        ::AngularJ64   ... Total angular momentum F
        + levelJ   ::Level        ... Atomic level with well-defined total (electronic) angular momentum J

    """
    struct IJF_Vector
        I          ::AngularJ64
        F          ::AngularJ64
        levelJ     ::Level
    end

    """
    `Hfs.IJF_Vector()`  ... constructor for an `empty` instance of IJF_Vector.
    """
    function IJF_Vector()
        IJF_Vector(AngularJ64(0), AngularJ64(0), Level())
    end

    # `Base.show(io::IO, ijfVector::Hfs.IJF_Vector)`  ... prepares a proper printout of the variable ijfVector.
    function Base.show(io::IO, ijfVector::Hfs.IJF_Vector)
        println(io, "I:           $(ijfVector.I)  ")
        println(io, "F:           $(ijfVector.F)  ")
        println(io, "levelJ:      $(ijfVector.levelJ)  ")
    end

   """
    `struct  Hfs.IJF_Basis`  ... defines a type for a IJF-coupled basis.

        + vectors   ::Array{IJF_Vector,1}   ... List of IJF_Vectors that form the basis.
        + basisJ    ::Basis                 ... Electronic basis that allows access to the electronic orbitals and CSF basis.
    """
    struct IJF_Basis
        vectors     ::Array{IJF_Vector,1}
        basisJ      ::Basis
    end

    """
   `Hfs.IJF_Basis()`  ... constructor for an `empty` instance of IJF_Basis.
    """
    function IJF_Basis()
        IJF_Basis(IJF_Vector[], Basis())
    end

    # `Base.show(io::IO, ijfBasis::Hfs.IJF_Basis)`  ... prepares a proper printout of the variable ijfBasis.
    function Base.show(io::IO, ijfBasis::Hfs.IJF_Basis)
        println(io, "levelFs:           $(ijfBasis.levelFs)  ")
        println(io, "basisJ:            $(ijfBasis.basisJ)  ")
    end

    """
    `struct  Hfs.IJF_Level`  ... defines a type for a IJF-coupled level.

        + I              ::AngularJ64       ... Nuclear spin I.
        + F              ::AngularJ64       ... Total angular momentum F.
        + M              ::AngularM64       ... Total projection M, only defined if a particular sublevel is referred to.
        + parity         ::Parity           ... Parity of the level which corresponds to the electronic system.
        + energy         ::Float64          ... energy
        + hasStateRep    ::Bool             ... Determines whether this IJF-level 'point' to a physical representation of
                                                an state (i.e. an basis and corresponding mixing coefficient vector), or
                                                just to empty instances, otherwise.
        + basis          ::IJF_Basis        ... basis for this level
        + mc             ::Vector{Float64}  ... Vector of mixing coefficients w.r.t basis.
    """
    struct IJF_Level
        I                ::AngularJ64
        F                ::AngularJ64
        M                ::AngularM64
        parity           ::Parity
        energy           ::Float64
        hasStateRep      ::Bool
        basis            ::IJF_Basis
        mc               ::Vector{Float64}
    end

    """
    `Hfs.IJF_Level()`  ... constructor for an `empty` instance of IJF_Level.
    """
    function IJF_Level()
        IJF_Level(AngularJ64(0), AngularJ64(0), AngularJ64(0), AngularM64(0), Basics.plus, 0., false, IJF_Basis(), Float64[])
    end

    # `Base.show(io::IO, ijfLevel::Hfs.IJF_Level)`  ... prepares a proper printout of the variable ijfLevel::Hfs.IJF_Level.
    function Base.show(io::IO, ijfLevel::Hfs.IJF_Level)
        println(io, "I:              $(ijfLevel.I)  ")
        println(io, "J:              $(ijfLevel.J)  ")
        println(io, "F:              $(ijfLevel.F)  ")
        println(io, "M:              $(ijfLevel.M)  ")
        println(io, "parity:         $(ijfLevel.parity)  ")
        println(io, "energy:         $(ijfLevel.energy)  ")
        println(io, "hasStateRep:    $(ijfLevel.hasStateRep)  ")
        println(io, "basis:          $(ijfLevel.basis)  ")
        println(io, "mc:             $(ijfLevel.mc)  ")
    end

    """
    `struct  Hfs.IJF_Multiplet`  ... defines a type for a multiplet of IJF-coupled levels.

        + name     ::String                ... A name associated to the multiplet.
        + levelFs  ::Array{IJF_Level,1}    ... List of IJF-coupled levels.

    """
    struct IJF_Multiplet
        name       ::String
        levelFs    ::Array{IJF_Level,1}
    end

    """
    `Hfs.IJF_Multiplet()`  ... constructor for an `empty` instance of Hfs.IJF_Multiplet.
    """
    function IJF_Multiplet()
        IJF_Multiplet("", IJF_Level[])
    end

    # `Base.show(io::IO, ijfMultiplet::Hfs.IJF_Multiplet)`  ... prepares a proper printout of the variable ijfMultiplet::Hfs.IJF_Multiplet.
    function Base.show(io::IO, ijfMultiplet::Hfs.IJF_Multiplet)
        println(io, "name:           $(ijfMultiplet.name)  ")
        println(io, "levelFs:        $(ijfMultiplet.levelFs)  ")
    end

    """
    `struct  Hfs.Outcome`
        ... defines a type to keep the outcome of a HFS computation, such as the HFS A and B coefficients as well
            other results.

        + Jlevel                    ::Level            ... Atomic level to which the outcome refers to.
        + AIoverMu                  ::Float64          ... HFS A * I / mu value.
        + BoverQ                    ::Float64          ... HFS B / Q value
        + amplitudeT1               ::Complex{Float64} ... T1 amplitude
        + amplitudeT2               ::Complex{Float64} ... T2 amplitude
        + nuclearI                  ::AngularJ64       ... nuclear spin
        + hasFlevels                ::Bool             ... True if the F-values and F-energies are specified explicitly.
        + Flevels                   ::Bool             ... Array to keep F-values and F-energies.
    """
    struct Outcome
        Jlevel                      ::Level
        AIoverMu                    ::Float64
        BoverQ                      ::Float64
        amplitudeT1                 ::Complex{Float64}
        amplitudeT2                 ::Complex{Float64}
        nuclearI                    ::AngularJ64
        hasFlevels                  ::Bool
        Flevels                     ::Bool
    end

    """
    `Hfs.Outcome()`  ... constructor for an `empty` instance of Hfs.Outcome for the computation of HFS properties.
    """
    function Outcome()
        Outcome(Level(), 0., 0., 0., 0., AngularJ64(0), false, false)
    end

    # `Base.show(io::IO, outcome::Hfs.Outcome)`  ... prepares a proper printout of the variable outcome::Hfs.Outcome.
    function Base.show(io::IO, outcome::Hfs.Outcome)
        println(io, "Jlevel:                    $(outcome.Jlevel)  ")
        println(io, "AIoverMu:                  $(outcome.AIoverMu)  ")
        println(io, "BoverQ:                    $(outcome.BoverQ)  ")
        println(io, "amplitudeT1:               $(outcome.amplitudeT1)  ")
        println(io, "amplitudeT2:               $(outcome.amplitudeT2)  ")
        println(io, "nuclearI:                  $(outcome.nuclearI)  ")
        println(io, "hasFlevels:                $(outcome.hasFlevels)  ")
        println(io, "Flevels:                   $(outcome.Flevels)  ")
    end

    """
    `struct  Settings`  ... defines a type for the details and parameters of computing HFS A and B coefficients.

        + calcT1                    ::Bool             ... True if T1-amplitudes (HFS A values) need to be calculated, and false otherwise.
        + calcT2                    ::Bool             ... True if T2-amplitudes (HFS B values) need to be calculated, and false otherwise.
        + calcNondiagonal           ::Bool             ... True if also (non-)diagonal hyperfine amplitudes are to be calculated and
                                                           printed, and false otherwise.
        + calcIJFexpansion          ::Bool             ... True if the selected atomic levels are to be represented in a IJF-coupled basis,
                                                           and false otherwise.
        + printBefore    ::Bool                        ... True if a list of selected levels is printed before the actual computations start.
        + printDeltaEF              ::Bool             ... True if the energy shift of E_F (with regard to E_J) is to be printed, and false otherwise.
        + levelSelection            ::LevelSelection   ... Specifies the selected levels, if any.
    """
    struct Settings
        calcT1                      ::Bool
        calcT2                      ::Bool
        calcNondiagonal             ::Bool
        calcIJFexpansion            ::Bool
        printBefore                 ::Bool
        printDeltaEF                ::Bool
        levelSelection              ::LevelSelection
    end

    """
    `Hfs.Settings(; calcT1::Bool=true,` calcT2::Bool=false, calcNondiagonal::Bool=false, calcIJFexpansion::Bool=false,
                        printBefore::Bool=false, printDeltaEF::Bool=false,
                        levelSelection::LevelSelection=LevelSelection())
        ... keyword constructor to overwrite selected value of Einstein line computations.
    """
    function Settings(; calcT1::Bool=true, calcT2::Bool=false, calcNondiagonal::Bool=false, calcIJFexpansion::Bool=false,
                        printBefore::Bool=false, printDeltaEF::Bool=false,
                        levelSelection::LevelSelection=LevelSelection())
        Settings(calcT1, calcT2, calcNondiagonal, calcIJFexpansion, printBefore, printDeltaEF, levelSelection)
    end

    # `Base.show(io::IO, settings::Hfs.Settings)`  ... prepares a proper printout of the variable settings::Hfs.Settings.
    function Base.show(io::IO, settings::Hfs.Settings)
        println(io, "calcT1:                   $(settings.calcT1)  ")
        println(io, "calcT2:                   $(settings.calcT2)  ")
        println(io, "calcNondiagonal:          $(settings.calcNondiagonal)  ")
        println(io, "calcIJFexpansion:         $(settings.calcIJFexpansion)  ")
        println(io, "printBefore:              $(settings.printBefore)  ")
        println(io, "printDeltaEF:             $(settings.printDeltaEF)  ")
        println(io, "levelSelection:           $(settings.levelSelection)  ")
    end

    """
    `Hfs.amplitude(kind::String, rLevel::Level, sLevel::Level, grid::Radial.Grid; printout::Bool=true)`
        ... to compute either the  T^(1) or T^(2) hyperfine amplitude <alpha_r J_r || T^(n)) || alpha_s J_s>
            for a given pair of levels. A value::ComplexF64 is returned.
    """
    function  amplitude(kind::String, rLevel::Level, sLevel::Level, grid::Radial.Grid; printout::Bool=true)
        #
        if  rLevel.parity != sLevel.parity   return( ComplexF64(0.) )   end
        nr = length(rLevel.basis.csfs);    ns = length(sLevel.basis.csfs);    matrix = zeros(ComplexF64, nr, ns)
        if  printout   printstyled("Compute hyperfine $(kind[1:5]) matrix of dimension $nr x $ns ... \n", color=:light_green)   end
        #
        for  r = 1:nr
            for  s = 1:ns
                me = 0.
                if  rLevel.basis.csfs[r].parity  != rLevel.parity    ||  sLevel.basis.csfs[s].parity  != sLevel.parity  ||
                    rLevel.parity != sLevel.parity    continue
                end
                #
                if      kind == "T^(1) amplitude"
                #--------------------------------
                    wa = compute("angular coefficients: 1-p, Grasp92", 0, 1, rLevel.basis.csfs[r], sLevel.basis.csfs[s])
                    for  coeff in wa
                        ja = Basics.subshell_2j(rLevel.basis.orbitals[coeff.a].subshell)
                        jb = Basics.subshell_2j(sLevel.basis.orbitals[coeff.b].subshell)
                        tamp  = InteractionStrength.hfs_t1(rLevel.basis.orbitals[coeff.a], sLevel.basis.orbitals[coeff.b], grid)
                        #
                        ## println("**  <$(coeff.a) || t1 || $(coeff.b)>  = $(coeff.T * tamp)   = $(coeff.T) * $tamp" )
                        me = me + coeff.T * tamp
                    end
                #
                elseif  kind == "T^(2) amplitude"
                #--------------------------------
                    wa = compute("angular coefficients: 1-p, Grasp92", 0, 2, rLevel.basis.csfs[r], sLevel.basis.csfs[s])
                    for  coeff in wa
                        ja = Basics.subshell_2j(rLevel.basis.orbitals[coeff.a].subshell)
                        jb = Basics.subshell_2j(sLevel.basis.orbitals[coeff.b].subshell)
                        tamp  = InteractionStrength.hfs_t2(rLevel.basis.orbitals[coeff.a], sLevel.basis.orbitals[coeff.b], grid)
                        #
                        ## println("**  <$(coeff.a) || t2 || $(coeff.b)>  = $(coeff.T * tamp)   = $(coeff.T) * $tamp" )
                        me = me + coeff.T * tamp
                    end
                #
                else    error("stop a")
                end
                #
                matrix[r,s] = me
                println("**                                                                    " *
                        "                      <$r || $kind || $s> = $me" )
            end
        end
        if  printout   printstyled("done.\n", color=:light_green)   end
        amplitude = transpose(rLevel.mc) * matrix * sLevel.mc
        #
        return( amplitude )
    end

    """
    `Hfs.computeAmplitudesProperties(outcome::Hfs.Outcome, grid::Radial.Grid, settings::Hfs.Settings, im::Hfs.InteractionMatrix)
        ... to compute all amplitudes and properties of for a given level; an outcome::Hfs.Outcome is returned for which the
            amplitudes and properties are now evaluated explicitly.
    """
    function  computeAmplitudesProperties(outcome::Hfs.Outcome, grid::Radial.Grid, settings::Hfs.Settings, im::Hfs.InteractionMatrix)
        AIoverMu = 0.0;   BoverQ = 0.0;   amplitudeT1 = 0.0;   amplitudeT2 = 0.0;    J = AngularMomentum.oneJ(outcome.Jlevel.J)
        if  settings.calcT1  &&  outcome.Jlevel.J != AngularJ64(0)
            if  im.calcT1   amplitudeT1 = transpose(outcome.Jlevel.mc) * im.matrixT1 * outcome.Jlevel.mc
            else            amplitudeT1 = Hfs.amplitude("T^(1) amplitude", outcome.Jlevel, outcome.Jlevel, grid)
            end

            AIoverMu = amplitudeT1 / sqrt(J * (J+1))
        end
        #
        if  settings.calcT2  &&  outcome.Jlevel.J != AngularJ64(0)
            if  im.calcT2   amplitudeT2 = transpose(outcome.Jlevel.mc) * im.matrixT2 * outcome.Jlevel.mc
            else            amplitudeT2 = Hfs.amplitude("T^(2) amplitude", outcome.Jlevel, outcome.Jlevel, grid)
            end

            BoverQ   = 2 * amplitudeT2 * sqrt( (2J-1) / ((J+1)*(2J+3)) )  ## * sqrt(J)
        end
        newOutcome = Hfs.Outcome( outcome.Jlevel, AIoverMu, BoverQ, amplitudeT1, amplitudeT2,
                                  outcome.nuclearI, outcome.hasFlevels, outcome.Flevels)
        return( newOutcome )
    end

    """
    `Hfs.computeHyperfineMultiplet(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid, settings::Hfs.Settings; output=true)`
        ... to compute a hyperfine multiplet, i.e. a representation of hyperfine levels within a hyperfine-coupled basis as defined by the
            given (electronic) multiplet; a hfsMultiplet::IJF_Multiplet is returned.
    """
    function computeHyperfineMultiplet(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid, settings::Hfs.Settings; output=true)
        println("")
        printstyled("Hfs.computeHyperfineMultiplet(): The computation of the hyperfine multiplet starts now ... \n", color=:light_green)
        printstyled("------------------------------------------------------------------------------------------ \n", color=:light_green)
        #
        hfsBasis     = Hfs.defineHyperfineBasis(multiplet, nm, settings)
        hfsMultiplet = Hfs.computeHyperfineRepresentation(hfsBasis, nm, grid, settings)
        # Print all results to screen
        hfsMultiplet = Basics.sortByEnergy(hfsMultiplet)
        Basics.tabulate("multiplet: energies",                                      hfsMultiplet)
        Basics.tabulate("multiplet: energy of each level relative to lowest level", hfsMultiplet)
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary
            Basics.tabulate("multiplet: energies",                                      hfsMultiplet, stream=iostream)
            Basics.tabulate("multiplet: energy of each level relative to lowest level", hfsMultiplet, stream=iostream)
        end
        #
        if    output    return( hfsMultiplet )
        else            return( nothing )
        end
    end

    """
    `Hfs.computeHyperfineRepresentation(hfsBasis::IJF_Basis, nm::Nuclear.Model, grid::Radial.Grid, settings::Hfs.Settings)`
        ... to set-up and diagonalized the Hamiltonian matrix of H^(DFB) + H^(hfs) within the atomic hyperfine (IJF-coupled) basis;
            a hfsMultiplet::IJF_Multiplet is returned.
    """
    function computeHyperfineRepresentation(hfsBasis::IJF_Basis, nm::Nuclear.Model, grid::Radial.Grid, settings::Hfs.Settings)
        n = length(hfsBasis.vectors);   matrix = zeros(n,n)
        for  r = 1:n
            for  s = 1:n
                matrix[r,s] = 0.
                if  r == s    matrix[r,s] = matrix[r,s] + hfsBasis.vectors[r].levelJ.energy                  end
                if  hfsBasis.vectors[r].F              !=  hfsBasis.vectors[s].F                continue     end
                if  hfsBasis.vectors[r].levelJ.parity  !=  hfsBasis.vectors[s].levelJ.parity    continue     end
                # Now add the hyperfine matrix element for K = 1 (magnetic) and K = 2 (electric) contributions
                spinI = nm.spinI;   Jr = hfsBasis.vectors[r].levelJ.J;   Js = hfsBasis.vectors[s].levelJ.J;   F = hfsBasis.vectors[r].F
                Ix    = AngularMomentum.oneJ(spinI)
                #
                wa = AngularMomentum.phaseFactor([spinI, +1, Jr, +1, F])  *
                     AngularMomentum.Wigner_6j(spinI, Jr, F, Js, spinI, AngularJ64(1))
                wb = Hfs.amplitude("T^(1) amplitude", hfsBasis.vectors[r].levelJ, hfsBasis.vectors[s].levelJ, grid::Radial.Grid)
                wc = nm.mu * sqrt( (Ix+1)/Ix )
                matrix[r,s] = matrix[r,s] + wa * wb * wc     * 1.0e-6 ## fudge-factor to keep HFS interaction small
                #
                if  spinI  in [AngularJ64(0), AngularJ64(1//2)]                                 continue     end
                wa = AngularMomentum.phaseFactor([spinI, +1, Jr, +1, F]) *
                     AngularMomentum.Wigner_6j(spinI, Jr, F, Js, spinI, AngularJ64(2))
                wb = Hfs.amplitude("T^(2) amplitude", hfsBasis.vectors[r].levelJ, hfsBasis.vectors[s].levelJ, grid::Radial.Grid)
                wc = nm.Q / 2. * sqrt( (Ix+1)*(2Ix+3)/ (Ix*(2Ix-1)) )
                matrix[r,s] = matrix[r,s] + wa * wb * wc     * 1.0e-6 ## fudge-factor to keep HFS interaction small
            end
        end
        #
        # Diagonalize the matrix and set-up the representation
        eigen    = Basics.diagonalize("matrix: LinearAlgebra", matrix)
        levelFs  = Hfs.IJF_Level[]
        for  ev = 1:length(eigen.values)
            # Construct the eigenvector with regard to the given basis (not w.r.t the symmetry block)
            evector   = eigen.vectors[ev];    en = eigen.values[ev]
            parity    = Basics.plus;    F = AngularJ64(0);     MF = AngularM64(0)
            for  r = 1:length(hfsBasis.vectors)
                if  abs(evector[r]) > 1.0e-6
                    parity = hfsBasis.vectors[r].levelJ.parity
                    F      = hfsBasis.vectors[r].F;     MF = AngularM64(hfsBasis.vectors[r].F);   break    end
            end
            newlevelF = Hfs.IJF_Level(nm.spinI, F, MF, parity, en, true, hfsBasis, evector)
            push!( levelFs, newlevelF)
        end
        hfsMultiplet = Hfs.IJF_Multiplet("hyperfine", levelFs)

        return( hfsMultiplet )
    end

    """
    `Hfs.computeInteractionMatrix(basis::Basis, grid::Radial.Grid, settings::Hfs.Settings)`
        ... to compute the T^1 and/or T^2 interaction matrices for the given basis, i.e. (<csf_r || T^(n)) || csf_s>).
            An im::Hfs.InteractionMatrix is returned.
    """
    function  computeInteractionMatrix(basis::Basis, grid::Radial.Grid, settings::Hfs.Settings)
        #
        ncsf = length(basis.csfs);    matrixT1 = zeros(ncsf,ncsf);    matrixT2 = zeros(ncsf,ncsf)
        #
        if  settings.calcT1
            calcT1 = true;    matrixT1 = zeros(ncsf,ncsf)
            for  r = 1:ncsf
                for  s = 1:ncsf
                    if  basis.csfs[r].parity  != basis.csfs[s].parity   continue    end
                    #
                    wa = compute("angular coefficients: 1-p, Grasp92", 0, 1, basis.csfs[r], basis.csfs[s])
                    for  coeff in wa
                        ja   = Basics.subshell_2j(basis.orbitals[coeff.a].subshell)
                        jb   = Basics.subshell_2j(basis.orbitals[coeff.b].subshell)
                        tamp = InteractionStrength.hfs_t1(basis.orbitals[coeff.a], basis.orbitals[coeff.b], grid)
                        matrixT1[r,s] = matrixT1[r,s] + coeff.T * tamp
                    end
                end
            end
        else
            calcT1 = false;    matrixT1 = zeros(2,2)
        end
        #
        if  settings.calcT2
            calcT2 = true;    matrixT2 = zeros(ncsf,ncsf)
            for  r = 1:ncsf
                for  s = 1:ncsf
                    if  basis.csfs[r].parity  != basis.csfs[s].parity   continue    end
                    #
                    wa = compute("angular coefficients: 1-p, Grasp92", 0, 2, basis.csfs[r], basis.csfs[s])
                    for  coeff in wa
                        ja   = Basics.subshell_2j(basis.orbitals[coeff.a].subshell)
                        jb   = Basics.subshell_2j(basis.orbitals[coeff.b].subshell)
                        tamp  = InteractionStrength.hfs_t2(basis.orbitals[coeff.a], basis.orbitals[coeff.b], grid)
                        matrixT2[r,s] = matrixT2[r,s] + coeff.T * tamp
                    end
                end
            end
        else
            calcT2 = false;    matrixT2 = zeros(2,2)
        end
        #
        im = Hfs.InteractionMatrix(calcT1, calcT2, matrixT1, matrixT2)
        #
        return( im )
    end

    """
    `Hfs.computeOutcomes(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid, settings::Hfs.Settings; output=true)`
        ... to compute (as selected) the HFS A and B parameters as well as hyperfine energy splittings for the levels
            of the given multiplet and as specified by the given settings. The results are printed in neat tables to
            screen and, if requested, an arrays{Hfs.Outcome,1} with all the results are returned.
    """
    function computeOutcomes(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid, settings::Hfs.Settings; output=true)
        println("")
        printstyled("Hfs.computeOutcomes(): The computation of the Hyperfine amplitudes and parameters starts now ... \n", color=:light_green)
        printstyled("------------------------------------------------------------------------------------------------ \n", color=:light_green)
        println("")
        outcomes = Hfs.determineOutcomes(multiplet, settings)
        # Display all selected levels before the computations start
        if  settings.printBefore    Hfs.displayOutcomes(outcomes)    end
        # Calculate all amplitudes and requested properties
        im = computeInteractionMatrix(multiplet.levels[1].basis, grid, settings)
        newOutcomes = Hfs.Outcome[]
        for  outcome in outcomes
            newOutcome = Hfs.computeAmplitudesProperties(outcome, grid, settings, im)
            push!( newOutcomes, newOutcome)
        end
        # Print all results to screen
        Hfs.displayResults(stdout, newOutcomes, nm, settings)
        # Compute and display the non-diagonal hyperfine amplitudes, if requested
        if  settings.calcNondiagonal    Hfs.displayNondiagonal(stdout, multiplet, grid, settings)   end
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary
            Hfs.displayResults(iostream, newOutcomes, nm, settings)
            if  settings.calcNondiagonal    Hfs.displayNondiagonal(iostream, multiplet, grid, settings)   end
        end
        #
        if    output    return( newOutcomes )
        else            return( nothing )
        end
    end

    """
    `Hfs.defineHyperfineBasis(multiplet::Multiplet, nm::Nuclear.Model, settings::Hfs.Settings)`
        ... to define/set-up an atomic hyperfine (IJF-coupled) basis for the given electronic multipet;
            a hfsBasis::IJF_Basis is returned.
    """
    function  defineHyperfineBasis(multiplet::Multiplet, nm::Nuclear.Model, settings::Hfs.Settings)
        function  display_ijfVector(i::Int64, vector::Hfs.IJF_Vector)
            si = string(i);   ni = length(si);    sa = repeat(" ", 5);    sa = sa[1:5-ni] * si * ")  "
            sa = sa * "[" * string( LevelSymmetry(vector.levelJ.J, vector.levelJ.parity) ) * "] " * string(vector.F) * repeat(" ", 4)
            return( sa )
        end

        vectors = IJF_Vector[]
        #
        for i = 1:length(multiplet.levels)
            Flist = Basics.oplus(nm.spinI, multiplet.levels[i].J)
            for  F in Flist
                push!(vectors,  IJF_Vector(nm.spinI, F, multiplet.levels[i]) )
            end
        end
        #
        hfsBasis = IJF_Basis(vectors, multiplet.levels[1].basis)
        #
        println(" ")
        println("  Construction of a atomic hyperfine (IJF-coupled) basis of dimension $(length(hfsBasis.vectors)), based on " *
                "$(length(multiplet.levels)) electronic levels and nuclear spin $(hfsBasis.vectors[1].I).")
        println("  Basis vectors are defined as [J^P] F: \n")
        sa = "  "
        for  k = 1:length(hfsBasis.vectors)
            if  length(sa) > 100    println(sa);   sa = "  "   end
            sa = sa * display_ijfVector(k, hfsBasis.vectors[k])
        end
        if   length(sa) > 5    println(sa)   end
        println(" ")

        return( hfsBasis )
    end

    """
    `Hfs.determineOutcomes(multiplet::Multiplet, settings::Hfs.Settings)`
        ... to determine a list of Outcomes's for the computation of HFS A- and B-parameters for the given multiplet.
            It takes into account the particular selections and settings. An Array{Hfs.Outcome,1} is returned. Apart from the
            level specification, all physical properties are set to zero during the initialization process.
    """
    function  determineOutcomes(multiplet::Multiplet, settings::Hfs.Settings)
        outcomes = Hfs.Outcome[]
        for  level  in  multiplet.levels
            if  Basics.selectLevel(level, settings.levelSelection)
                push!( outcomes, Hfs.Outcome(level, 0., 0., 0., 0., AngularJ64(0), false, false) )
            end
        end
        return( outcomes )
    end

    """
    `Hfs.displayNondiagonal(stream::IO, multiplet::Multiplet, grid::Radial.Grid, settings::Hfs.Settings)`
        ... to compute and display all non-diagonal hyperfine amplitudes for the selected levels. A small neat table of
            all (pairwise) hyperfine amplitudes is printed but nothing is returned otherwise.
    """
    function  displayNondiagonal(stream::IO, multiplet::Multiplet, grid::Radial.Grid, settings::Hfs.Settings)
        # Determine pairs to be calculated
        pairs = Tuple{Int64,Int64}[]
        for  (f, fLevel)  in  enumerate(multiplet.levels)
            for  (i, iLevel)  in  enumerate(multiplet.levels)
                if  Basics.selectLevel(fLevel, settings.levelSelection)  &&   Basics.selectLevel(iLevel, settings.levelSelection)
                    push!( pairs, (f,i) )
                end
            end
        end
        #
        nx = 107
        println(stream, " ")
        println(stream, "  Selected (non-) diagonal hyperfine amplitudes:")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  "
        sa = sa * TableStrings.center(10, "Level_f"; na=2)
        sa = sa * TableStrings.center(10, "Level_i"; na=2)
        sa = sa * TableStrings.center(10, "J^P_f";   na=3)
        sa = sa * TableStrings.center(57, "T1   --   Amplitudes   --   T2"; na=4);
        sa = sa * TableStrings.center(10, "J^P_f";   na=4)
        println(stream, sa);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  (f,i) in pairs
            sa   = "  ";
            sa   = sa * TableStrings.center(10, string(f); na=2)
            sa   = sa * TableStrings.center(10, string(i); na=2)
            symf = LevelSymmetry( multiplet.levels[f].J, multiplet.levels[f].parity)
            symi = LevelSymmetry( multiplet.levels[i].J, multiplet.levels[i].parity)
            sa   = sa * TableStrings.center(10, string(symf); na=4)
            T1   = Hfs.amplitude("T^(1) amplitude", multiplet.levels[f], multiplet.levels[i], grid, printout=false)
            T2   = Hfs.amplitude("T^(2) amplitude", multiplet.levels[f], multiplet.levels[i], grid, printout=false)
            sa   = sa * @sprintf("%.5e %s %.5e", T1.re, "  ", T1.im) * "    "
            sa   = sa * @sprintf("%.5e %s %.5e", T2.re, "  ", T2.im) * "    "
            sa   = sa * TableStrings.center(10, string(symi); na=4)
            println(stream, sa )
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        return( nothing )
    end

    """
    `Hfs.displayOutcomes(outcomes::Array{Hfs.Outcome,1})`
        ... to display a list of levels that have been selected for the computations A small neat table of all
            selected levels and their energies is printed but nothing is returned otherwise.
    """
    function  displayOutcomes(outcomes::Array{Hfs.Outcome,1})
        nx = 43
        println(" ")
        println("  Selected HFS levels:")
        println(" ")
        println("  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(10, "Level"; na=2);                             sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center(10, "J^P";   na=4);                             sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(14, "Energy"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        println(sa);    println(sb);    println("  ", TableStrings.hLine(nx))
        #
        for  outcome in outcomes
            sa  = "  ";    sym = LevelSymmetry( outcome.Jlevel.J, outcome.Jlevel.parity)
            sa = sa * TableStrings.center(10, TableStrings.level(outcome.Jlevel.index); na=2)
            sa = sa * TableStrings.center(10, string(sym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", outcome.Jlevel.energy)) * "    "
            println( sa )
        end
        println("  ", TableStrings.hLine(nx))
        println(" ")
        #
        return( nothing )
    end

    """
    `Hfs.displayResults(stream::IO, outcomes::Array{Hfs.Outcome,1}, nm::Nuclear.Model, settings::Hfs.Settings)`
        ... to display the energies, A- and B-values, Delta E_F energy shifts, etc. for the selected levels. All nuclear
            parameters are taken from the nuclear model. A neat table is printed but nothing is returned otherwise.
    """
    function  displayResults(stream::IO, outcomes::Array{Hfs.Outcome,1}, nm::Nuclear.Model, settings::Hfs.Settings)
        nx = 117
        println(stream, " ")
        println(stream, "  HFS parameters and amplitudes:")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(10, "Level"; na=2);                             sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center(10, "J^P";   na=4);                             sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(14, "Energy"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        sa = sa * TableStrings.center(31, "A/mu [mu_N]  --  HFS  --  B/Q [barn]"; na=4);
        sb = sb * TableStrings.center(31, TableStrings.inUnits("energy"); na=4)
        sa = sa * TableStrings.center(32, "T1 -- Amplitudes -- T2"    ; na=4);        sb = sb * TableStrings.hBlank(36)
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  outcome in outcomes
            sa  = "  ";    sym = LevelSymmetry( outcome.Jlevel.J, outcome.Jlevel.parity)
            sa = sa * TableStrings.center(10, TableStrings.level(outcome.Jlevel.index); na=2)
            sa = sa * TableStrings.center(10, string(sym); na=4)
            energy = outcome.Jlevel.energy
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", energy))           * "      "
            wa = Defaults.convertUnits("energy: from atomic", outcome.AIoverMu) / nm.spinI.num * nm.spinI.den
            wa = wa / (2 * 1836.15267 ) * Defaults.getDefaults("alpha")  # take 2m_p and alpha into account
            sa = sa * @sprintf("%.8e", wa)                                                   * "    "
            wa = Defaults.convertUnits("energy: from atomic", outcome.BoverQ)
            wa = wa / Defaults.convertUnits("cross section: from atomic to barn", 1.0)  # take Q [barn] into account
            sa = sa * @sprintf("%.8e", wa)                                                   * "    "
            sa = sa * @sprintf("%.8e %s %.8e", outcome.amplitudeT1.re, "  ", outcome.amplitudeT2.re) * "    "
            println(stream, sa )
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        #
        if !settings.printDeltaEF   return( nothing )   end
        #
        # Printout the Delta E_F energy shifts of the hyperfine levels |alpha F> with regard to the (electronic) levels |alpha J>
        nx = 90
        println(stream, " ")
        println(stream, "  HFS Delta E_F energy shifts with regard to the (electronic) level energies E_J:")
        println(stream, " ")
        println(stream, "    Nuclear spin I:                        $(nm.spinI) ")
        println(stream, "    Nuclear magnetic-dipole moment    mu = $(nm.mu)    ")
        println(stream, "    Nuclear electric-quadrupole moment Q = $(nm.Q)     ")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(10, "Level"; na=2);                             sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center(10, "J^P";   na=4);                             sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(14, "Energy"; na=4);                            sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(10, "F^P";   na=4);                             sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(14, "Delta E_F"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        sa = sa * TableStrings.center(14, "C factor"; na=4);
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  outcome in outcomes
            sa  = "  ";    sym = LevelSymmetry( outcome.Jlevel.J, outcome.Jlevel.parity)
            sa = sa * TableStrings.center(10, TableStrings.level(outcome.Jlevel.index); na=2)
            sa = sa * TableStrings.center(10, string(sym); na=4)
            energy = outcome.Jlevel.energy
            #
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", energy))           * "    "
            J     = AngularMomentum.oneJ(outcome.Jlevel.J);   spinI = AngularMomentum.oneJ(nm.spinI)
            Flist = Basics.oplus(nm.spinI, outcome.Jlevel.J)
            first = true
            for  Fang in Flist
                Fsym    = LevelSymmetry(Fang, outcome.Jlevel.parity)
                F       = AngularMomentum.oneJ(Fang)
                Cfactor = F*(F+1) - J*(J+1) - spinI*(spinI+1)
                energy  = outcome.AIoverMu * nm.mu / spinI * Cfactor / 2.
                if  abs(outcome.BoverQ) > 1.0e-10
                    energy  = energy +  outcome.BoverQ * nm.Q * 3/4 * (Cfactor*(Cfactor+1) - spinI*(spinI+1)*J*(J+1) ) /
                                        ( 2spinI*(2spinI-1)*J*(2J-1) )
                end
                sb = TableStrings.center(10, string(Fsym); na=2)
                sb = sb * TableStrings.flushright(16, @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", energy))) * "    "
                sb = sb * TableStrings.flushright(12, @sprintf("%.5e", Cfactor))
                #
                if   first    println(stream,  sa*sb );   first = false
                else          println(stream,  TableStrings.hBlank( length(sa) ) * sb )
                end
            end
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        return( nothing )
    end

end # module


module Atmos

export AtmosModel, AtmosAcousticLinearModel, AtmosAcousticGravityLinearModel

using CLIMAParameters
using CLIMAParameters.Planet: grav, cp_d
using CLIMAParameters.Atmos.SubgridScale: C_smag
using DocStringExtensions
using LinearAlgebra, StaticArrays
using ..ConfigTypes
using ..Orientations
import ..Orientations:
    vertical_unit_vector,
    altitude,
    latitude,
    longitude,
    projection_normal,
    gravitational_potential,
    ∇gravitational_potential,
    projection_tangential

using ..VariableTemplates
using ..Thermodynamics
using ..TemperatureProfiles

using ..TurbulenceClosures
import ..TurbulenceClosures: turbulence_tensors
using ..TurbulenceConvection

import ..Thermodynamics: internal_energy
using ..MPIStateArrays: MPIStateArray
using ..Mesh.Grids:
    VerticalDirection,
    HorizontalDirection,
    min_node_distance,
    EveryDirection,
    Direction

using ClimateMachine.BalanceLaws
using ClimateMachine.Problems

import ClimateMachine.BalanceLaws:
    vars_state,
    flux_first_order!,
    flux_second_order!,
    source!,
    wavespeed,
    boundary_state!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    transform_post_gradient_laplacian!,
    init_state_auxiliary!,
    init_state_prognostic!,
    update_auxiliary_state!,
    indefinite_stack_integral!,
    reverse_indefinite_stack_integral!,
    integral_load_auxiliary_state!,
    integral_set_auxiliary_state!,
    reverse_integral_load_auxiliary_state!,
    reverse_integral_set_auxiliary_state!

import ClimateMachine.DGMethods:
    LocalGeometry, lengthscale, resolutionmetric, DGModel

import ..DGMethods.NumericalFluxes:
    boundary_state!,
    boundary_flux_second_order!,
    normal_boundary_flux_second_order!,
    NumericalFluxFirstOrder,
    NumericalFluxGradient,
    NumericalFluxSecondOrder,
    CentralNumericalFluxHigherOrder,
    CentralNumericalFluxDivergence,
    CentralNumericalFluxFirstOrder,
    numerical_flux_first_order!,
    NumericalFluxFirstOrder
using ..DGMethods.NumericalFluxes:
    RoeNumericalFlux, HLLCNumericalFlux, RusanovNumericalFlux

import ..Courant: advective_courant, nondiffusive_courant, diffusive_courant

"""
    AtmosModel <: BalanceLaw

A `BalanceLaw` for atmosphere modeling. Users may over-ride prescribed
default values for each field.

# Usage

    AtmosModel(
        param_set,
        problem,
        orientation,
        ref_state,
        turbulence,
        hyperdiffusion,
        spongelayer,
        moisture,
        radiation,
        source,
        tracers,
        data_config,
    )

# Fields
$(DocStringExtensions.FIELDS)
"""
struct AtmosModel{FT, PS, PR, O, RS, T, TC, HD, VS, M, P, R, S, TR, DC} <:
       BalanceLaw
    "Parameter Set (type to dispatch on, e.g., planet parameters. See CLIMAParameters.jl package)"
    param_set::PS
    "Problem (initial and boundary conditions)"
    problem::PR
    "An orientation model"
    orientation::O
    "Reference State (For initial conditions, or for linearisation when using implicit solvers)"
    ref_state::RS
    "Turbulence Closure (Equations for dynamics of under-resolved turbulent flows)"
    turbulence::T
    "Turbulence Convection Closure (e.g., EDMF)"
    turbconv::TC
    "Hyperdiffusion Model (Equations for dynamics of high-order spatial wave attenuation)"
    hyperdiffusion::HD
    "Viscous sponge layers"
    viscoussponge::VS
    "Moisture Model (Equations for dynamics of moist variables)"
    moisture::M
    "Precipitation Model (Equations for dynamics of precipitating species)"
    precipitation::P
    "Radiation Model (Equations for radiative fluxes)"
    radiation::R
    "Source Terms (Problem specific source terms)"
    source::S
    "Tracer Terms (Equations for dynamics of active and passive tracers)"
    tracers::TR
    "Data Configuration (Helper field for experiment configuration)"
    data_config::DC
end

"""
    AtmosModel{FT}()

Constructor for `AtmosModel` (where `AtmosModel <: BalanceLaw`) for LES
and single stack configurations.
"""
function AtmosModel{FT}(
    ::Union{Type{AtmosLESConfigType}, Type{SingleStackConfigType}},
    param_set::AbstractParameterSet;
    init_state_prognostic::ISP = nothing,
    problem::PR = AtmosProblem(init_state_prognostic = init_state_prognostic),
    orientation::O = FlatOrientation(),
    ref_state::RS = HydrostaticState(DecayingTemperatureProfile{FT}(param_set),),
    turbulence::T = SmagorinskyLilly{FT}(0.21),
    turbconv::TC = NoTurbConv(),
    hyperdiffusion::HD = NoHyperDiffusion(),
    viscoussponge::VS = NoViscousSponge(),
    moisture::M = EquilMoist{FT}(),
    precipitation::P = NoPrecipitation(),
    radiation::R = NoRadiation(),
    source::S = (
        Gravity(),
        Coriolis(),
        GeostrophicForcing{FT}(7.62e-5, 0, 0),
        turbconv_sources(turbconv)...,
    ),
    tracers::TR = NoTracers(),
    data_config::DC = nothing,
) where {FT <: AbstractFloat, ISP, PR, O, RS, T, TC, HD, VS, M, P, R, S, TR, DC}

    atmos = (
        param_set,
        problem,
        orientation,
        ref_state,
        turbulence,
        turbconv,
        hyperdiffusion,
        viscoussponge,
        moisture,
        precipitation,
        radiation,
        source,
        tracers,
        data_config,
    )

    return AtmosModel{FT, typeof.(atmos)...}(atmos...)
end

"""
    AtmosModel{FT}()

Constructor for `AtmosModel` (where `AtmosModel <: BalanceLaw`) for GCM
configurations.
"""
function AtmosModel{FT}(
    ::Type{AtmosGCMConfigType},
    param_set::AbstractParameterSet;
    init_state_prognostic::ISP = nothing,
    problem::PR = AtmosProblem(init_state_prognostic = init_state_prognostic),
    orientation::O = SphericalOrientation(),
    ref_state::RS = HydrostaticState(DecayingTemperatureProfile{FT}(param_set),),
    turbulence::T = SmagorinskyLilly{FT}(C_smag(param_set)),
    turbconv::TC = NoTurbConv(),
    hyperdiffusion::HD = NoHyperDiffusion(),
    viscoussponge::VS = NoViscousSponge(),
    moisture::M = EquilMoist{FT}(),
    precipitation::P = NoPrecipitation(),
    radiation::R = NoRadiation(),
    source::S = (Gravity(), Coriolis(), turbconv_sources(turbconv)...),
    tracers::TR = NoTracers(),
    data_config::DC = nothing,
) where {FT <: AbstractFloat, ISP, PR, O, RS, T, TC, HD, VS, M, P, R, S, TR, DC}

    atmos = (
        param_set,
        problem,
        orientation,
        ref_state,
        turbulence,
        turbconv,
        hyperdiffusion,
        viscoussponge,
        moisture,
        precipitation,
        radiation,
        source,
        tracers,
        data_config,
    )

    return AtmosModel{FT, typeof.(atmos)...}(atmos...)
end

"""
    vars_state(m::AtmosModel, ::Prognostic, FT)

Conserved state variables (prognostic variables).
"""
function vars_state(m::AtmosModel, st::Prognostic, FT)
    @vars begin
        ρ::FT
        ρu::SVector{3, FT}
        ρe::FT
        turbulence::vars_state(m.turbulence, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
        moisture::vars_state(m.moisture, st, FT)
        radiation::vars_state(m.radiation, st, FT)
        tracers::vars_state(m.tracers, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::Gradient, FT)

Pre-transform gradient variables.
"""
function vars_state(m::AtmosModel, st::Gradient, FT)
    @vars begin
        u::SVector{3, FT}
        h_tot::FT
        turbulence::vars_state(m.turbulence, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
        moisture::vars_state(m.moisture, st, FT)
        tracers::vars_state(m.tracers, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::GradientFlux, FT)

Post-transform gradient variables.
"""
function vars_state(m::AtmosModel, st::GradientFlux, FT)
    @vars begin
        ∇h_tot::SVector{3, FT}
        turbulence::vars_state(m.turbulence, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
        moisture::vars_state(m.moisture, st, FT)
        tracers::vars_state(m.tracers, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::GradientLaplacian, FT)

Pre-transform hyperdiffusive variables.
"""
function vars_state(m::AtmosModel, st::GradientLaplacian, FT)
    @vars begin
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::Hyperdiffusive, FT)

Post-transform hyperdiffusive variables.
"""
function vars_state(m::AtmosModel, st::Hyperdiffusive, FT)
    @vars begin
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::Auxiliary, FT)

Auxiliary variables, such as vertical (stack) integrals, coordinates,
orientation information, reference states, subcomponent auxiliary vars,
debug variables.
"""
function vars_state(m::AtmosModel, st::Auxiliary, FT)
    @vars begin
        ∫dz::vars_state(m, UpwardIntegrals(), FT)
        ∫dnz::vars_state(m, DownwardIntegrals(), FT)
        coord::SVector{3, FT}
        orientation::vars_state(m.orientation, st, FT)
        ref_state::vars_state(m.ref_state, st, FT)
        turbulence::vars_state(m.turbulence, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
        moisture::vars_state(m.moisture, st, FT)
        tracers::vars_state(m.tracers, st, FT)
        radiation::vars_state(m.radiation, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::UpwardIntegrals, FT)
"""
function vars_state(m::AtmosModel, st::UpwardIntegrals, FT)
    @vars begin
        radiation::vars_state(m.radiation, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::DownwardIntegrals, FT)
"""
function vars_state(m::AtmosModel, st::DownwardIntegrals, FT)
    @vars begin
        radiation::vars_state(m.radiation, st, FT)
    end
end

####
#### Forward orientation methods
####
projection_normal(bl, aux, u⃗) =
    projection_normal(bl.orientation, bl.param_set, aux, u⃗)
projection_tangential(bl, aux, u⃗) =
    projection_tangential(bl.orientation, bl.param_set, aux, u⃗)
latitude(bl, aux) = latitude(bl.orientation, aux)
longitude(bl, aux) = longitude(bl.orientation, aux)
altitude(bl, aux) = altitude(bl.orientation, bl.param_set, aux)
vertical_unit_vector(bl, aux) =
    vertical_unit_vector(bl.orientation, bl.param_set, aux)
gravitational_potential(bl, aux) = gravitational_potential(bl.orientation, aux)
∇gravitational_potential(bl, aux) =
    ∇gravitational_potential(bl.orientation, aux)

turbulence_tensors(atmos::AtmosModel, args...) =
    turbulence_tensors(atmos.turbulence, atmos, args...)

include("problem.jl")
include("ref_state.jl")
include("moisture.jl")
include("thermo_states.jl")
include("precipitation.jl")
include("radiation.jl")
include("source.jl")
include("tracers.jl")
include("linear.jl")
include("courant.jl")
include("filters.jl")

"""
    flux_first_order!(
        m::AtmosModel,
        flux::Grad,
        state::Vars,
        aux::Vars,
        t::Real
    )

Computes and assembles non-diffusive fluxes in the model
equations.
"""
@inline function flux_first_order!(
    m::AtmosModel,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    ρ = state.ρ
    ρinv = 1 / ρ
    ρu = state.ρu
    u = ρinv * ρu

    # advective terms
    flux.ρ = ρ * u
    flux.ρu = ρ * u .* u'
    flux.ρe = u * state.ρe

    # pressure terms
    ts = recover_thermo_state(m, state, aux)
    p = air_pressure(ts)
    if m.ref_state isa HydrostaticState
        flux.ρu += (p - aux.ref_state.p) * I
    else
        flux.ρu += p * I
    end
    flux.ρe += u * p
    flux_radiation!(m.radiation, m, flux, state, aux, t)
    flux_moisture!(m.moisture, m, flux, state, aux, t)
    flux_tracers!(m.tracers, m, flux, state, aux, t)
    flux_first_order!(m.turbconv, m, flux, state, aux, t)
end

function compute_gradient_argument!(
    atmos::AtmosModel,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ρinv = 1 / state.ρ
    transform.u = ρinv * state.ρu
    ts = recover_thermo_state(atmos, state, aux)
    e_tot = state.ρe * (1 / state.ρ)
    transform.h_tot = total_specific_enthalpy(ts, e_tot)

    compute_gradient_argument!(atmos.moisture, transform, state, aux, t)
    compute_gradient_argument!(atmos.turbulence, transform, state, aux, t)
    compute_gradient_argument!(
        atmos.hyperdiffusion,
        atmos,
        transform,
        state,
        aux,
        t,
    )
    compute_gradient_argument!(atmos.tracers, transform, state, aux, t)
    compute_gradient_argument!(atmos.turbconv, atmos, transform, state, aux, t)
end

function compute_gradient_flux!(
    atmos::AtmosModel,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    diffusive.∇h_tot = ∇transform.h_tot

    # diffusion terms required for SGS turbulence computations
    compute_gradient_flux!(
        atmos.turbulence,
        atmos.orientation,
        diffusive,
        ∇transform,
        state,
        aux,
        t,
    )
    # diffusivity of moisture components
    compute_gradient_flux!(atmos.moisture, diffusive, ∇transform, state, aux, t)
    compute_gradient_flux!(atmos.tracers, diffusive, ∇transform, state, aux, t)
    compute_gradient_flux!(
        atmos.turbconv,
        atmos,
        diffusive,
        ∇transform,
        state,
        aux,
        t,
    )
end

function transform_post_gradient_laplacian!(
    atmos::AtmosModel,
    hyperdiffusive::Vars,
    hypertransform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform_post_gradient_laplacian!(
        atmos.hyperdiffusion,
        atmos,
        hyperdiffusive,
        hypertransform,
        state,
        aux,
        t,
    )
end

"""
    flux_second_order!(
        atmos::AtmosModel,
        flux::Grad,
        state::Vars,
        diffusive::Vars,
        hyperdiffusive::Vars,
        aux::Vars,
        t::Real
    )
Diffusive fluxes in AtmosModel. Viscosity, diffusivity are calculated
in the turbulence subcomponent and accessed within the diffusive flux
function. Contributions from subcomponents are then assembled (pointwise).
"""
@inline function flux_second_order!(
    atmos::AtmosModel,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    ν, D_t, τ = turbulence_tensors(atmos, state, diffusive, aux, t)
    ν, D_t, τ = sponge_viscosity_modifier(atmos, atmos.viscoussponge, ν, D_t, τ, aux)
    d_h_tot = -D_t .* diffusive.∇h_tot
    flux_second_order!(atmos, flux, state, τ, d_h_tot)
    flux_second_order!(atmos.moisture, flux, state, diffusive, aux, t, D_t)
    flux_second_order!(
        atmos.hyperdiffusion,
        flux,
        state,
        diffusive,
        hyperdiffusive,
        aux,
        t,
    )
    flux_second_order!(atmos.tracers, flux, state, diffusive, aux, t, D_t)
    flux_second_order!(atmos.turbconv, atmos, flux, state, diffusive, aux, t)
end

#TODO: Consider whether to not pass ρ and ρu (not state), foc BCs reasons
@inline function flux_second_order!(
    atmos::AtmosModel,
    flux::Grad,
    state::Vars,
    τ,
    d_h_tot,
)
    flux.ρu += τ * state.ρ
    flux.ρe += τ * state.ρu
    flux.ρe += d_h_tot * state.ρ
end

@inline function wavespeed(
    m::AtmosModel,
    nM,
    state::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    ρinv = 1 / state.ρ
    u = ρinv * state.ρu
    uN = abs(dot(nM, u))
    ts = recover_thermo_state(m, state, aux)
    ss = soundspeed_air(ts)

    FT = typeof(state.ρ)
    ws = fill(uN + ss, MVector{number_states(m, Prognostic()), FT})
    vars_ws = Vars{vars_state(m, Prognostic(), FT)}(ws)

    wavespeed_tracers!(m.tracers, vars_ws, nM, state, aux, t)

    return ws
end

function update_auxiliary_state!(
    dg::DGModel,
    m::AtmosModel,
    Q::MPIStateArray,
    t::Real,
    elems::UnitRange,
)
    FT = eltype(Q)
    state_auxiliary = dg.state_auxiliary

    if number_states(m, UpwardIntegrals()) > 0
        indefinite_stack_integral!(dg, m, Q, state_auxiliary, t, elems)
        reverse_indefinite_stack_integral!(dg, m, Q, state_auxiliary, t, elems)
    end

    update_auxiliary_state!(nodal_update_auxiliary_state!, dg, m, Q, t, elems)

    # TODO: Remove this hook. This hook was added for implementing
    # the first draft of EDMF, and should be removed so that we can
    # rely on a single vertical element traversal. This hook allows
    # us to compute globally vertical quantities specific to EDMF
    # until we're able to remove them or somehow incorporate them
    # into a higher level hierarchy.
    update_auxiliary_state!(dg, m.turbconv, m, Q, t, elems)

    return true
end

function nodal_update_auxiliary_state!(
    m::AtmosModel,
    state::Vars,
    aux::Vars,
    t::Real,
)
    atmos_nodal_update_auxiliary_state!(m.moisture, m, state, aux, t)
    atmos_nodal_update_auxiliary_state!(m.radiation, m, state, aux, t)
    atmos_nodal_update_auxiliary_state!(m.tracers, m, state, aux, t)
    turbulence_nodal_update_auxiliary_state!(m.turbulence, m, state, aux, t)
    turbconv_nodal_update_auxiliary_state!(m.turbconv, m, state, aux, t)
end

function integral_load_auxiliary_state!(
    m::AtmosModel,
    integ::Vars,
    state::Vars,
    aux::Vars,
)
    integral_load_auxiliary_state!(m.radiation, integ, state, aux)
    integral_load_auxiliary_state!(m.turbconv, m, integ, state, aux)
end

function integral_set_auxiliary_state!(m::AtmosModel, aux::Vars, integ::Vars)
    integral_set_auxiliary_state!(m.radiation, aux, integ)
    integral_set_auxiliary_state!(m.turbconv, m, aux, integ)
end

function reverse_integral_load_auxiliary_state!(
    m::AtmosModel,
    integ::Vars,
    state::Vars,
    aux::Vars,
)
    reverse_integral_load_auxiliary_state!(m.radiation, integ, state, aux)
end

function reverse_integral_set_auxiliary_state!(
    m::AtmosModel,
    aux::Vars,
    integ::Vars,
)
    reverse_integral_set_auxiliary_state!(m.radiation, aux, integ)
end

function atmos_nodal_init_state_auxiliary!(
    m::AtmosModel,
    aux::Vars,
    tmp::Vars,
    geom::LocalGeometry,
)
    aux.coord = geom.coord
    init_aux_turbulence!(m.turbulence, m, aux, geom)
    atmos_init_aux!(m.ref_state, m, aux, tmp, geom)
    init_aux_hyperdiffusion!(m.hyperdiffusion, m, aux, geom)
    atmos_init_aux!(m.tracers, m, aux, geom)
    init_aux_turbconv!(m.turbconv, m, aux, geom)
    m.problem.init_state_auxiliary(m.problem, m, aux, geom)
end

"""
    init_state_auxiliary!(
        m::AtmosModel,
        aux::Vars,
        grid,
        direction
    )

Initialise auxiliary variables for each AtmosModel subcomponent.
Store Cartesian coordinate information in `aux.coord`.
"""
function init_state_auxiliary!(
    m::AtmosModel,
    state_auxiliary::MPIStateArray,
    grid,
    direction,
)
    init_aux!(m, m.orientation, state_auxiliary, grid, direction)

    init_state_auxiliary!(
        m,
        (m, aux, tmp, geom) ->
            atmos_init_ref_state_pressure!(m.ref_state, m, aux, geom),
        state_auxiliary,
        grid,
        direction,
    )

    ∇p = ∇reference_pressure(m.ref_state, state_auxiliary, grid)

    init_state_auxiliary!(
        m,
        atmos_nodal_init_state_auxiliary!,
        state_auxiliary,
        grid,
        direction;
        state_temporary = ∇p,
    )
end

"""
    source!(
        m::AtmosModel,
        source::Vars,
        state::Vars,
        diffusive::Vars,
        aux::Vars,
        t::Real,
        direction::Direction,
    )

Computes (and assembles) source terms `S(Y)` in:
```
∂Y
-- = - ∇ • F + S(Y)
∂t
```
"""
function source!(
    m::AtmosModel,
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    atmos_source!(m.source, m, source, state, diffusive, aux, t, direction)
end

"""
    init_state_prognostic!(
        m::AtmosModel,
        state::Vars,
        aux::Vars,
        coords,
        t,
        args...,
    )

Initialise state variables. `args...` provides an option to include
configuration data (current use cases include problem constants,
spline-interpolants).
"""
function init_state_prognostic!(
    m::AtmosModel,
    state::Vars,
    aux::Vars,
    coords,
    t,
    args...,
)
    m.problem.init_state_prognostic(
        m.problem,
        m,
        state,
        aux,
        coords,
        t,
        args...,
    )
end

roe_average(ρ⁻, ρ⁺, var⁻, var⁺) =
    (sqrt(ρ⁻) * var⁻ + sqrt(ρ⁺) * var⁺) / (sqrt(ρ⁻) + sqrt(ρ⁺))

function numerical_flux_first_order!(
    numerical_flux::RoeNumericalFlux,
    balance_law::AtmosModel,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    direction,
) where {S, A}
    @assert balance_law.moisture isa DryModel

    numerical_flux_first_order!(
        CentralNumericalFluxFirstOrder(),
        balance_law,
        fluxᵀn,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )

    FT = eltype(fluxᵀn)
    param_set = balance_law.param_set
    _cv_d::FT = cv_d(param_set)
    _T_0::FT = T_0(param_set)

    Φ = gravitational_potential(balance_law, state_auxiliary⁻)

    ρ⁻ = state_prognostic⁻.ρ
    ρu⁻ = state_prognostic⁻.ρu
    ρe⁻ = state_prognostic⁻.ρe
    ts⁻ = recover_thermo_state(
        balance_law,
        balance_law.moisture,
        state_prognostic⁻,
        state_auxiliary⁻,
    )

    u⁻ = ρu⁻ / ρ⁻
    uᵀn⁻ = u⁻' * normal_vector
    e⁻ = ρe⁻ / ρ⁻
    h⁻ = total_specific_enthalpy(ts⁻, e⁻)
    p⁻ = air_pressure(ts⁻)
    c⁻ = soundspeed_air(ts⁻)

    ρ⁺ = state_prognostic⁺.ρ
    ρu⁺ = state_prognostic⁺.ρu
    ρe⁺ = state_prognostic⁺.ρe

    # TODO: state_auxiliary⁺ is not up-to-date
    # with state_prognostic⁺ on the boundaries
    ts⁺ = recover_thermo_state(
        balance_law,
        balance_law.moisture,
        state_prognostic⁺,
        state_auxiliary⁺,
    )

    u⁺ = ρu⁺ / ρ⁺
    uᵀn⁺ = u⁺' * normal_vector
    e⁺ = ρe⁺ / ρ⁺
    h⁺ = total_specific_enthalpy(ts⁺, e⁺)
    p⁺ = air_pressure(ts⁺)
    c⁺ = soundspeed_air(ts⁺)

    ρ̃ = sqrt(ρ⁻ * ρ⁺)
    ũ = roe_average(ρ⁻, ρ⁺, u⁻, u⁺)
    h̃ = roe_average(ρ⁻, ρ⁺, h⁻, h⁺)
    c̃ = sqrt(roe_average(ρ⁻, ρ⁺, c⁻^2, c⁺^2))

    ũᵀn = ũ' * normal_vector

    Δρ = ρ⁺ - ρ⁻
    Δp = p⁺ - p⁻
    Δu = u⁺ - u⁻
    Δuᵀn = Δu' * normal_vector

    w1 = abs(ũᵀn - c̃) * (Δp - ρ̃ * c̃ * Δuᵀn) / (2 * c̃^2)
    w2 = abs(ũᵀn + c̃) * (Δp + ρ̃ * c̃ * Δuᵀn) / (2 * c̃^2)
    w3 = abs(ũᵀn) * (Δρ - Δp / c̃^2)
    w4 = abs(ũᵀn) * ρ̃

    fluxᵀn.ρ -= (w1 + w2 + w3) / 2
    fluxᵀn.ρu -=
        (
            w1 * (ũ - c̃ * normal_vector) +
            w2 * (ũ + c̃ * normal_vector) +
            w3 * ũ +
            w4 * (Δu - Δuᵀn * normal_vector)
        ) / 2
    fluxᵀn.ρe -=
        (
            w1 * (h̃ - c̃ * ũᵀn) +
            w2 * (h̃ + c̃ * ũᵀn) +
            w3 * (ũ' * ũ / 2 + Φ - _T_0 * _cv_d) +
            w4 * (ũ' * Δu - ũᵀn * Δuᵀn)
        ) / 2

    if !(balance_law.tracers isa NoTracers)
        ρχ⁻ = state_prognostic⁻.tracers.ρχ
        χ⁻ = ρχ⁻ / ρ⁻

        ρχ⁺ = state_prognostic⁺.tracers.ρχ
        χ⁺ = ρχ⁺ / ρ⁺

        χ̃ = roe_average(ρ⁻, ρ⁺, χ⁻, χ⁺)
        Δρχ = ρχ⁺ - ρχ⁻

        wt = abs(ũᵀn) * (Δρχ - χ̃ * Δp / c̃^2)

        fluxᵀn.tracers.ρχ -= ((w1 + w2) * χ̃ + wt) / 2
    end
end

"""
    NumericalFluxFirstOrder()
        ::HLLCNumericalFlux,
        balance_law::AtmosModel,
        fluxᵀn,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )

An implementation of the numerical flux based on the HLLC method for
the AtmosModel. For more information on this particular implementation,
see Chapter 10.4 in the provided reference below.

## References
    @book{toro2013riemann,
        title={Riemann solvers and numerical methods for fluid dynamics: a practical introduction},
        author={Toro, Eleuterio F},
        year={2013},
        publisher={Springer Science & Business Media}
    }
"""
function numerical_flux_first_order!(
    ::HLLCNumericalFlux,
    balance_law::AtmosModel,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    direction,
) where {S, A}
    FT = eltype(fluxᵀn)
    num_state_prognostic = number_states(balance_law, Prognostic())
    param_set = balance_law.param_set

    # Extract the first-order fluxes from the AtmosModel (underlying BalanceLaw)
    # and compute normals on the positive + and negative - sides of the
    # interior facets
    flux⁻ = similar(parent(fluxᵀn), Size(3, num_state_prognostic))
    fill!(flux⁻, -zero(FT))
    flux_first_order!(
        balance_law,
        Grad{S}(flux⁻),
        state_prognostic⁻,
        state_auxiliary⁻,
        t,
        direction,
    )
    fluxᵀn⁻ = flux⁻' * normal_vector

    flux⁺ = similar(flux⁻)
    fill!(flux⁺, -zero(FT))
    flux_first_order!(
        balance_law,
        Grad{S}(flux⁺),
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )
    fluxᵀn⁺ = flux⁺' * normal_vector

    # Extract relevant fields and thermodynamic variables defined on
    # the positive + and negative - sides of the interior facets
    ρ⁻ = state_prognostic⁻.ρ
    ρu⁻ = state_prognostic⁻.ρu
    ρe⁻ = state_prognostic⁻.ρe
    ts⁻ = recover_thermo_state(
        balance_law,
        balance_law.moisture,
        state_prognostic⁻,
        state_auxiliary⁻,
    )

    u⁻ = ρu⁻ / ρ⁻
    c⁻ = soundspeed_air(ts⁻)

    uᵀn⁻ = u⁻' * normal_vector
    e⁻ = ρe⁻ / ρ⁻
    p⁻ = air_pressure(ts⁻)

    ρ⁺ = state_prognostic⁺.ρ
    ρu⁺ = state_prognostic⁺.ρu
    ρe⁺ = state_prognostic⁺.ρe
    ts⁺ = recover_thermo_state(
        balance_law,
        balance_law.moisture,
        state_prognostic⁺,
        state_auxiliary⁺,
    )

    u⁺ = ρu⁺ / ρ⁺
    uᵀn⁺ = u⁺' * normal_vector
    e⁺ = ρe⁺ / ρ⁺
    p⁺ = air_pressure(ts⁺)
    c⁺ = soundspeed_air(ts⁺)

    # Wave speeds estimates S⁻ and S⁺
    S⁻ = min(uᵀn⁻ - c⁻, uᵀn⁺ - c⁺)
    S⁺ = max(uᵀn⁻ + c⁻, uᵀn⁺ + c⁺)

    # Compute the middle wave speed S⁰ in the contact/star region
    S⁰ =
        (p⁺ - p⁻ + ρ⁻ * uᵀn⁻ * (S⁻ - uᵀn⁻) - ρ⁺ * uᵀn⁺ * (S⁺ - uᵀn⁺)) /
        (ρ⁻ * (S⁻ - uᵀn⁻) - ρ⁺ * (S⁺ - uᵀn⁺))

    p⁰ =
        (
            p⁺ +
            p⁻ +
            ρ⁻ * (S⁻ - uᵀn⁻) * (S⁰ - uᵀn⁻) +
            ρ⁺ * (S⁺ - uᵀn⁺) * (S⁰ - uᵀn⁺)
        ) / 2

    # Compute p * D = p * (0, n₁, n₂, n₃, S⁰)
    pD = @MVector zeros(FT, num_state_prognostic)
    if balance_law.ref_state isa HydrostaticState
        # pressure should be continuous but it doesn't hurt to average
        ref_p⁻ = state_auxiliary⁻.ref_state.p
        ref_p⁺ = state_auxiliary⁺.ref_state.p
        ref_p⁰ = (ref_p⁻ + ref_p⁺) / 2

        momentum_p = p⁰ - ref_p⁰
    else
        momentum_p = p⁰
    end

    pD[2] = momentum_p * normal_vector[1]
    pD[3] = momentum_p * normal_vector[2]
    pD[4] = momentum_p * normal_vector[3]
    pD[5] = p⁰ * S⁰

    # Computes both +/- sides of intermediate flux term flux⁰
    flux⁰⁻ =
        (S⁰ * (S⁻ * parent(state_prognostic⁻) - fluxᵀn⁻) + S⁻ * pD) / (S⁻ - S⁰)
    flux⁰⁺ =
        (S⁰ * (S⁺ * parent(state_prognostic⁺) - fluxᵀn⁺) + S⁺ * pD) / (S⁺ - S⁰)

    if 0 <= S⁻
        parent(fluxᵀn) .= fluxᵀn⁻
    elseif S⁻ < 0 <= S⁰
        parent(fluxᵀn) .= flux⁰⁻
    elseif S⁰ < 0 <= S⁺
        parent(fluxᵀn) .= flux⁰⁺
    else # 0 > S⁺
        parent(fluxᵀn) .= fluxᵀn⁺
    end
end

end # module


#=##############################################################################
# DESCRIPTION
    Vehicles based on FLOWVLM (VLM and blade-element) models.

# AUTHORSHIP
  * Author    : Eduardo J. Alvarez
  * Email     : Edo.AlvarezR@gmail.com
  * Created   : Apr 2020
  * License   : MIT
=###############################################################################

################################################################################
# ABSTRACT VLM VEHICLE TYPE
################################################################################
"""
    `AbstractVLMVehicle{N, M, R} <: AbstractVehicle{N, M, R}`

Implementations of `AbstractUVLMVehicle` are expected to have the following
fields:

# Required inputs
* `system::vlm.WingSystem`

# Optional inputs
* `tilting_systems::NTuple{N, vlm.WingSystem}`
* `rotor_systems::NTuple{M, Array{vlm.Rotor, 1}}`
* `vlm_system::vlm.WingSystem`
* `wake_system::vlm.WingSystem`
* `grids::Array{gt.GridTypes, 1}`

# Internal properties
* `V::Array{R, 1}`                          # Current vehicle velocity
* `W::Array{R, 1}`                          # Current vehicle angular velocity
* `prev_data::Array{Any, 1}`                # Information about previous step
* `grid_O::Array{Array{R, 1}, 1}`           # Origin of every grid

They also need to implement the following functions required by
`AbstractVehicle`:

* `shed_wake(...)`
* `generate_static_particle_fun(...)`
* `save_vtk(...)`

See the documentation and code of `AbstractVehicle` for more details.

In general, in order for implementation to work correctly, it is required that
the components of `wake_system` are also components of either `vlm_system` or
`rotor_systems`, and that none of the components of `rotor_systems` are also
components of `vlm_system`.
"""
abstract type AbstractVLMVehicle{N, M, R} <: AbstractVehicle{N, M, R} end

# Implementations
for header_name in ["unsteady", "quasisteady"]
  include("FLOWUnsteady_vehicle_vlm_"*header_name*".jl")
end

##### FUNCTIONS  ###############################################################
get_ntltsys(self::AbstractVLMVehicle) = typeof(self).parameters[1]

get_nrtrsys(self::AbstractVLMVehicle) = typeof(self).parameters[2]

function add_dV(self::AbstractVLMVehicle, dV)
    self.V .+= dV
    return nothing
end

function add_dW(self::AbstractVLMVehicle, dW)
    self.W .+= dW
    return nothing
end

function set_V(self::AbstractVLMVehicle, V)
    self.V .= V
    return nothing
end

function set_W(self::AbstractVLMVehicle, W)
    self.W .= W
    return nothing
end

function tilt_systems(self::AbstractVLMVehicle{N,M,R}, angles::NTuple{N, Array{R2, 1}}
                                                    ) where{N, M, R, R2<:Real}
    # Iterate over tilting system
    for i in 1:get_ntltsys(self)
        sys = self.tilting_systems[i]
        Oaxis = gt.rotation_matrix2([-a for a in angles[i]]...)
        vlm.setcoordsystem(sys, sys.O, Oaxis)
    end
    return nothing
end
function tilt_systems(self::AbstractVLMVehicle{0,M,R}, angles::Tuple{})  where{M, R}
    return nothing
end

function nextstep_kinematic(self::AbstractVLMVehicle, dt::Real)
    dX = dt*self.V                  # Translation
    dA = 180/pi * dt*self.W         # Angular rotation (degrees)

    O = self.system.O + dX          # New origin of the system (translation)
    M = gt.rotation_matrix2([-a for a in dA]...) # Rotation matrix
    Oaxis = M*self.system.Oaxis     # New axes of the system (rotation)

    # Save state of current vehicle
    _update_prev_vlm_system(self, deepcopy(self.vlm_system))
    _update_prev_wake_system(self, deepcopy(self.wake_system))
    _update_prev_rotor_systems(self, deepcopy(self.rotor_systems))

    # Translation and rotation
    vlm.setcoordsystem(self.system, O, Oaxis)

    # Translation and rotation of every grid
    for i in 1:length(self.grids)
        self.grid_O[i] .+= dX

        # Translation
        gt.lintransform!(self.grids[i], Array(1.0I, 3, 3) , dX)

        # Brings the grid back to the global origin
        gt.lintransform!(self.grids[i], Array(1.0I, 3, 3) , -self.grid_O[i])

        # Rotation and brings the grid back to its position
        gt.lintransform!(self.grids[i], M, self.grid_O[i])
    end

    return nothing
end

"""
Precalculations before calling the solver.

Calculates kinematic velocity on VLM an adds them as a solution field
"""
function precalculations(self::AbstractVLMVehicle, Vinf::Function,
                                pfield::vpm.ParticleField, t::Real, dt::Real,
                                nt::Int)

    if nt!=0
        # ---------- 1) Recalculate wake horseshoes with kinematic velocity -
        # Calculate kinematic velocity at every control point
        Vkin = _Vkinematic_wake(self, dt; t=t, targetX="CP")
        vlm._addsolution(self.wake_system, "Vkin", Vkin; t=t)

        Vkin = _Vkinematic_vlm(self, dt; t=t, targetX="CP")
        vlm._addsolution(self.vlm_system, "Vkin", Vkin; t=t)

        # Recalculate horseshoes
        vlm.getHorseshoes(self.vlm_system; t=t, extraVinf=_extraVinf1)
        vlm.getHorseshoes(self.wake_system; t=t, extraVinf=_extraVinf1)

        # ---------- 2) Paste previous Gamma solution ----------------------
        for i in 1:length(self.wake_system.wings)
            wing = vlm.get_wing(self.wake_system, i)
            prev_wing = vlm.get_wing(_get_prev_wake_system(self), i)

            if typeof(prev_wing) != vlm.Rotor
                sol = deepcopy(prev_wing.sol["Gamma"])
            else
                sol = deepcopy(prev_wing._wingsystem.sol["Gamma"])
            end

            vlm._addsolution(wing, "Gamma", sol; t=t)
        end
        for i in 1:length(self.vlm_system.wings)
            wing = vlm.get_wing(self.vlm_system, i)
            prev_wing = vlm.get_wing(_get_prev_vlm_system(self), i)

            if typeof(prev_wing) != vlm.Rotor
                sol = deepcopy(prev_wing.sol["Gamma"])
            else
                sol = deepcopy(prev_wing._wingsystem.sol["Gamma"])
            end

            vlm._addsolution(wing, "Gamma", sol; t=t)
        end
    end

    return nothing
end

function save_vtk_base(self::AbstractVLMVehicle, filename; path=nothing,
                                   num=nothing, save_wopwopin=false, optargs...)
    strn = vlm.save(self.system, filename; path=path, num=num, optargs...)

    for (i, grid) in enumerate(self.grids)
        strn *= gt.save(grid, filename*"_Grid$i"; path=path, num=num)
    end

    # Generate inputs for PSU-WOPWOP
    if save_wopwopin
        for (si, rotors) in enumerate(self.rotor_systems)
            # Compact patch for loading
            generate_vtkliftinglines(rotors, filename*"_Sys$(si)", path; num=num,
                                                                suf="_compact")
            # Loft for thickness
            # NOTE: this is not needed since FLOWNoise can read the lofted VTK
            # straight up
            # for (ri, rotors) enumerate(rotors)
            #     vlm.save(rotor, filename*"_Sys$(si)_Rotor$(ri)"; path=path, addtiproot=true,
            #                             wopwop=true, wopbin=true, wopv=1.0, num=num)
            # end

            # Loading file
            save_gammas = [[vlm.get_blade(rotor, i).sol["Gamma"]
                                      for i in 1:rotor.B] for rotor in rotors]
            # NOTE TO SELF: Ftot is a force per unit length
            save_Ftot = [[vlm.get_blade(rotor, i).sol["Ftot"]
                                      for i in 1:rotor.B] for rotor in rotors]
            loadfname = "loading_Sys$(si)"*(num!=nothing ? ".$num" : "")*".jld"
            JLD.save(joinpath(path, loadfname), "rotorgammas", save_gammas,
                                                        "Ftot", save_Ftot)
        end
    end

    return strn
end

##### INTERNAL FUNCTIONS  ######################################################

_get_prev_vlm_system(self::AbstractVLMVehicle) = self.prev_data[1]
_get_prev_wake_system(self::AbstractVLMVehicle) = self.prev_data[2]
_get_prev_rotor_systems(self::AbstractVLMVehicle) = self.prev_data[3]
function _update_prev_vlm_system(self::AbstractVLMVehicle, system)
    self.prev_data[1] = system
end
function _update_prev_wake_system(self::AbstractVLMVehicle, system)
    self.prev_data[2] = system
end
function _update_prev_rotor_systems(self::AbstractVLMVehicle, rotor_systems)
    self.prev_data[3] = rotor_systems
end

"""
Return the maximum number of static particles.
"""
function _get_m_static(self::AbstractVLMVehicle)
    m = vlm.get_m(self.vlm_system)

    for rotors in self.rotor_systems
        for rotor in rotors
            m += vlm.get_m(rotor)
        end
    end

    return m
end

"""
Returns the local translational velocity of every control point in `vlm_system`.
"""
function _Vkinematic_vlm(self::AbstractVLMVehicle, args...; optargs...)
    return _Vkinematic(self.vlm_system, _get_prev_vlm_system(self), args...;
                                                                    optargs...)
end

"""
Returns the local translational velocity of every control point in `wake_system`.
"""
function _Vkinematic_wake(self::AbstractVLMVehicle, args...; optargs...)
    return _Vkinematic(self.wake_system, _get_prev_wake_system(self), args...;
                                                                    optargs...)
end

function _Vkinematic(system, prev_system, dt::Real; t=0.0, targetX="CP")

    cur_Xs = _get_Xs(system, targetX; t=t)
    prev_Xs = _get_Xs(prev_system, targetX; t=t)

    return [-(cur_Xs[i]-prev_Xs[i])/dt for i in 1:size(cur_Xs, 1)]
end

"""
Returns the local translational velocity of every midpoint in the ri-th rotor in
the si-th system formatted as the output of `_get_midXs()`.
"""
function _Vkinematic_rotor(rotor_systems::NTuple{M, Array{vlm.Rotor, 1}},
                           prev_rotor_systems::NTuple{M, Array{vlm.Rotor, 1}},
                           si, ri, dt::Real
                          ) where{M}

    cur_Xs = _get_midXs(rotor_systems[si][ri])
    prev_Xs = _get_midXs(prev_rotor_systems[si][ri])

    Vkin = [-(cur_Xs[i]-prev_Xs[i])/dt for i in 1:size(cur_Xs, 1)]

    return Vkin

end

function _get_Xs(system::Union{vlm.Wing, vlm.WingSystem, vlm.Rotor},
                                                        targetX::String; t=0.0)
    if targetX=="CP"
        Xs = [vlm.getControlPoint(system, i
                                        ) for i in 1:vlm.get_m(system)]
    else
        Xs = [vlm.getHorseshoe(system, i; t=t)[vlm.VLMSolver.HS_hash[targetX]
                                        ] for i in 1:vlm.get_m(system)]
    end
    return Xs
end

function _get_Xs(system::Union{vlm.Wing, vlm.WingSystem, vlm.Rotor},
                                        targetXs::Array{T,1}; t=0.0) where{T}
    return vcat([_get_Xs(system, targetX; t=t) for targetX in targetXs]...)
end

"Returns the i-th bound-vortex midpoint of the j-th blade in `rotor`"
function _get_midX(rotor::vlm.Rotor, j::Int, i::Int)
    return (vlm.getHorseshoe(vlm.get_blade(rotor, j), i)[2] +
            vlm.getHorseshoe(vlm.get_blade(rotor, j), i)[3]) / 2
end

"Returns all bound-vortex midpoints of the j-th blade in `rotor`"
function _get_midXs(rotor::vlm.Rotor, j::Int)
    return [_get_midX(rotor, j, i)
                            for i in 1:vlm.get_m(vlm.get_blade(rotor, j)) ]
end
"Returns all bound-vortex midpoints in `rotor`"
_get_midXs(rotor::vlm.Rotor) = vcat([_get_midXs(rotor, j) for j in 1:rotor.B]...)

"Returns all bound-vortex midpoints in an array of rotors"
_get_midXs(rotors::Array{vlm.Rotor, 1}) = vcat([_get_midXs(rotor) for rotor in rotors]...)

"Returns all bound-vortex midpoints in multiple rotor systems"
function _get_midXs(rotor_systems::NTuple{M, Array{vlm.Rotor, 1}}) where{M}
    return vcat([_get_midXs(rotors) for rotors in rotor_systems]...)
end

"Receives an array `midXs` formatted as the output of `_get_midXs(rotor_systems)`
and returns the section of the array that corresponds to the ri-th rotor in the
si-th system (formatted linearly).
"
function _parse_midXs(rotor_systems::NTuple{M, Array{vlm.Rotor, 1}},
                        midXs::Array{Array{R, 1}}, si, ri) where{R<:Real, M}

    if si>length(rotor_systems) || si<=0
        error("Got invalid system index $si; max is $(length(rotor_systems))")
    end

    # Find lower bound of midXs belonging to rotor ri
    i_low = 1
    for (sj, rotors) in enumerate(rotor_systems[1:si])       # Iterate over rotor systems
        for (rj, rotor) in enumerate(rotors)                 # Iterative over rotors
            if sj==si && rj==ri
                break
            end
            i_low += vlm.get_m(rotor)
        end
    end

    # Find upper bound of midXs belonging to rotor ri
    rotor = rotor_systems[si][ri]
    i_up = i_low-1 + vlm.get_m(rotor)

    # # Format it as blades
    out = midXs[i_low:i_up]
    # out = [out[ ((i-1)*rotor.m+1):(i*rotor.m) ] for i in 1:rotor.B]
    return out
end

"""
    Receives the output of `parse_midXs(rotor_systems, midXs, si, ri)` and
outputs the array formatted as blades.
"""
function _format_blades(arr::Array{Array{R, 1}, 1},
                        rotor_systems::NTuple{M, Array{vlm.Rotor, 1}},
                        si, ri) where{R, M}

    rotor = rotor_systems[si][ri]
    arr = [arr[ ((i-1)*rotor.m+1):(i*rotor.m) ] for i in 1:rotor.B]

    return arr
end

"""
Function for FLOWVLM to fetch and add extra velocities to what every
control point imposes as the boundary condition
"""
function _extraVinf1(i, t; wing=nothing)
    if wing==nothing; error("Logic error!"); end;
    return wing.sol["Vkin"][i]
end
function _extraVinf2(i, t; wing=nothing)
    if wing==nothing; error("Logic error!"); end;
    return wing.sol["Vvpm"][i] + wing.sol["Vkin"][i]
end
function _extraVinf3(i, t; wing=nothing)
    if wing==nothing; error("Logic error!"); end;
    return wing.sol["Vind"][i] + wing.sol["Vkin"][i]
end

"""
Receives the FLOWVLM object `system` (Wing/WingSystem/Rotor), and adds vortex
particles to the particle field `pfield` at each trailing edge position where
an infinite vortex starts. `dt` indicates the length of the interval of time
that the vortex shedding represents.

Give it a previous system to detect differences in circulation and add
unsteady particles.
"""
function VLM2VPM(system::Union{vlm.Wing, vlm.WingSystem, vlm.Rotor}, pfield, dt,
                    Vinf;
                    t=0.0, prev_system=nothing, unsteady_shedcrit=-1.0,
                    p_per_step=1, sigmafactor=1.0, overwrite_sigma=nothing,
                    check=true, debug=false, tol=1e-6, omit_shedding=[])

  m = vlm.get_m(system)   # Number of lattices

  # Velocity at horseshoe points Ap and Bp
  if prev_system==nothing
      Vinfs_Ap = vlm.getVinfs(system; t=t, target="Ap")
      Vinfs_Bp = vlm.getVinfs(system; t=t, target="Bp")
  else
      Vinfs_Ap = []
      Vinfs_Bp = []
      for i in 1:m
          Ap = vlm.getHorseshoe(system, i)[1]
          prev_Ap = vlm.getHorseshoe(prev_system, i)[1]
          push!(Vinfs_Ap, -(Ap-prev_Ap)/dt + Vinf(Ap, t))

          Bp = vlm.getHorseshoe(system, i)[4]
          prev_Bp = vlm.getHorseshoe(prev_system, i)[4]
          push!(Vinfs_Bp, -(Bp-prev_Bp)/dt + Vinf(Bp, t))
      end
  end

  # Adds a particle at each infinite vortex
  prev_HS = [nothing for i in 1:8]
  for i in 1:m  # Iterates over lattices
    HS = vlm.getHorseshoe(system, i)

    Ap, A, B, Bp, CP, infDA, infDB, Gamma = HS
    (prev_Ap, prev_A, prev_B, prev_Bp, prev_CP,
      prev_infDA, prev_infDB, prev_Gamma) = prev_HS
    cntgs = true          # Contiguous horseshoes flag

    if nothing in HS; error("Logic error! $HS"); end;
    if true in [isnan.(elem) for elem in HS]; error("Logic error! $HS"); end;

    # ----------- Case of left wing tip ---------------------------
    if i==1
      # Adds particle at Ap
      X = Ap                                    # Particle position
      gamma = Gamma                             # Infinite vortex circulation
      V = norm(Vinfs_Ap[i])                     # Freestream at X
      infD = -infDA                             # Direction of vorticity
      sigma = sigmafactor*V*dt                  # Vortex blob radius
      vol = pi*(norm(Bp-Ap)/2)^2*V*dt           # Volume of particle
      l = -infD*V*dt                             # Distance the TE travels

      # Check if the system loops around
      if norm(Ap - vlm.getHorseshoe(system, m)[4]) / norm(Bp - Ap) <= tol
          gamma -= vlm.getHorseshoe(system, m)[8]
      end

      if unsteady_shedcrit<=0
          if !(i in omit_shedding)
              add_particle(pfield, X, gamma, dt, V, infD, sigma, vol,
                                l, p_per_step; overwrite_sigma=overwrite_sigma)
          end
      end

    # ----------- Case of wing tip on discontinuous wing --------
    elseif norm(Ap - prev_Bp) / norm(Bp - Ap) > tol
      if debug; println("Wing tip found at i=$i. $Ap!=$prev_Bp"); end;

      cntgs = false

      # Adds particle at previous Bp
      X = prev_Bp                               # Particle position
      gamma = -prev_Gamma                       # Infinite vortex circulation
      V = norm(Vinfs_Bp[i-1])                   # Freestream at X
      infD = -prev_infDB                        # Direction of vorticity
      sigma = sigmafactor*V*dt                  # Vortex blob radius
      vol = pi*(norm(prev_Bp-prev_Ap)/2)^2*V*dt # Volume of particle
      l = -infD*V*dt                             # Distance the TE travels

      if unsteady_shedcrit<=0
          if !(i-1 in omit_shedding)
              add_particle(pfield, X, gamma, dt, V, infD, sigma, vol,
                              l, p_per_step; overwrite_sigma=overwrite_sigma)
          end
      end

      # Adds particle at Ap
      X = Ap                                    # Particle position
      gamma = Gamma                             # Infinite vortex circulation
      V = norm(Vinfs_Ap[i])                     # Freestream at X
      infD = -infDA                             # Direction of vorticity
      sigma = sigmafactor*V*dt                  # Vortex blob radius
      vol = pi*(norm(Bp-Ap)/2)^2*V*dt           # Volume of particle
      l = -infD*V*dt                             # Distance the TE travels

      if unsteady_shedcrit<=0
          if !(i in omit_shedding)
              add_particle(pfield, X, gamma, dt, V, infD, sigma, vol,
                                l, p_per_step; overwrite_sigma=overwrite_sigma)
          end
      end

    # ----------- Case of contiguous horseshoes -----------------
    else

      # Verify logic
      if check
        crit1 = norm(Vinfs_Ap[i]-Vinfs_Bp[i-1]) / norm((Vinfs_Ap[i]+Vinfs_Bp[i-1])/2)
        crit2 = norm(infDA-prev_infDB) / norm((infDA+prev_infDB)/2)
        if crit1 > tol
          error("Logic error! Vinfs_Ap[i]!=Vinfs_Bp[i-1] "*
                  "( $(Vinfs_Ap[i])!=$(Vinfs_Bp[i-1]) )")
        elseif crit2 > tol
          error("Logic error! infDA!=prev_infDB "*
                  "( $(infDA)!=$(prev_infDB) )")
        elseif norm(infDA) - 1 > tol
          error("Logic error! norm(infDA)!= 1 "*
                  "( $(norm(infDA)) )")
        end
      end

      # Adds particle at Ap
      X = Ap                                    # Particle position
      gamma = Gamma - prev_Gamma                # Infinite vortex circulation
      V = norm(Vinfs_Ap[i])                     # Freestream at X
      infD = -infDA                             # Direction of vorticity
      sigma = sigmafactor*V*dt                  # Vortex blob radius
      vol = pi*(norm(Bp-Ap)/2)^2*V*dt           # Volume of particle
      l = -infD*V*dt                             # Distance the TE travels

      if unsteady_shedcrit<=0
          if !(i in omit_shedding)
              add_particle(pfield, X, gamma, dt, V, infD, sigma, vol,
                              l, p_per_step; overwrite_sigma=overwrite_sigma)
          end
      end

      # ----------- Case of discontinuous right wing tip -----------------------
      if i==m && norm(Bp - vlm.getHorseshoe(system, 1)[1]) / norm(Bp - Ap) > tol
        # Adds particle at Bp
        X = Bp                                    # Particle position
        gamma = -Gamma                            # Infinite vortex circulation
        V = norm(Vinfs_Bp[i])                     # Freestream at X
        infD = -infDB                             # Direction of vorticity
        sigma = sigmafactor*V*dt                  # Vortex blob radius
        vol = pi*(norm(Bp-Ap)/2)^2*V*dt           # Volume of particle
        l = -infD*V*dt                             # Distance the TE travels

        if unsteady_shedcrit<=0
            if !(i in omit_shedding)
                add_particle(pfield, X, gamma, dt, V, infD, sigma, vol,
                            l, p_per_step; overwrite_sigma=overwrite_sigma)
            end
        end
      end
    end

    # Adds unsteady particle
    if prev_system!=nothing && unsteady_shedcrit>0

      (p_Ap, p_A, p_B, p_Bp, p_CP,
          p_infDA, p_infDB, p_Gamma) = vlm.getHorseshoe(prev_system, i)
      X = (Ap+Bp)/2                               # LE midpoint position
      p_X = (p_Ap+p_Bp)/2                         # Previous LE midpoint position
      gamma = Gamma - p_Gamma                     # Bound circulation increase
      infD = (!cntgs || i==1 ? Ap : (prev_Ap+prev_Bp)/2) - X # Direction of circulation
      sigma = sigmafactor*norm(B-A)               # Vortex blob radius
      vol = 4/3*pi*(sigma/2)^3                    # Volume of particle
      l = -(X-p_X) + Vinf(X, t)*dt                # Distance the TE travels

      # Adds particle only if difference is greater than 1%
      if abs(gamma/p_Gamma) > unsteady_shedcrit && !(i in omit_shedding)
        add_particle(pfield, X, gamma, 1.0, 1.0, infD, sigma, vol,
                                        l, 1; overwrite_sigma=overwrite_sigma)
      end
    end

    prev_HS = HS
  end
end
##### END OF ABSTRACT VLM VEHICLE ##############################################