text:
"""
`EntropicSwitch` - Entropically switching potential with three local minima.  It
is symmetric about x=0

### Fields
* `x` - Position x in R²
"""
function EntropicSwitch(x)
    return (3 * exp(-x[1]^2 - (x[2]-1/3)^2)
     - 3 * exp(-x[1]^2 - (x[2]-5/3)^2)
     - 5 * exp(-(x[1]-1)^2 - x[2]^2)
     - 5 * exp(-(x[1]+1)^2 - x[2]^2)
     + 1/5 * x[1]^4 + 1/5 * (x[2]-1/3)^4);
end

"""
`SymmetricTwoChannel` - Double well potential in 2D with two, symmetric channels
joining them.

### Fields
* `x` - Position x in R²
"""
function SymmetricTwoChannel(x)
    return 1/6 * (4 * (1-x[1]^2-x[2]^2)^2 + 2 *(x[1]^2-2)^2
    + ((x[1]+x[2])^2 - 1 )^2+ ((x[1]-x[2])^2 - 1 )^2);
end

"""
`Muller` - The Muller potential with three distinct minima and highy asymmetric.

### Fields
* `x` - Position x in R²
"""
function Muller(x)

    aa = (-1, -1, -6.5, 0.7);
    bb = (0., 0., 11., 0.6);
    cc = (-10., -10., -6.5, 0.7);
    AA = (-200., -100., -170., 15.);
    XX = (1., 0., -0.5, -1.);
    YY = (0., 0.5, 1.5, 1.);

    return ( AA[1]*exp(aa[1]*(x[1]-XX[1])^2+bb[1]*(x[1]-XX[1])*(x[2]-YY[1])+cc[1]*(x[2]-YY[1])^2)
                 +AA[2]*exp(aa[2]*(x[1]-XX[2])^2+bb[2]*(x[1]-XX[2])*(x[2]-YY[2])+cc[2]*(x[2]-YY[2])^2)
                 +AA[3]*exp(aa[3]*(x[1]-XX[3])^2+bb[3]*(x[1]-XX[3])*(x[2]-YY[3])+cc[3]*(x[2]-YY[3])^2)
                 +AA[4]*exp(aa[4]*(x[1]-XX[4])^2+bb[4]*(x[1]-XX[4])*(x[2]-YY[4])+cc[4]*(x[2]-YY[4])^2));

end

"""
`Rosenbrock` - Banana shaped Rosenbrock potentials with global minimum is located
at (a,a²).

### Fields
* `x` - Position x in R²
### Optional Fields
* `a = 1.0` - Rosenbrock parameter
* `b = 100.0` - Rosenbrock parameter
"""
function Rosenbrock(x; a = 1.0, b = 100.0)
    return (a-x[1])^2 + b * (x[2]-x[1]^2)^2

end

"""
`Zpotential` - Z shaped potential.

### Fields
* `x` - Position x in R²
"""
function Zpotential(x)
    return (x[1]^4 + x[2]^4)/20480 -
    3 * exp(-0.01 * (x[1]+5)^2 -0.2 * (x[2]+5)^2) -
    3 * exp(-0.01 * (x[1]-5)^2 -0.2 * (x[2]-5)^2) +
    5 * exp(-0.2 * (x[1]+3*(x[2]-3))^2)/(1+exp(-x[1]-3)) +
    5 * exp(-0.2 * (x[1]+3*(x[2]+3))^2)/(1+exp(x[1]-3)) +
    3 * exp(-0.01 *(x[1]^2 + x[2]^2))
end

"""
`EntropicBox` - A potential concentrated in [0,1]² with internal entropic
barriers.  Formulated by D. Aristoff (Colorado State).

### Fields
* `x` - Position x in R²
"""
function EntropicBox(x)

    c = (50.5, 49.5, 10^5, 51, 49);

    return exp(-(c[1]*(x[1]-0.25).^2+c[1]*(x[2]-0.75).^2+2*c[2]*(x[1]-0.25).*(x[2]-0.75))) +
            exp(-c[3]*(x[1]^2*(1-x[1])^2*x[2]^2*(1-x[2])^2)) +
            0.5*exp(-(c[4]*x[1]^2+c[4]*x[2]^2-2*c[5]*x[1]*x[2]));
end

max_stars_repo_name: gideonsimpson/TestLandscapes

max_issues_repo_name: gideonsimpson/TestLandscapes

max_forks_repo_name: gideonsimpson/TestLandscapes


text:
"""
`module  JAC.StrongField
    ... a submodel of JAC that contains all methods to set-up and perform strong-field computations.
"""
module StrongField

    using    GSL, ## HypergeometricFunctions, Plots,
             Printf, SpecialFunctions, ..AngularMomentum, ..Basics, ..Defaults, ..InteractionStrength, ..Radial, ..ManyElectron,
             ..Nuclear, ..Pulse, ..TableStrings

    export   aaaa

    """
    `abstract type StrongField.AbstractSFAObservable`
        ... defines an abstract and a number of types for the observables that can be computed with SFA amplitudes.

        + struct SfaNoObservable            ... an empty instance of an observable that does not help compute anything.
        + struct SfaEnergyDistribution      ... to compute the energy distribution of the photoelectrons.
        + struct SfaMomentumDistribution    ... to compute the momentum distribution of the photoelectrons.
    """
    abstract type  AbstractSFAObservable                                 end
    struct   SfaNoObservable     <: StrongField.AbstractSFAObservable    end

    """
    `struct  StrongField.SfaEnergyDistribution  <: StrongField.AbstractSFAObservable`
        ... to compute in SFA the energy distribution of photoelectrons at given energies and angles.

        + theta         ::Float64             ... polar angle of the energy distribution
        + phi           ::Float64             ... azimuthal angle of the energy distribution
        + energies      ::Array{Float64,1}    ... specifies the photoelectron energy in the current units.
    """
    struct   SfaEnergyDistribution  <: StrongField.AbstractSFAObservable
        theta           ::Float64
        phi             ::Float64
        energies        ::Array{Float64,1}
    end

    """
    `StrongField.SfaEnergyDistribution(theta::Float64, phi::Float64, NoEnergies::Int64, maxEnergy::Float64)`
        ... defines an energy distribution for given (theta,phi) and for NoEnergies between 0 < energy <= maxEnergy;
            a dist::SfaEnergyDistribution is returned.
    """
    function SfaEnergyDistribution(theta::Float64, phi::Float64, NoEnergies::Int64, maxEnergy::Float64)
        energies = Float64[];      for  i=1:NoEnergies   push!(energies, i*maxEnergy / NoEnergies)   end
        SfaEnergyDistribution(theta, phi, energies)
    end

    function Base.string(obs::SfaEnergyDistribution)
        sa = "Compute photo-electron energy distribution at angles (theta,phi) = ($(obs.theta), $(obs.phi)) for the energies [a.u.]:"
        return( sa )
    end

    function Base.show(io::IO, obs::SfaEnergyDistribution)
        sa = string(obs);       print(io, sa, "\n");     print(io, "   ", obs.energies)
    end

    """
    `struct  StrongField.SfaMomentumDistribution  <: StrongField.AbstractSFAObservable`
        ... to compute in SFA the momentum distribution of photoelectrons for given polar angle theta and at given energies and azimuthal angles.

        + theta         ::Float64             ... polar angle of the energy distribution
        + phi           ::Array{Float64,1}    ... specifies the azimuthal angles
        + energies      ::Array{Float64,1}    ... specifies the photoelectron energy in the current units.
    """
    struct   SfaMomentumDistribution  <: StrongField.AbstractSFAObservable
        theta           ::Float64
        phis            ::Array{Float64,1}
        energies        ::Array{Float64,1}
    end

    """
    `StrongField.SfaMomentumDistribution(theta::Float64, NoPhi::Int64, NoEnergies::Int64, maxEnergy::Float64)`
        ... defines a momentum distribution for given theta and for NoEnergies between 0 < energy <= maxEnergy and NoPhi between 0 <= phi < 2pi;
            a dist::SfaMomentumDistribution is returned.
    """
    function SfaMomentumDistribution(theta::Float64, NoPhi::Int64, NoEnergies::Int64, maxEnergy::Float64)
        energies = Float64[];      for  i=1:NoEnergies   push!(energies, i*maxEnergy / NoEnergies)   end
        phis = Float64[];          for  i=1:NoPhi        push!(phis, i*2pi / NoPhi)                  end
        SfaMomentumDistribution(theta, phis, energies)
    end

    function Base.string(obs::SfaMomentumDistribution)
        sa = "Compute photo-electron momentum distribution at polar angle theta = $(obs.theta) for (energy [a.u.], phi):"
        return( sa )
    end

    function Base.show(io::IO, obs::SfaMomentumDistribution)
        sa = string(obs);       print(io, sa, "\n");     print(io, "   ", obs.energies);    print(io, "   ", obs.phi)
    end

    """
    `abstract type StrongField.AbstractVolkovState`
        ... defines an abstract and a number of types for dealing with the Volkov states in the computation of the SFA amplitudes.

        + struct FreeVolkov       ... to apply the free-Volkov states.
        + struct CoulombVolkov    ... apply the Coulomb-Volkov states.
        + struct DistortedVolkov  ... apply the Distorted-Volkov states.
    """
    abstract type  AbstractVolkovState                         end
    struct         FreeVolkov        <:  AbstractVolkovState   end

    """
    `struct StrongField.CoulombVolkov     <:  StrongField.AbstractVolkovState`

        + Z ::Float64   ... Charge that generates the Coulomb potential
    """
    struct  CoulombVolkov   <: StrongField.AbstractVolkovState
        Z   ::Float64
    end

    struct         DistortedVolkov   <:  AbstractVolkovState   end

    function Base.string(volkov::FreeVolkov)         return( "free-Volkov states" )         end
    function Base.string(volkov::CoulombVolkov)      return( "Coulomb-Volkov states in potential of charge Z = $(volkov.Z)" )      end
    function Base.string(volkov::DistortedVolkov)    return( "distorted-Volkov states" )    end

    """
    `struct  StrongField.Settings`
        ... defines a type for the details and parameters of computing SFA amplitudes.

        + multipoles            ::Array{EmMultipoles}  ... Multipoles of the radiation field that are to be included.
        + gauges                ::Array{UseGauge}      ... Specifies the gauges to be included into the computations.
        + printBefore           ::Bool                 ... True, if all the requested amplitudes are printed before the computation start.
        + printAmplitudes       ::Bool                 ... True, if the amplitudes are to be printed.
    """
    struct Settings
        multipoles              ::Array{EmMultipole,1}
        gauges                  ::Array{UseGauge}
        printBefore             ::Bool
        printAmplitudes         ::Bool
    end

    """
    `StrongField.Settings()`
        ... constructor for the default values of SFA computations.
    """
    function Settings()
        Settings([E1], [UseCoulomb], false, false)
    end

    # `Base.show(io::IO, settings::StrongField.Settings)`  ... prepares a proper printout of the variable settings::StrongField.Settings.
    function Base.show(io::IO, settings::StrongField.Settings)
        println(io, "multipoles:                 $(settings.multipoles)  ")
        println(io, "use-gauges:                 $(settings.gauges)  ")
        println(io, "printBefore:                $(settings.printBefore)  ")
        println(io, "printAmplitudes:            $(settings.printAmplitudes)  ")
    end

    """
    `struct  StrongField.Computation`
        ... defines a type for the computation of strong-field amplitudes and observables (properties).

        + observable         ::AbstractSFAObservable        ... The obserable to be calculated in SFA.
        + nuclearModel       ::Nuclear.Model                ... Model, charge and parameters of the nucleus.
        + grid               ::Radial.Grid                  ... The radial grid to be used for the computation.
        + initialLevel       ::Level                        ... Initial level of the atom
        + finalLevel         ::Level                        ... Final level of the atom
        + beam               ::Pulse.AbstractBeam           ... Beam character and properties of the incident light pulse.
        + envelope           ::Pulse.AbstractEnvelope       ... Envelope of the incident light pulse.
        + polarization       ::Basics.AbstractPolarization  ... Envelope of the incident light pulse.
        + volkov             ::AbstractVolkovState          ... Specify the treatment/approach of the Volkov states.
        + settings           ::StrongField.Settings         ... Settings for the controlling the SFA amplitudes.
    """
    struct  Computation
        observable           ::AbstractSFAObservable
        nuclearModel         ::Nuclear.Model
        grid                 ::Radial.Grid
        initialLevel         ::Level
        finalLevel           ::Level
        beam                 ::Pulse.AbstractBeam
        envelope             ::Pulse.AbstractEnvelope
        polarization         ::Basics.AbstractPolarization
        volkov               ::AbstractVolkovState
        settings             ::StrongField.Settings
    end

    """
    `StrongField.Computation()`  ... constructor for an `empty` instance of StrongField.Computation().
    """
    function Computation()
        Computation( SfaNoObservable(), Nuclear.Model(1.0), Radial.Grid(true), Level(), Level(),
                     Pulse.PlaneWaveBeam(), Pulse.InfiniteEnvelope(), Basics.RightCircular(), FreeVolkov(), Settings()  )
    end

    # `Base.string(comp::StrongField.Computation)`  ... provides a String notation for the variable comp::StrongField.Computation.
    function Base.string(comp::StrongField.Computation)
        sa = "Strong-field computation:  " * string(comp.observable) * " for Z = $(comp.nuclearModel.Z)  with " * string(comp.volkov) * "\n"
        sa = sa * " initial level with (J, P, energy) = ($(comp.initialLevel.J), $(comp.initialLevel.parity), $(comp.initialLevel.energy)) \n"
        sa = sa * " final level with   (J, P, energy) = ($(comp.finalLevel.J), $(comp.finalLevel.parity), $(comp.finalLevel.energy))   \n\n"
        sa = sa * "The incident laser pulse is described by the: \n "
        sa = sa * string(comp.beam) * "\n " * string(comp.polarization) * "\n " * string(comp.envelope)
        return( sa )
    end

    # `Base.show(io::IO, comp::StrongField.Computation)`  ... prepares a proper printout of the variable comp::StrongField.Computation.
    function Base.show(io::IO, comp::StrongField.Computation)
        sa = Base.string(comp);             print(io, sa, "\n\n")
        println(io, "Settings:              \n$(comp.settings)    ")
        println(io, "nuclearModel:          $(comp.nuclearModel)  ")
        println(io, "grid:                  $(comp.grid)  ")
    end

    """
    `struct  StrongField.SphericalAmplitude`
        ... to keep the amplitude at a given  energy-angular point (energy, theta, phi) in momentum space.

        + energy        ::Float64      ... Kinetic energy of the (outgoing) photoelectron.
        + theta         ::Float64      ... Polar angle of the (outgoing) photoelectron.
        + phi           ::Float64      ... Azimuthal angle of the (outgoing) photoelectron.
        + value         ::ComplexF64   ... (Total) Amplitude.
    """
    struct   SphericalAmplitude
        energy          ::Float64
        theta           ::Float64
        phi             ::Float64
        value           ::ComplexF64
    end

    function Base.string(amp::StrongField.SphericalAmplitude)
        sa = "Total SFA amplitude at (energy, theta, phi) = ($(amp.energy),$(amp.theta),$(amp.phi)) is $(amp.value)."
        return( sa )
    end

    function Base.show(io::IO, amp::StrongField.SphericalAmplitude)
        sa = string(amp);       print(io, sa)
    end

    """
    `struct  StrongField.OutcomeEnergyDistribution`
        ... to comprise the energy distribution of photoelectrons at given energies and angles, for instance,
            for later graphical representation.

        + theta         ::Float64             ... polar angle of the energy distribution
        + phi           ::Float64             ... azimuthal angle of the energy distribution
        + energies      ::Array{Float64,1}    ... selected energies of the distribution.
        + probabilities ::Array{Float64,1}    ... calculated probabilities of the energy distribution.
    """
    struct   OutcomeEnergyDistribution
        theta           ::Float64
        phi             ::Float64
        energies        ::Array{Float64,1}
        probabilities   ::Array{Float64,1}
    end

    # `Base.show(io::IO, outcome::StrongField.OutcomeEnergyDistribution)`  ... prepares a proper printout of the variable outcome::StrongField.OutcomeEnergyDistribution.
    function Base.show(io::IO, outcome::StrongField.OutcomeEnergyDistribution)
        println(io, "theta:             $(outcome.theta)  ")
        println(io, "phi:               $(outcome.phi)  ")
        println(io, "energies:          $(outcome.energies)  ")
        println(io, "probabilities:     $(outcome.probabilities)  ")
    end

    """
    `struct  StrongField.OutcomeMomentumDistribution`
        ... to comprise the momentum distribution of photoelectrons at given energies and angles, for instance,
            for later graphical representation.

        + theta         ::Float64             ... polar angle of the momentum distribution
        + momenta       ::Array{Float64}      ... pairs [px py] of momenta of the momentum distribution
        + probabilities ::Array{Float64,1}    ... calculated probabilities of the energy distribution.
    """
    struct   OutcomeMomentumDistribution
        theta           ::Float64
        momenta         ::Array{Float64}
        probabilities   ::Array{Float64,1}
    end

    # `Base.show(io::IO, outcome::StrongField.OutcomeMomentumDistribution)`  ... prepares a proper printout of the variable outcome::StrongField.OutcomeMomentumDistribution.
    function Base.show(io::IO, outcome::StrongField.OutcomeMomentumDistribution)
        println(io, "theta:             $(outcome.theta)  ")
        println(io, "momenta:           $(outcome.momenta)  ")
        println(io, "probabilities:     $(outcome.probabilities)  ")
    end

    ##################################################################################################################################################
    ##################################################################################################################################################
    ##################################################################################################################################################

    ####################################################################################################################
    ######The following four functions will be replaced later and serve only for comparison with previous results#######
    ####################################################################################################################

    """
    `StrongField.HydrogenPnl(epsiloni::Float64, n::Int, l::Int, rGrid::Array{Float64,1})`
        ... returns the non-relativistic hydrogen-like radial wave function with modified binding energy epsiloni
                at the radial grid points rGrid
    """
    function  HydrogenPnl(epsiloni::Float64, n::Int, l::Int, rGrid::Array{Float64,1})
        Pnl = Float64[]
        Z = n*sqrt(-2*epsiloni)

        for j = 1:length(rGrid)
            r = rGrid[j]
            p = r * sqrt( (2*Z/n)^3 * factorial(n-l-1)/(2*n*factorial(n+l)) ) * exp(-Z*r/n) * (2*Z*r/n)^l * GSL.sf_laguerre_n(n-l-1,2*l+1,2*Z*r/n)
            push!( Pnl, p )
        end

        return( Pnl )
    end

    """
    `StrongField.HydrogenDPnlDr(epsiloni::Float64, n::Int, l::Int, rGrid::Array{Float64,1})`
        ... returns the r-derivative of the non-relativistic hydrogen-like radial wave function with modified binding energy epsiloni
                at the radial grid points rGrid
    """
    function  HydrogenDPnlDr(epsiloni::Float64, n::Int, l::Int, rGrid::Array{Float64,1})
        DPnl = Float64[]
        Z = n*sqrt(-2*epsiloni)

        for j = 1:length(rGrid)
            r = rGrid[j]
            if  n-l-2 >= 0
                p = 1/n * 2^(l+1) * exp(-r*Z/n) * (r*Z/n)^l * sqrt( Z^3 * factorial(n-l-1) / (n^4*factorial(l+n)) ) * ( (n+l*n-r*Z) * GSL.sf_laguerre_n(n-l-1,2*l+1,2*r*Z/n) - 2*r*Z * GSL.sf_laguerre_n(n-l-2,2+2*l,2*r*Z/n) )
            elseif  n == 1 && l == 0
                p = -2 * exp(-r*Z) * sqrt(Z^3) * (r*Z - 1)
            else #not implemented
                p = 0
            end
            push!( DPnl, p )
        end

        return( DPnl )
    end

    """
    `StrongField.VolkovP(epsilonp::Float64, lp::Int, rGrid::Array{Float64,1})`
        ... returns the (plane-wave) Volkov radial wave function
                at the radial grid points rGrid
    """
    function  VolkovP(epsilonp::Float64, lp::Int, rGrid::Array{Float64,1})
        P = Float64[]

        for j = 1:length(rGrid)
            r = rGrid[j]
            p = r * GSL.sf_bessel_jl( lp, sqrt(2*epsilonp)*r )
            push!( P, p )
        end

        return( P )
    end

    """
    `StrongField.CoulombVolkovP(epsilonp::Float64, lp::Int, ZContinuum::Float64, rGrid::Array{Float64,1})`
        ... returns the Coulomb-Volkov radial wave function
                at the radial grid points rGrid
    """
    function  CoulombVolkovP(epsilonp::Float64, lp::Int, ZContinuum::Float64, rGrid::Array{Float64,1})
        sqrtTwoEps = sqrt(2*epsilonp)

        P = ComplexF64[]

        for j = 1:length(rGrid)
            r = rGrid[j]
            rho = sqrtTwoEps*r
            eta = -ZContinuum/sqrtTwoEps
            GammaValue = gamma( lp + 1 + eta * im )
            sigma = angle( GammaValue )
            Fl = 2.0^lp * exp(-0.5*pi*eta) * abs( GammaValue ) / factorial(2*lp+1) * rho^(lp+1) * exp( rho * im ) * HypergeometricFunctions.drummond1F1(lp+1+eta*im,2*lp+2,-2*rho*im)

            p = Fl * exp( sigma * im ) / sqrtTwoEps

            push!( P, p )
        end

        return( P )
    end

    """
    `StrongField.pReducedME(epsilonp::Float64, lp::Int, n::Int, l::Int, epsiloni::Float64, volkov::AbstractVolkovState)`
        ... computes the reduced matrix elements of the momentum operator <epsilonp lp ||p||n l> in the one-particle picture
    """
    function  pReducedME(epsilonp::Float64, lp::Int, n::Int, l::Int, epsiloni::Float64, volkov::AbstractVolkovState)
        rmax = 100.
        orderGL = 1000
        gaussLegendre = Radial.GridGL("Finite",0.0,rmax,orderGL)
        rgrid = gaussLegendre.t
        weights = gaussLegendre.wt

        Pnl = HydrogenPnl( epsiloni, n, l, rgrid )
        if  typeof(volkov) == FreeVolkov
            Pepsplp = VolkovP( epsilonp, lp, rgrid )
        elseif  typeof(volkov) == CoulombVolkov
            Pepsplp = CoulombVolkovP( epsilonp, lp, volkov.Z, rgrid )
        end

        DPnl = HydrogenDPnlDr( epsiloni, n, l, rgrid )

        integral = 0. * im

        #Sum over grid and compute Gauss-Legendre sum
        for    j = 1:orderGL
            r = rgrid[j]
            integrand = conj( Pepsplp[j] )/r * ( r*DPnl[j] - ((lp-l)*(lp+l+1))/2 * Pnl[j] )

            #Gauss-Legendre sum
            integral = integral + weights[j] * integrand
        end

        #Note that GSL.sf_coupling_3j takes takes the input (2*j1,2*j2,2*j3,2*m1,2*m2,2*m3)
        integral = integral * (-im)^(lp+1) * (-1)^lp * GSL.sf_coupling_3j( 2*lp, 2*1, 2*l, 0, 0, 0 )

        return( integral )
    end

    """
    `StrongField.scalarProdBoundCont(epsilonp::Float64, lp::Int, n::Int, l::Int, epsiloni::Float64, volkov::AbstractVolkovState)`
        ... computes the scalar product of the bound (hydrogenic) and continuum (Volkov) states in the one-particle picture
    """
    function  scalarProdBoundCont(epsilonp::Float64, lp::Int, n::Int, l::Int, m::Int, epsiloni::Float64, volkov::AbstractVolkovState)
        rmax = 1000.
        orderGL = 1000
        gaussLegendre = Radial.GridGL("Finite",0.0,rmax,orderGL)
        rgrid = gaussLegendre.t
        weights = gaussLegendre.wt

        Pnl = HydrogenPnl( epsiloni, n, l, rgrid )

        if  typeof(volkov) == FreeVolkov
            Pepsplp = VolkovP( epsilonp, lp, rgrid )
        elseif  typeof(volkov) == CoulombVolkov
            Pepsplp = CoulombVolkovP( epsilonp, lp, volkov.Z, rgrid )
        end

        DPnl = HydrogenDPnlDr( epsiloni, n, l, rgrid )

        integral = 0. * im

        #Sum over grid and compute Gauss-Legendre sum
        for    j = 1:orderGL
            r = rgrid[j]
            integrand = conj( Pepsplp[j] ) * Pnl[j]

            #Gauss-Legendre sum
            integral = integral + weights[j] * integrand
        end

        return((-im)^l * integral)
    end

    ####################################################################################################################

    """
    `StrongField.computeEnvelopeIntegrals(envelope::Pulse.AbstractEnvelope, beam::AbstractBeam, polarization::Basics.AbstractPolarization,
                                          thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)`
        ... computes the pulse-envelope integrals for the given laser pulse; this pulse is completely specified by its beam character
            and parameters, the pulse-envelope parameters and the polarization of the light. A tuple of the two integrals
            (fVolkov, fVolkovSquared)::Tuple{Complex{Float64},Complex{Float64}} is returned.
    """
    function computeEnvelopeIntegrals(envelope::Pulse.AbstractEnvelope, beam::AbstractBeam, polarization::Basics.AbstractPolarization,
                                      thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)
        fVolkovPlus    = Pulse.envelopeVolkovIntegral(true,  envelope, beam, polarization, thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)
        fVolkovMinus   = Pulse.envelopeVolkovIntegral(false, envelope, beam, polarization, thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)
        fVolkovSquared = Pulse.envelopeQuadVolkovIntegral(envelope, beam, polarization,    thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)

        #println("> Envelope integrals for an $(string(envelope)) with A0=$(beam.A0), omega=$(beam.omega) [a.u.], cep=$(beam.cep)" *
        #        " and for $(string(polarization)) light are:" )
        #println("    F^(Volkov) [+; omega; f^(env); A]   = $fVolkovPlus" )
        #println("    F^(Volkov) [-; omega; f^(env); A]   = $fVolkovMinus" )
        #println("    F^(quad, Volkov) [f^(env); A]       = $fVolkovSquared" )

        return( (fVolkovPlus, fVolkovMinus, fVolkovSquared) )
    end

    """
    `StrongField.computeOutcome(observable::StrongField.SfaEnergyDistribution, amplitudes::Array{SphericalAmplitude,1})`
        ... computes the requested photoelectron energy distribution and returns the results in the variable
            outcome::StrongField.OutcomeEnergyDistribution
    """
    function  computeOutcome(obs::StrongField.SfaEnergyDistribution, amplitudes::Array{SphericalAmplitude,1})
        probabilities = Float64[]
        for  amp  in  amplitudes   pp = sqrt(2*amp.energy);  push!(probabilities, pp * (amp.value * conj(amp.value)) )      end
        outcome = OutcomeEnergyDistribution(obs.theta, obs.phi, obs.energies, probabilities)
        return( outcome )
    end

    """
    `StrongField.computeOutcome(observable::StrongField.SfaMomentumDistribution, amplitudes::Array{SphericalAmplitude,1})`
        ... computes the requested photoelectron momentum distribution and returns the results in the variable
            outcome::StrongField.OutcomeMomentumDistribution
    """
    function  computeOutcome(obs::StrongField.SfaMomentumDistribution, amplitudes::Array{SphericalAmplitude,1})
        probabilities = Float64[]
        momenta       = Array{Float64}(undef,0,2)
        for  amp  in  amplitudes
            pp = sqrt(2*amp.energy)
            px = pp*sin(amp.theta)*cos(amp.phi)
            py = pp*sin(amp.theta)*sin(amp.phi)
            momenta = [momenta ; [px py]]
            push!( probabilities, pp * (amp.value * conj(amp.value)) )
        end
        outcome = OutcomeMomentumDistribution(obs.theta, momenta, probabilities)
        return( outcome )
    end

    """
    `StrongField.computeSphericalAmplitudes(comp::StrongField.Computation)`
        ... computes all necessary spherical amplitudes for the given initial and final levels, the Volkov states
            and polarization and by using the pulse envelope integrals. A newAmplitudes::Array{SphericalAmplitude,1} is returned.
    """
    function  computeSphericalAmplitudes(comp::StrongField.Computation)
        newAmplitudes = SphericalAmplitude[]
        #
        # First determine which spherical SFA amplitudes need to be computed before the actual computation starts
        sfaAmplitudes = StrongField.determineSphericalAmplitudes(comp.observable)
        # Determine quantum numbers of the initial and final state
        nqn = 1;    lqn = 0;    n = 1;      l = 0;      m = 0;    initialEn = convertUnits("energy: from eV to atomic", comp.initialLevel.energy)
        #

        lpMin = abs(l-1)
        lpMax = l+1

        # Compute the requested amplitudes
        for  amp in  sfaAmplitudes
            thetap        = amp.theta;   phip = amp.phi;     energyp = amp.energy
            envIntegrals  = StrongField.computeEnvelopeIntegrals(comp.envelope, comp.beam, comp.polarization, thetap, phip, energyp, initialEn)
            fVolkovPlus, fVolkovMinus, fVolkovSquared = envIntegrals
            #
            reducedMEArray = ComplexF64[]
            for lp = lpMin:lpMax
                push!( reducedMEArray, pReducedME(energyp, lp, n, l, initialEn, comp.volkov) )
            end
            #

            wminus = 0.0im;    wplus = 0.0im
            # Collect contributions from all l_p, q terms; this summation will change in a (nljm) representation
            for  lp = lpMin:lpMax
                reducedME = reducedMEArray[lp-lpMin+1]
                for  q in [-1, 0, 1]
                    wplus = wplus + AngularMomentum.sphericalYlm(lp, m-q, amp.theta, amp.phi) * (-1)^q  *
                             Basics.determinePolarizationVector(q, comp.polarization, star=false)         *
                             AngularMomentum.ClebschGordan(l, m, 1, -q, lp, m-q) * reducedME
                    wminus = wminus + AngularMomentum.sphericalYlm(lp, m+q, amp.theta, amp.phi)           *
                             Basics.determinePolarizationVector(q, comp.polarization, star=true)          *
                             AngularMomentum.ClebschGordan(l, m, 1, q, lp, m+q) * reducedME
                end
            end
            scalarProd = scalarProdBoundCont(energyp, l, n, l, m, initialEn, comp.volkov)
            wa = -im * sqrt(2/pi) * (fVolkovPlus * wminus + fVolkovMinus * wplus) - im / sqrt(2*pi) * fVolkovSquared * AngularMomentum.sphericalYlm(l, m, amp.theta, amp.phi) * scalarProd
            push!(newAmplitudes, SphericalAmplitude(amp.energy, amp.theta, amp.phi, wa))

            println(">> $(SphericalAmplitude(amp.energy, amp.theta, amp.phi, wa))")
        end

        return( newAmplitudes )
    end

    """
    `StrongField.determineSphericalAmplitudes(observable::StrongField.SfaEnergyDistribution)`
        ... determines which direct (and other) SFA amplitudes need to be computed; these amplitudes are not yet computed here but can be arranged
            so that only a minimum number of Volkov states and/or reduced many-electron matrix elements need to be computed.
            A list of amplitudes::Array{StrongField.SphericalAmplitude,1} is returned.
    """
    function  determineSphericalAmplitudes(observable::StrongField.SfaEnergyDistribution)
        amplitudes = StrongField.SphericalAmplitude[]
        for  energy in observable.energies      push!(amplitudes, SphericalAmplitude(energy, observable.theta, observable.phi, 0.))     end

        println("> A total of $(length(amplitudes)) spherical amplitudes need to be calculated.")
        return( amplitudes )
    end

    """
    `StrongField.determineSphericalAmplitudes(observable::StrongField.SfaMomentumDistribution)`
        ... determines which direct (and other) SFA amplitudes need to be computed; these amplitudes are not yet computed here but can be arranged
            so that only a minimum number of Volkov states and/or reduced many-electron matrix elements need to be computed.
            A list of amplitudes::Array{StrongField.SphericalAmplitude,1} is returned.
    """
    function  determineSphericalAmplitudes(observable::StrongField.SfaMomentumDistribution)
        amplitudes = StrongField.SphericalAmplitude[]
        for  energy in observable.energies
            for  phi in observable.phis
                push!(amplitudes, SphericalAmplitude(energy, observable.theta, phi, 0.))
            end
        end

        println("> A total of $(length(amplitudes)) spherical amplitudes need to be calculated.")
        return( amplitudes )
    end

    """
    `StrongField.plot( comp::StrongField.Computation, results::Dict{String,Any}, energyScale::String = "atomic", probabilityScaling::String = "linear" )`
        ... generates a graphical representation of the observable (either StrongField.SfaEnergyDistribution or StrongField.SfaMomentumDistribution).
        energyScale determines the scaling of the energy axis either in atomic units (energyScale = "atomic") or in units of hbar*omega (energyScale = "omega").
        The y-axis is scaled either linearly (probabilityScaling = "linear") or logarithmically (probabilityScaling = "log")
    """
    function plot( comp::StrongField.Computation, results::Dict{String,Any}, energyScale::String = "atomic", probabilityScaling::String = "linear" )

        #---Photoelectron energy spectrum---
        if  typeof(comp.observable) == StrongField.SfaEnergyDistribution
                #prepare data and rescale the x-axis if neccessary
                if  energyScale == "atomic"
                    energyLabel = "E (a.u.)"
                    energies = results["energy distribution"].energies
                elseif  energyScale == "omega"
                    energyLabel = "E/omega"
                    energies = results["energy distribution"].energies / comp.beam.omega
                end
                probabilities = results["energy distribution"].probabilities
                gr() #sets the plotting backend to the package "GR"

                #set scaling of the y-axis
                scaling = :identity
                if  probabilityScaling == "log"
                            scaling = :log10
                end

                #generate the plot
                p = Plots.plot(energies, probabilities,
                                title = "Photoelectron energy spectrum",
                                xlabel = energyLabel,
                                ylabel = "P(E)",
                                xscale = :identity,
                                yscale = scaling,
                                framestyle = :box,
                                legend = :none,
                                #markershape = :circle,
                                #markershape = :none,
                                line = 2,
                                linecolor = :black,
                                gridlinewidth = 2,
                                tickfontsize = 10,
                                labelfontsize = 10,
                                labelfontfamily = "Latin Modern Roman",
                                titlefontfamily = "Latin Modern Roman"

                              )

                #export the plot as png-file
                png("energy_spectrum")
        #-----------------------------------

        #---Photoelectron momentum distribution---
        elseif  typeof(comp.observable) == StrongField.SfaMomentumDistribution
                #prepare data
                pList = results["momentum distribution"].momenta
                probList = results["momentum distribution"].probabilities

                #generate the plot
                #pyplot()
                #r = range(0,stop=10,length=11)
                #theta = range(0,stop=360,length=361)
                #f(r,theta) = r^2
                #println("$(f.(r,theta'))")
                #Plots.plot( heatmap( f.(r,theta'), proj=:polar ) )
                #println("$(transpose(probList))")
                #Plots.plot( heatmap( pList, transpose(probList), proj=:polar ) )
                #imshow(pList,transpose(probList))
                #probList = probList[2:-1]
                #matplotlib.pyplot.pcolormesh(pList[:,1], pList[:,2], transpose(probList))
                contourf( pList[:,1], pList[:,2], probList )
                #matplotlib.pyplot.imshow(pList)
                #heatmap( pList[:,1], pList[:,2], probList, proj=:polar, legend=true
                                #heatmap( probList
                #                )

                #p = Plots.plot( #heatmap( pList[:,1], pList[:,2], probList
                #               Plots.GR.polarheatmap( probList
                            ##heatmap( probList
                #                 )
                #              )

                #export the plot as png-file
                png("momentum_distribution")
        #-----------------------------------------

        #---Not a valid obserable---
        else
                error("Undefined observable for strong-field computations.")
        end
        #---------------------------

    end

    """
    `StrongField.perform(comp::StrongField.Computation; output::Bool=false)`
        ... to perform a computation of (one) selected observable that is related to SFA, such as a photoelectron energy
            distribution, momentum distribution or the computation of sidebands.
    """
    function perform(comp::StrongField.Computation; output::Bool=false)
        if  output    results = Dict{String, Any}()    else    results = nothing    end
        nModel = comp.nuclearModel

        if       typeof(comp.observable) == StrongField.SfaEnergyDistribution
            sfaAmplitudes = StrongField.computeSphericalAmplitudes(comp)
            sfaOutcome    = StrongField.computeOutcome(comp.observable, sfaAmplitudes)
            if output    results = Base.merge( results, Dict("computation" => comp, "energy distribution" => sfaOutcome) )  end
        elseif   typeof(comp.observable) == StrongField.SfaMomentumDistribution
            sfaAmplitudes = StrongField.computeSphericalAmplitudes(comp)
            sfaOutcome    = StrongField.computeOutcome(comp.observable, sfaAmplitudes)
            if output    results = Base.merge( results, Dict("computation" => comp, "momentum distribution" => sfaOutcome) )  end
        else     error("Undefined observable for strong-field computations.")
        end

        if  output   return( results )   else    return( nothing )   end
    end

end # module

ext: jl


text:
using FMMTrees
using StaticArrays
using AdaptiveCrossApproximation

p = rand(SVector{2,Float64}, 100)
q, tree = AdaptiveCrossApproximation.clustertree(p)

using Plots
gr()

rectangle(ll, ur) = [ll[1],ur[1],ur[1],ll[1],ll[1]], [ll[2],ll[2],ur[2],ur[2],ll[2]]

plot()
level = 1
clrs = [:blue, :red, :green, :black, :cyan, :yellow]
state = [(0,tree[1].num_children)]
for (i,b) in enumerate(tree)
    ll, ur = FMMTrees.boundingbox(q[b.data.begin_idx:b.data.end_idx-1])
    x, y = rectangle(ll, ur)
    level == 5 && plot!(x,y,color=clrs[level])
    level == 4 && println(level, ": ", b.data.begin_idx:b.data.end_idx-1)
    state[end] = (state[end][1] + b.num_children + 1, state[end][2])
    if b.num_children != 0
         level += 1
         push!(state, (0,b.num_children))
     else
         while state[end][1] == state[end][2]
             level -= 1
             pop!(state)
         end
     end
     plot!()
end

x = [r[1] for r in q]
y = [r[2] for r in q]
scatter!(x,y)
plot!(axis=:none,legend=:none)
plot!(title="level 4 ACA clusters")

FMMTrees.print_tree(tree)


text:
using Test
using RNGTest
using Random

include("tests.jl")


text:
export QBasedPolicy, TabularRandomPolicy

using MacroTools: @forward
using Flux
using Setfield

"""
    QBasedPolicy(;learner::Q, explorer::S)

Use a Q-`learner` to generate estimations of action values.
Then an `explorer` is applied on the estimations to select an action.
"""
Base.@kwdef mutable struct QBasedPolicy{Q<:AbstractLearner,E<:AbstractExplorer} <:
                           AbstractPolicy
    learner::Q
    explorer::E
end

Flux.functor(x::QBasedPolicy) = (learner = x.learner,), y -> @set x.learner = y.learner

(π::QBasedPolicy)(env) = π(env, ActionStyle(env))
(π::QBasedPolicy)(env, ::MinimalActionSet) = get_actions(env)[π.explorer(π.learner(env))]
(π::QBasedPolicy)(env, ::FullActionSet) =
    get_actions(env)[π.explorer(π.learner(env), get_legal_actions_mask(env))]

RLBase.get_prob(p::QBasedPolicy, env) = get_prob(p, env, ActionStyle(env))
RLBase.get_prob(p::QBasedPolicy, env, ::MinimalActionSet) =
    get_prob(p.explorer, p.learner(env))
RLBase.get_prob(p::QBasedPolicy, env, ::FullActionSet) =
    get_prob(p.explorer, p.learner(env), get_legal_actions_mask(env))

@forward QBasedPolicy.learner RLBase.get_priority, RLBase.update!

function Flux.testmode!(p::QBasedPolicy, mode = true)
    testmode!(p.learner, mode)
    testmode!(p.explorer, mode)
end

#####
# TabularRandomPolicy
#####

const TabularRandomPolicy = QBasedPolicy{<:TabularLearner,<:WeightedExplorer}

function TabularRandomPolicy(;
    rng = Random.GLOBAL_RNG,
    is_normalized = true,
    table = Dict{String,Vector{Float64}}(),
)
    QBasedPolicy(;
        learner = TabularLearner(table),
        explorer = WeightedExplorer(; is_normalized = is_normalized, rng = rng),
    )
end

function (p::TabularRandomPolicy)(env::AbstractEnv)
    if ChanceStyle(env) === EXPLICIT_STOCHASTIC
        if get_current_player(env) == get_chance_player(env)
            # this should be faster. we don't need to allocate memory to store the probability of chance node
            return rand(p.explorer.rng, get_actions(env))
        end
    end
    p(env, ActionStyle(env))  # fall back to general implementation above
end

function RLBase.get_prob(p::TabularRandomPolicy, env, ::FullActionSet)
    m = get_legal_actions_mask(env)
    prob = zeros(length(m))
    prob[m] .= p.learner(env)
    prob
end


text:
# from Float64 to Posit8,16,32
Posit8(x::Float64) = ccall((:convertDoubleToP8, SoftPositPath), Posit8, (Float64,),x)
Posit16(x::Float64) = ccall((:convertDoubleToP16, SoftPositPath), Posit16, (Float64,),x)
Posit32(x::Float64) = ccall((:convertDoubleToP32, SoftPositPath), Posit32, (Float64,),x)

# from Float32/16 to Posit8,16,32
Posit8(x::T where {T <: Float16or32}) = Posit8(Float64(x))
Posit16(x::T where {T <: Float16or32}) = Posit16(Float64(x))
Posit32(x::T where {T <: Float16or32}) = Posit32(Float64(x))

# from Float64 to Posit_1 (convertDoubleToPX1 currently not available)
# Posit8_1(x::Float64) = ccall((:convertDoubleToPX1, SoftPositPath), Posit8_1, (Float64,Int64),x,8)
# Posit16_1(x::Float64) = ccall((:convertDoubleToPX1, SoftPositPath), Posit16_1, (Float64,Int64),x,16)
# Posit24_1(x::Float64) = ccall((:convertDoubleToPX1, SoftPositPath), Posit24_1, (Float64,Int64),x,24)

# use detour
Posit8_1(x::Float64) = Posit8_1(Posit32(x))
Posit16_1(x::Float64) = Posit16_1(Posit32(x))
Posit24_1(x::Float64) = Posit24_1(Posit32(x))

# from Float32/16 to Posit_1
Posit8_1(x::T where {T <: Float16or32}) = Posit8_1(Float64(x))
Posit16_1(x::T where {T <: Float16or32}) = Posit16_1(Float64(x))
Posit24_1(x::T where {T <: Float16or32}) = Posit24_1(Float64(x))

# from Float64 to Posit_2
Posit8_2(x::Float64) = ccall((:convertDoubleToPX2, SoftPositPath), Posit8_2, (Float64,Int64),x,8)
Posit16_2(x::Float64) = ccall((:convertDoubleToPX2, SoftPositPath), Posit16_2, (Float64,Int64),x,16)
Posit24_2(x::Float64) = ccall((:convertDoubleToPX2, SoftPositPath), Posit24_2, (Float64,Int64),x,24)

# from Float32/16 to Posit_2
Posit8_2(x::T where {T <: Float16or32}) = Posit8_2(Float64(x))
Posit16_2(x::T where {T <: Float16or32}) = Posit16_2(Float64(x))
Posit24_2(x::T where {T <: Float16or32}) = Posit24_2(Float64(x))


text:
using Statistics

# Quick and dirty Anderson implementation ... lacks important things like
# control of the condition number of the anderson matrix.
# Not particularly optimised. Also should be moved to NLSolve ...
#
# Accelerates the iterative solution of f(x) = 0 according to a
# damped preconditioned scheme
#    xₙ₊₁ = xₙ + αₙ P⁻¹ f(xₙ)
# Where f(x) computes the residual (e.g. SCF(x) - x)
# Further define
#    preconditioned residual  Pf(x) = P⁻¹ f(x)
#    fixed-point map          g(x)  = V + α Pf(x)
# where the α may vary between steps.
#
# Finds the linear combination xₙ₊₁ = g(xₙ) + ∑ᵢ βᵢ (g(xᵢ) - g(xₙ))
# such that |Pf(xₙ) + ∑ᵢ βᵢ (Pf(xᵢ) - Pf(xₙ))|² is minimal
struct AndersonAcceleration
    m::Int                  # maximal history size
    iterates::Vector{Any}   # xₙ
    residuals::Vector{Any}  # Pf(xₙ)
    maxcond::Real           # Maximal condition number for Anderson matrix
end
AndersonAcceleration(;m=10, maxcond=1e6) = AndersonAcceleration(m, [], [], maxcond)

function Base.push!(anderson::AndersonAcceleration, xₙ, αₙ, Pfxₙ)
    push!(anderson.iterates,  vec(xₙ))
    push!(anderson.residuals, vec(Pfxₙ))
    if length(anderson.iterates) > anderson.m
        popfirst!(anderson.iterates)
        popfirst!(anderson.residuals)
    end
    @assert length(anderson.iterates) <= anderson.m
    @assert length(anderson.iterates) == length(anderson.residuals)
    anderson
end

# Gets the current xₙ, Pf(xₙ) and damping αₙ
function (anderson::AndersonAcceleration)(xₙ, αₙ, Pfxₙ)
    xs   = anderson.iterates
    Pfxs = anderson.residuals

    # Special cases with fast exit
    anderson.m == 0 && return xₙ .+ αₙ .* Pfxₙ
    if isempty(xs)
        push!(anderson, xₙ, αₙ, Pfxₙ)
        return xₙ .+ αₙ .* Pfxₙ
    end

    M = hcat(Pfxs...) .- vec(Pfxₙ)  # Mᵢⱼ = (Pfxⱼ)ᵢ - (Pfxₙ)ᵢ
    # We need to solve 0 = M' Pfxₙ + M'M βs <=> βs = - (M'M)⁻¹ M' Pfxₙ

    # Ensure the condition number of M stays below maxcond, else prune the history
    Mfac = qr(M)
    while size(M, 2) > 1 && cond(Mfac.R) > anderson.maxcond
        M = M[:, 2:end]  # Drop oldest entry in history
        popfirst!(anderson.iterates)
        popfirst!(anderson.residuals)
        Mfac = qr(M)
    end

    xₙ₊₁ = vec(xₙ) .+ αₙ .* vec(Pfxₙ)
    βs   = -(Mfac \ vec(Pfxₙ))
    for (iβ, β) in enumerate(βs)
        xₙ₊₁ .+= β .* (xs[iβ] .- vec(xₙ) .+ αₙ .* (Pfxs[iβ] .- vec(Pfxₙ)))
    end

    push!(anderson, xₙ, αₙ, Pfxₙ)
    reshape(xₙ₊₁, size(xₙ))
end

function ScfAcceptStepAll()
    accept_step(info, info_next) = true
end

"""
Accept a step if the energy is at most increasing by `max_energy` and the residual
is at most `max_relative_residual` times the residual in the previous step.
"""
function ScfAcceptImprovingStep(;max_energy_change=1e-12, max_relative_residual=1.0)
    function accept_step(info, info_next)
        energy_change = info_next.energies.total - info.energies.total
        relative_residual = norm(info_next.Pinv_δV) / norm(info.Pinv_δV)

        # Accept if energy goes down or residual decreases
        accept = energy_change < max_energy_change || relative_residual < max_relative_residual
        mpi_master() && @debug "Step $(accept ? "accepted" : "discarded")" energy_change relative_residual
        accept
    end
end

"""
Use the two iteration states `info` and `info_next` to find a damping value
from a quadratic model for the SCF energy. Returns `nothing` if the constructed
model is not considered trustworthy, else returns the suggested damping.
"""
function scf_damping_quadratic_model(info, info_next; modeltol=0.1)
    T     = eltype(info.Vin)
    dvol  = info.basis.dvol

    Vin   = info.Vin
    ρin   = info.ρout      # = ρ(Vin)
    Vout  = info.Vout      # = step(Vin), where step(V) = (Vext + Vhxc(ρ(V)))
    α0    = info_next.α
    Vnext = info_next.Vin  # = Vin + α0 * (Anderson(Vin, P⁻¹( Vout - Vin )) - Vin)
    ρnext = info_next.ρout # = ρ(Vnext)
    δρ    = ρnext - ρin
    # α0 * δV = α0 * (Vnext - Vin) = α0 * (Anderson(Vin, P⁻¹( Vout - Vin )) - Vin)

    # We build a quadratic model for
    #   ϕ(α) = E(Vin  + α δV)  at α = 0
    #        = E(Vin) + α ∇E|_(V=Vin) ⋅ δV + ½ α^2 <δV | ∇²E|_(V=Vin) | δV>
    #        = E(Vin) + (α/α0) ∇E|_(V=Vin) ⋅ (α0 * δV) + ½ (α/α0)² <α0 * δV | ∇²E|_(V=Vin) | α0 * δV>
    #
    # Now
    #      ∇E|_(V=Vin)  = - χ₀(Vout - Vin)
    #      ∇²E|_(V=Vin) ≃ - χ₀ (1 - K χ₀)        (only true if Vin is an SCF minimum, K taken at Vin)
    # and therefore using the self-adjointness of χ₀
    #      ∇E|_(V=Vin) ⋅ δV         = -(Vout - Vin) ⋅ χ₀(δV) = - (Vout - Vin) ⋅ δρ
    #      <δV | ∇²E|_(V=Vin) | δV> = - δV ⋅ δρ + δρ ⋅ K(δρ)
    #
    slope = dot(Vout .- Vin, δρ) / α0 * dvol
    Kδρ   = apply_kernel(info.basis, δρ; ρ=ρin)
    curv  = dvol * (-dot(Vnext .- Vin, δρ) + dot(δρ, Kδρ)) / α0^2
    Emodel(α) = info.energies.total + slope * α + curv * α^2 / 2

    # Relative error of the model at α0 (the damping we used to get info_next)
    Etotal_next = info_next.energies.total

    # TODO Is this a good error measure ... for larger problems it seems to be over-demanding
    model_relerror = abs(Etotal_next - Emodel(α0)) / abs(Etotal_next - info.energies.total)

    minimum_exists = curv > eps(T)  # Does the model predict a minimum along the search direction
    trusted_model  = model_relerror < modeltol      # Model fits observation
    tight_model    = model_relerror < modeltol / 5  # Model fits observation very well

    # Accept model if it leads to minimum and is either tight or shows a negative slope
    α_model = -slope / curv
    if minimum_exists && (tight_model || (slope < -eps(T) && trusted_model))
        mpi_master() && @debug "Quadratic model accepted" model_relerror slope curv α_model
        (α=α_model, relerror=model_relerror)
    else
        mpi_master() && @debug "Quadratic model discarded" model_relerror slope curv α_model
        (α=nothing, relerror=model_relerror) # Model not trustworthy ...
    end
end

# Adaptive damping using a quadratic model
@kwdef struct AdaptiveDamping
    α_min = 0.05        # Minimal damping
    α_max = 1.0         # Maximal damping
    α_trial_init = 0.8  # Initial trial damping used (i.e. in the first SCF step)
    α_trial_min = 0.2   # Minimal trial damping used in a step
    α_trial_enhancement = 1.1  # Enhancement factor to α_trial in case a step is immediately successful
    modeltol = 0.1      # Maximum relative error on the predicted energy for model
    #                     to be considered trustworthy
end
function AdaptiveDamping(α_trial_min; kwargs...)
    # Select some reasonable defaults.
    # The free tweaking parameter here should be increased a bit for cases,
    # where the Anderson does weird stuff in case of too small damping.
    AdaptiveDamping(;α_min=α_trial_min / 4,
                     α_max=max(1.25α_trial_min, 1.0),
                     α_trial_init=max(α_trial_min, 0.8),
                     α_trial_min=α_trial_min,
                     kwargs...)
end

function ensure_damping_within_range(damping::AdaptiveDamping, α, α_next)
    α_sign = sign(α_next)
    abs(α_next) ≤ damping.α_min / 5 && (α_sign = +1.0)
    if α_sign > 0.0
        # Avoid getting stuck
        α_next = min(0.95abs(α), abs(α_next))
    else
        # Don't move too far backwards (where model validity cannot be ensured)
        α_next = min(0.50abs(α), abs(α_next))
    end
    α_next = clamp(α_next, damping.α_min, damping.α_max)
    α_sign * α_next
end

function propose_backtrack_damping(damping::AdaptiveDamping, info, info_next)
    if abs(info_next.α) < 1.75damping.α_min
        # Too close to α_min to be worth making another step ... just give up
        return info_next.α
    end

    α_next, relerror = scf_damping_quadratic_model(info, info_next; modeltol=damping.modeltol)
    if isnothing(α_next)
        # Model failed ... use heuristics: Half for small model error, else use a quarter
        α_next = info_next.α / (relerror < 10 ? 2 : 4)
    end
    ensure_damping_within_range(damping, info_next.α, α_next)
end

trial_damping(damping::AdaptiveDamping) = damping.α_trial_init
function trial_damping(damping::AdaptiveDamping, info, info_next, step_successful)
    n_backtrack = length(info_next.diagonalization)

    α_trial = abs(info_next.α)  # By default use the α that worked in this step
    if step_successful && n_backtrack == 1  # First step was good => speed things up
        α_trial ≥ damping.α_max && return damping.α_max  # No need to compute model
        α_model = scf_damping_quadratic_model(info, info_next; modeltol=damping.modeltol).α
        if !isnothing(α_model)  # Model is meaningful
            α_trial = max(damping.α_trial_enhancement * abs(α_model), α_trial)
        end
    end

    # Ensure returned damping is in valid range
    clamp(α_trial, damping.α_trial_min, damping.α_max)
end

struct FixedDamping
    α
end
FixedDamping() = FixedDamping(0.8)
propose_backtrack_damping(damping::FixedDamping) = damping.α
trial_damping(damping::FixedDamping, args...) = damping.α

# Notice: For adaptive damping to run smoothly, multiple defaults need to be changed.
#         See the function scf_potential_mixing_adaptive for that use case.
@timing function scf_potential_mixing(
    basis::PlaneWaveBasis;
    damping=FixedDamping(0.8),
    n_bands=default_n_bands(basis.model),
    ρ=guess_density(basis),
    V=nothing,
    ψ=nothing,
    tol=1e-6,
    maxiter=100,
    eigensolver=lobpcg_hyper,
    n_ep_extra=3,
    diag_miniter=1,
    determine_diagtol=ScfDiagtol(),
    mixing=SimpleMixing(),
    is_converged=ScfConvergenceEnergy(tol),
    callback=ScfDefaultCallback(),
    acceleration=AndersonAcceleration(;m=10),
    accept_step=ScfAcceptStepAll(),
    max_backtracks=3,  # Maximal number of backtracking line searches
)
    # TODO Test other mixings and lift this
    @assert (   mixing isa SimpleMixing
             || mixing isa KerkerMixing
             || mixing isa KerkerDosMixing)
    damping isa Number && (damping = FixedDamping(damping))

    if ψ !== nothing
        @assert length(ψ) == length(basis.kpoints)
        for ik in 1:length(basis.kpoints)
            @assert size(ψ[ik], 2) == n_bands + n_ep_extra
        end
    end

    # Initial guess for V (if none given)
    energies, ham = energy_hamiltonian(basis, nothing, nothing; ρ=ρ)
    isnothing(V) && (V = total_local_potential(ham))

    function EVρ(Vin; diagtol=tol / 10, ψ=nothing)
        ham_V = hamiltonian_with_total_potential(ham, Vin)
        res_V = next_density(ham_V; n_bands=n_bands, ψ=ψ, n_ep_extra=n_ep_extra,
                             miniter=diag_miniter, tol=diagtol, eigensolver=eigensolver)
        new_E, new_ham = energy_hamiltonian(basis, res_V.ψ, res_V.occupation;
                                            ρ=res_V.ρout, eigenvalues=res_V.eigenvalues,
                                            εF=res_V.εF)
        (basis=basis, ham=new_ham, energies=new_E, Vin=Vin,
         Vout=total_local_potential(new_ham), res_V...)
    end

    n_iter    = 1
    converged = false
    α_trial   = trial_damping(damping)
    diagtol   = determine_diagtol((ρin=ρ, Vin=V, n_iter=n_iter))
    info      = EVρ(V; diagtol=diagtol, ψ=ψ)
    Pinv_δV   = mix_potential(mixing, basis, info.Vout - info.Vin; n_iter, info...)
    info      = merge(info, (α=NaN, diagonalization=[info.diagonalization], ρin=ρ,
                             n_iter=n_iter, Pinv_δV=Pinv_δV))
    ΔEdown    = 0.0

    while n_iter < maxiter
        info = merge(info, (stage=:iterate, algorithm="SCF", converged=converged))
        callback(info)
        if MPI.bcast(is_converged(info), 0, MPI.COMM_WORLD)
            # TODO Debug why these MPI broadcasts are needed
            converged = true
            break
        end
        n_iter += 1
        info = merge(info, (n_iter=n_iter, ))

        # Ensure same α on all processors
        α_trial = MPI.bcast(α_trial, 0, MPI.COMM_WORLD)
        δV = (acceleration(info.Vin, α_trial, info.Pinv_δV) - info.Vin) / α_trial

        # Determine damping and take next step
        guess   = ψ
        α       = α_trial
        successful  = false  # Successful line search (final step is considered good)
        n_backtrack = 1
        diagonalization = empty(info.diagonalization)
        info_next = info
        while n_backtrack ≤ max_backtracks
            diagtol = determine_diagtol(info_next)
            mpi_master() && @debug "Iteration $n_iter linesearch step $n_backtrack   α=$α diagtol=$diagtol"
            Vnext = info.Vin .+ α .* δV

            info_next    = EVρ(Vnext; ψ=guess, diagtol=diagtol)
            Pinv_δV_next = mix_potential(mixing, basis, info_next.Vout - info_next.Vin;
                                         n_iter, info_next...)
            push!(diagonalization, info_next.diagonalization)
            info_next = merge(info_next, (α=α, diagonalization=diagonalization,
                                          ρin=info.ρout, n_iter=n_iter,
                                          Pinv_δV=Pinv_δV_next))

            successful = accept_step(info, info_next)
            successful = MPI.bcast(successful, 0, MPI.COMM_WORLD)  # Ensure same successful
            if successful || n_backtrack ≥ max_backtracks
                break
            end
            n_backtrack += 1

            # Adjust α to try again ...
            α_next = propose_backtrack_damping(damping, info, info_next)
            α_next = MPI.bcast(α_next, 0, MPI.COMM_WORLD)  # Ensure same α on all processors
            if α_next == α  # Backtracking further not useful ...
                break
            end

            # Adjust to guess fitting α best:
            guess = α_next > α / 2 ? info_next.ψ : ψ
            α = α_next
        end

        # Switch off acceleration in case of very bad steps
        ΔE = info_next.energies.total - info.energies.total
        ΔE < 0 && (ΔEdown = -max(abs(ΔE), tol))

        # Update α_trial and commit the next state
        α_trial = trial_damping(damping, info, info_next, successful)
        info = info_next
    end

    ham  = hamiltonian_with_total_potential(ham, info.Vout)
    info = (ham=ham, basis=basis, energies=info.energies, converged=converged,
            ρ=info.ρout, eigenvalues=info.eigenvalues, occupation=info.occupation,
            εF=info.εF, n_iter=n_iter, n_ep_extra=n_ep_extra, ψ=info.ψ,
            diagonalization=info.diagonalization, stage=:finalize, algorithm="SCF")
    callback(info)
    info
end

# Wrapper function setting a few good defaults for adaptive damping
function scf_potential_mixing_adaptive(basis; tol=1e-6, damping=AdaptiveDamping(), kwargs...)
    @assert damping isa AdaptiveDamping
    scf_potential_mixing(basis; tol=tol, diag_miniter=2,
                         accept_step=ScfAcceptImprovingStep(max_energy_change=tol),
                         determine_diagtol=ScfDiagtol(ratio_ρdiff=0.03, diagtol_max=5e-3),
                         damping=damping, kwargs...)
end


text:
# Unless you have a reason not to precompile, you should set your package to precompile.

# Reasons to not precompile:
# 1) You have global pointers. This includes any globally defined Bigs! These
#    will cause segfaults.
# 2) You are calling `@eval` in a function that modifies your modules' global scope
#    This may have bad interactions with precompilation.
# 3) You are importing/using a package which has precompilation disabled

__precompile__()

# If you are not going to precompile your package, it's good practice to disable it:
# __precompile__(false)

module ExamplePackage

# The main module file is for outlining the structure of the package.
# At the top of your package module, you should import your dependencies:

using Compat

# Then define your abstract types:

abstract AbstractMyType
abstract AbstractMyType2 <: AbstractMyType

# Some people make sure all export statements are at the top of the main module as well
# Only values which are exported enter the namespace, anything else must be
# qualified by the package name (ExamplePackage.not_exported). By doing this all
# at the top file, it is easy to know what functionality is exported by a package

# Note: You should try to keep exports to a minimum. Exports are the public API:
# internal functionality should not be exported!

export output_string, added_function

# Now include the real code
# Write your code in other files, otherwise the package outline gets muddled

include("output_string.jl")
include("my_new_file.jl")

end # module


text:
### Imports ###

import Base: size, show, getindex, setindex!, merge, merge!

### Exports ###

export AbstractContainer, AbstractNode
export size, show, relocate!, register!, nupdate!, eforward!, reinit!
export inputs, outputs, getindex, setindex!
export params
export merge, merge!


text:
# doubly periodic dipole
# https://dx.doi.org/10.1103/PhysRevLett.112.145301
# should be eigenstate of Hamiltonian evolution

using Test, Plots, VortexDistributions

# Methods
H(x) = x > 0. ? 1.0 : 0.0
T(x,xi) = x - xi
tans(x,xk) = tan((T(x,xk) - π)*0.5)
tanhs(x,xk,j) = tanh((T(x,xk) + 2*π*j)*0.5)

function Krow(x,y,xp,yp,xn,yn,j)
    return atan(tanhs(y,yn,j)*tans(x,xn)) -
    atan(tanhs(y,yp,j)*tans(x,xp))
end

# Dimensionless form
function θd(x,y,dip)
    vp,vn = dip
    xp,yp,_ = rawData(vp)
    xn,yn,_ = rawData(vn)
    s = 0.0
    for j = -5:5
        s += Krow(x,y,xp,yp,xn,yn,j)
    end
    return s + π*(H(T(x,xp)) - H(T(x,xn))) - y*(xp - xn)/(2*π)
end

# Now for arbitrary domains and dipole sizes:
function thetad(x,y,xp,yp,xn,yn)
    s = 0.0
    for j = -5:5
        s += Krow(x,y,xp,yp,xn,yn,j)
    end
    s += π*(H(T(x,xp)) - H(T(x,xn))) - y*(xp - xn)/(2*π)
    return s - x*H(abs(yp - yn) - π) + y*H(abs(xp - xn) - π)
end

function Thetad(x,y,xp,yp,xn,yn)
    Lx = x[end]-x[1]
    Ly = y[end]-y[1]
    return @. angle(exp(im*thetad.(x*2*pi/Lx,y'*2*pi/Ly,xp*2*pi/Lx,yp*2*pi/Ly,xn*2*pi/Lx,yn*2*pi/Ly)))
end

# Sukla's parameters
g=0.1
μ=15.
n0 = μ/g
γ=0.01
x = LinRange(-50,50,1024)
y = LinRange(-50,50,1024)

# dipole
xp = 0.
yp = 5

xn = 0.
yn = -5

vp = PointVortex(xp,yp,1)
vn = PointVortex(xn,yn,-1)
ξ = sqrt(1/μ)
dip = ScalarVortex(ξ,[vp;vn])

phase_periodic = Thetad(x,y,xp,yp,xn,yn)
heatmap(x,y,phase_periodic,transpose=true)

# make dipole wavefunction
psi = Torus(sqrt(n0)*one.(x*y'),x,y)
vortex!(psi,dip)

psip = @. abs(psi.ψ)*exp(im*phase_periodic)

heatmap(x,y,abs2.(psip),transpose=true)
heatmap(x,y,angle.(psip),transpose=true)

using FourierGPE

N = (1024,1024)
L = (100,100)
sim = Sim(L,N)
@unpack X,K = sim

phi = kspace(psip,sim) |> fftshift

DX,DK = dfftall(X,K)
kx,ky = K .|> fftshift

data = @. log(abs2(phi)+eps())
heatmap(kx,ky,data)

# define polar coordinates
Nkx = length(kx)
Nk = Nkx/2 |> Int
Nθk = 2*Nkx
k = LinRange(0,last(kx),Nk)'
θk = LinRange(0,2*pi,Nθk+1)[1:Nθk]

N = 100
ω = .3
plot(kx,hermite.(kx,N,ω))
hx,hy = init_polar(x,N,ω)
phiP = polar(psi.ψ,kx,ω,hx,hy)

densFP = abs2.(phiP)
heatmap(θk,k',densFP,transpose=true)
xlabel!(L"\theta_k");ylabel!(L"k")

dk = k[2] - k[1]
dθk = θk[2] - θk[1]

densk = sum(densFP;dims=1)'*dθk

plot(k',densk)


text:
### From array.jl

@ngenerate N Nothing function checksize(A::AbstractArray, I::NTuple{N, Any}...)
    @nexprs N d->(size(A, d) == length(I_d) || throw(DimensionMismatch("index $d has length $(length(I_d)), but size(A, $d) = $(size(A,d))")))
    nothing
end

unsafe_getindex(v::Real, ind::Int) = v
unsafe_getindex(v::Range, ind::Int) = first(v) + (ind-1)*step(v)
unsafe_getindex(v::BitArray, ind::Int) = Base.unsafe_bitgetindex(v.chunks, ind)
unsafe_getindex(v::AbstractArray, ind::Int) = v[ind]
unsafe_getindex(v, ind::Real) = unsafe_getindex(v, to_index(ind))

unsafe_setindex!{T}(v::AbstractArray{T}, x::T, ind::Int) = (v[ind] = x; v)
unsafe_setindex!(v::BitArray, x::Bool, ind::Int) = (Base.unsafe_bitsetindex!(v.chunks, x, ind); v)
unsafe_setindex!{T}(v::AbstractArray{T}, x::T, ind::Real) = unsafe_setindex!(v, x, to_index(ind))

# Version that uses cartesian indexing for src
@ngenerate N typeof(dest) function _getindex!(dest::Array, src::AbstractArray, I::NTuple{N,Union(Int,AbstractVector)}...)
    checksize(dest, I...)
    k = 1
    @nloops N i dest d->(@inbounds j_d = unsafe_getindex(I_d, i_d)) begin
        @inbounds dest[k] = (@nref N src j)
        k += 1
    end
    dest
end

# Version that uses linear indexing for src
@ngenerate N typeof(dest) function _getindex!(dest::Array, src::Array, I::NTuple{N,Union(Int,AbstractVector)}...)
    checksize(dest, I...)
    stride_1 = 1
    @nexprs N d->(stride_{d+1} = stride_d*size(src,d))
    @nexprs N d->(offset_d = 1)  # only really need offset_$N = 1
    k = 1
    @nloops N i dest d->(@inbounds offset_{d-1} = offset_d + (unsafe_getindex(I_d, i_d)-1)*stride_d) begin
        @inbounds dest[k] = src[offset_0]
        k += 1
    end
    dest
end

# It's most efficient to call checkbounds first, then to_index, and finally
# allocate the output. Hence the different variants.
_getindex(A, I::(Union(Int,AbstractVector)...)) =
    _getindex!(similar(A, index_shape(I...)), A, I...)

@nsplat N function getindex(A::Array, I::NTuple{N,Union(Real,AbstractVector)}...)
    checkbounds(A, I...)
    _getindex(A, to_index(I...))
end

# Also a safe version of getindex!
@nsplat N function getindex!(dest, src, I::NTuple{N,Union(Real,AbstractVector)}...)
    checkbounds(src, I...)
    _getindex!(dest, src, to_index(I...)...)
end

@ngenerate N typeof(A) function setindex!(A::Array, x, J::NTuple{N,Union(Real,AbstractArray)}...)
    @ncall N checkbounds A J
    @nexprs N d->(I_d = to_index(J_d))
    stride_1 = 1
    @nexprs N d->(stride_{d+1} = stride_d*size(A,d))
    @nexprs N d->(offset_d = 1)  # really only need offset_$N = 1
    if !isa(x, AbstractArray)
        @nloops N i d->(1:length(I_d)) d->(@inbounds offset_{d-1} = offset_d + (unsafe_getindex(I_d, i_d)-1)*stride_d) begin
            @inbounds A[offset_0] = x
        end
    else
        X = x
        @ncall N setindex_shape_check X I
        # TODO? A variant that can use cartesian indexing for RHS
        k = 1
        @nloops N i d->(1:length(I_d)) d->(@inbounds offset_{d-1} = offset_d + (unsafe_getindex(I_d, i_d)-1)*stride_d) begin
            @inbounds A[offset_0] = X[k]
            k += 1
        end
    end
    A
end

@ngenerate N NTuple{N,Vector{Int}} function findn{T,N}(A::AbstractArray{T,N})
    nnzA = countnz(A)
    @nexprs N d->(I_d = Array(Int, nnzA))
    k = 1
    @nloops N i A begin
        @inbounds if (@nref N A i) != zero(T)
            @nexprs N d->(I_d[k] = i_d)
            k += 1
        end
    end
    @ntuple N I
end

### subarray.jl

# Here we want to skip creating the dict-based cached version,
# so use the ngenerate function
function gen_getindex_body(N::Int)
    quote
        strd_1 = 1
        @nexprs $N d->(@inbounds strd_{d+1} = strd_d*s.dims[d])
        ind -= 1
        indp = s.first_index
        @nexprs $N d->begin
            i = div(ind, strd_{$N-d+1})
            @inbounds indp += i*s.strides[$N-d+1]
            ind -= i*strd_{$N-d+1}
        end
        s.parent[indp]
    end
end

eval(ngenerate(:N, nothing, :(getindex{T}(s::SubArray{T,N}, ind::Integer)), gen_getindex_body, 2:5, false))

function gen_setindex!_body(N::Int)
    quote
        strd_1 = 1
        @nexprs $N d->(@inbounds strd_{d+1} = strd_d*s.dims[d])
        ind -= 1
        indp = s.first_index
        @nexprs $N d->begin
            i = div(ind, strd_{$N-d+1})
            @inbounds indp += i*s.strides[$N-d+1]
            ind -= i*strd_{$N-d+1}
        end
        s.parent[indp] = v
    end
end

eval(ngenerate(:N, nothing, :(setindex!{T}(s::SubArray{T,N}, v, ind::Integer)), gen_setindex!_body, 2:5, false))

 cumsum(A::AbstractArray, axis::Integer) =  cumsum!(similar(A, Base._cumsum_type(A)), A, axis)
cumprod(A::AbstractArray, axis::Integer) = cumprod!(similar(A), A, axis)

for (f, op) in ((:cumsum!, :+),
                (:cumprod!, :*))
    @eval begin
        @ngenerate N typeof(B) function ($f){T,N}(B, A::AbstractArray{T,N}, axis::Integer)
            if size(B, axis) < 1
                return B
            end
            size(B) == size(A) || throw(DimensionMismatch("Size of B must match A"))
            if axis == 1
                # We can accumulate to a temporary variable, which allows register usage and will be slightly faster
                @inbounds @nloops N i d->(d > 1 ? (1:size(A,d)) : (1:1)) begin
                    tmp = convert(eltype(B), @nref(N, A, i))
                    @nref(N, B, i) = tmp
                    for i_1 = 2:size(A,1)
                        tmp = ($op)(tmp, @nref(N, A, i))
                        @nref(N, B, i) = tmp
                    end
                end
            else
                @nexprs N d->(isaxis_d = axis == d)
                # Copy the initial element in each 1d vector along dimension `axis`
                @inbounds @nloops N i d->(d == axis ? (1:1) : (1:size(A,d))) @nref(N, B, i) = @nref(N, A, i)
                # Accumulate
                @inbounds @nloops N i d->((1+isaxis_d):size(A, d)) d->(j_d = i_d - isaxis_d) begin
                    @nref(N, B, i) = ($op)(@nref(N, B, j), @nref(N, A, i))
                end
            end
            B
        end
    end
end

### from abstractarray.jl

@ngenerate N typeof(A) function fill!{T,N}(A::AbstractArray{T,N}, x)
    @nloops N i A begin
        @inbounds (@nref N A i) = x
    end
    A
end

@ngenerate N typeof(dest) function copy!{T,N}(dest::AbstractArray{T,N}, src::AbstractArray{T,N})
    if @nall N d->(size(dest,d) == size(src,d))
        @nloops N i dest begin
            @inbounds (@nref N dest i) = (@nref N src i)
        end
    else
        invoke(copy!, (typeof(dest), Any), dest, src)
    end
    dest
end

### BitArrays

## getindex

# general scalar indexing with two or more indices
# (uses linear indexing, which - in the safe version - performs the final
# bounds check and is defined in bitarray.jl)
# (code is duplicated for safe and unsafe versions for performance reasons)

@ngenerate N Bool function unsafe_getindex(B::BitArray, I_0::Int, I::NTuple{N,Int}...)
    stride = 1
    index = I_0
    @nexprs N d->begin
        stride *= size(B,d)
        index += (I_d - 1) * stride
    end
    return unsafe_getindex(B, index)
end

@ngenerate N Bool function getindex(B::BitArray, I_0::Int, I::NTuple{N,Int}...)
    stride = 1
    index = I_0
    @nexprs N d->begin
        l = size(B,d)
        stride *= l
        1 <= I_{d-1} <= l || throw(BoundsError())
        index += (I_d - 1) * stride
    end
    return B[index]
end

# contiguous multidimensional indexing: if the first dimension is a range,
# we can get some performance from using copy_chunks!

function unsafe_getindex(B::BitArray, I0::UnitRange{Int})
    X = BitArray(length(I0))
    copy_chunks!(X.chunks, 1, B.chunks, first(I0), length(I0))
    return X
end

function getindex(B::BitArray, I0::UnitRange{Int})
    checkbounds(B, I0)
    return unsafe_getindex(B, I0)
end

getindex{T<:Real}(B::BitArray, I0::UnitRange{T}) = getindex(B, to_index(I0))

@ngenerate N BitArray{length(index_shape(I0, I...))} function unsafe_getindex(B::BitArray, I0::UnitRange{Int}, I::NTuple{N,Union(Int,UnitRange{Int})}...)
    X = BitArray(index_shape(I0, I...))

    f0 = first(I0)
    l0 = length(I0)

    gap_lst_1 = 0
    @nexprs N d->(gap_lst_{d+1} = length(I_d))
    stride = 1
    ind = f0
    @nexprs N d->begin
        stride *= size(B, d)
        stride_lst_d = stride
        ind += stride * (first(I_d) - 1)
        gap_lst_{d+1} *= stride
    end

    storeind = 1
    @nloops(N, i, d->I_d,
        d->nothing, # PRE
        d->(ind += stride_lst_d - gap_lst_d), # POST
        begin # BODY
            copy_chunks!(X.chunks, storeind, B.chunks, ind, l0)
            storeind += l0
        end)
    return X
end

# general multidimensional non-scalar indexing

@ngenerate N BitArray{length(index_shape(I...))} function unsafe_getindex(B::BitArray, I::NTuple{N,Union(Int,AbstractVector{Int})}...)
    X = BitArray(index_shape(I...))
    Xc = X.chunks

    ind = 1
    @nloops N i d->I_d begin
        unsafe_bitsetindex!(Xc, (@ncall N unsafe_getindex B i), ind)
        ind += 1
    end
    return X
end

# general version with Real (or logical) indexing which dispatches on the appropriate method

@ngenerate N BitArray{length(index_shape(I...))} function getindex(B::BitArray, I::NTuple{N,Union(Real,AbstractVector)}...)
    checkbounds(B, I...)
    return unsafe_getindex(B, to_index(I...)...)
end

## setindex!

# general scalar indexing with two or more indices
# (uses linear indexing, which - in the safe version - performs the final
# bounds check and is defined in bitarray.jl)
# (code is duplicated for safe and unsafe versions for performance reasons)

@ngenerate N typeof(B) function unsafe_setindex!(B::BitArray, x::Bool, I_0::Int, I::NTuple{N,Int}...)
    stride = 1
    index = I_0
    @nexprs N d->begin
        stride *= size(B,d)
        index += (I_d - 1) * stride
    end
    unsafe_setindex!(B, x, index)
    return B
end

@ngenerate N typeof(B) function setindex!(B::BitArray, x::Bool, I_0::Int, I::NTuple{N,Int}...)
    stride = 1
    index = I_0
    @nexprs N d->begin
        l = size(B,d)
        stride *= l
        1 <= I_{d-1} <= l || throw(BoundsError())
        index += (I_d - 1) * stride
    end
    B[index] = x
    return B
end

# contiguous multidimensional indexing: if the first dimension is a range,
# we can get some performance from using copy_chunks!

function unsafe_setindex!(B::BitArray, X::BitArray, I0::UnitRange{Int})
    l0 = length(I0)
    l0 == 0 && return B
    f0 = first(I0)
    copy_chunks!(B.chunks, f0, X.chunks, 1, l0)
    return B
end

function unsafe_setindex!(B::BitArray, x::Bool, I0::UnitRange{Int})
    l0 = length(I0)
    l0 == 0 && return B
    f0 = first(I0)
    fill_chunks!(B.chunks, x, f0, l0)
    return B
end

@ngenerate N typeof(B) function unsafe_setindex!(B::BitArray, X::BitArray, I0::UnitRange{Int}, I::NTuple{N,Union(Int,UnitRange{Int})}...)
    length(X) == 0 && return B
    f0 = first(I0)
    l0 = length(I0)

    gap_lst_1 = 0
    @nexprs N d->(gap_lst_{d+1} = length(I_d))
    stride = 1
    ind = f0
    @nexprs N d->begin
        stride *= size(B, d)
        stride_lst_d = stride
        ind += stride * (first(I_d) - 1)
        gap_lst_{d+1} *= stride
    end

    refind = 1
    @nloops(N, i, d->I_d,
        d->nothing, # PRE
        d->(ind += stride_lst_d - gap_lst_d), # POST
        begin # BODY
            copy_chunks!(B.chunks, ind, X.chunks, refind, l0)
            refind += l0
        end)

    return B
end

@ngenerate N typeof(B) function unsafe_setindex!(B::BitArray, x::Bool, I0::UnitRange{Int}, I::NTuple{N,Union(Int,UnitRange{Int})}...)
    f0 = first(I0)
    l0 = length(I0)
    l0 == 0 && return B
    @nexprs N d->(length(I_d) == 0 && return B)

    gap_lst_1 = 0
    @nexprs N d->(gap_lst_{d+1} = length(I_d))
    stride = 1
    ind = f0
    @nexprs N d->begin
        stride *= size(B, d)
        stride_lst_d = stride
        ind += stride * (first(I_d) - 1)
        gap_lst_{d+1} *= stride
    end

    @nloops(N, i, d->I_d,
        d->nothing, # PRE
        d->(ind += stride_lst_d - gap_lst_d), # POST
        begin # BODY
            fill_chunks!(B.chunks, x, ind, l0)
        end)

    return B
end

# general multidimensional non-scalar indexing

@ngenerate N typeof(B) function unsafe_setindex!(B::BitArray, X::AbstractArray, I::NTuple{N,Union(Int,AbstractArray{Int})}...)
    refind = 1
    @nloops N i d->I_d @inbounds begin
        @ncall N unsafe_setindex! B convert(Bool,X[refind]) i
        refind += 1
    end
    return B
end

@ngenerate N typeof(B) function unsafe_setindex!(B::BitArray, x::Bool, I::NTuple{N,Union(Int,AbstractArray{Int})}...)
    @nloops N i d->I_d begin
        @ncall N unsafe_setindex! B x i
    end
    return B
end

# general versions with Real (or logical) indexing which dispatch on the appropriate method

# this one is for disambiguation only
function setindex!(B::BitArray, x, i::Real)
    checkbounds(B, i)
    return unsafe_setindex!(B, convert(Bool,x), to_index(i))
end

@ngenerate N typeof(B) function setindex!(B::BitArray, x, I::NTuple{N,Union(Real,AbstractArray)}...)
    checkbounds(B, I...)
    #return unsafe_setindex!(B, convert(Bool,x), to_index(I...)...) # segfaults! (???)
    @nexprs N d->(J_d = to_index(I_d))
    return @ncall N unsafe_setindex! B convert(Bool,x) J
end

# this one is for disambiguation only
function setindex!(B::BitArray, X::AbstractArray, i::Real)
    checkbounds(B, i)
    j = to_index(i)
    setindex_shape_check(X, j)
    return unsafe_setindex!(B, X, j)
end

@ngenerate N typeof(B) function setindex!(B::BitArray, X::AbstractArray, I::NTuple{N,Union(Real,AbstractArray)}...)
    checkbounds(B, I...)
    @nexprs N d->(J_d = to_index(I_d))
    @ncall N setindex_shape_check X J
    return @ncall N unsafe_setindex! B X J
end

## findn

@ngenerate N NTuple{N,Vector{Int}} function findn{N}(B::BitArray{N})
    nnzB = countnz(B)
    I = ntuple(N, x->Array(Int, nnzB))
    if nnzB > 0
        count = 1
        @nloops N i B begin
            if (@nref N B i) # TODO: should avoid bounds checking
                @nexprs N d->(I[d][count] = i_d)
                count += 1
            end
        end
    end
    return I
end

## isassigned

@ngenerate N Bool function isassigned(B::BitArray, I_0::Int, I::NTuple{N,Int}...)
    stride = 1
    index = I_0
    @nexprs N d->begin
        l = size(B,d)
        stride *= l
        1 <= I_{d-1} <= l || return false
        index += (I_d - 1) * stride
    end
    return isassigned(B, index)
end

## permutedims

for (V, PT, BT) in {((:N,), BitArray, BitArray), ((:T,:N), Array, StridedArray)}
    @eval @ngenerate N typeof(P) function permutedims!{$(V...)}(P::$PT{$(V...)}, B::$BT{$(V...)}, perm)
        dimsB = size(B)
        length(perm) == N || error("expected permutation of size $N, but length(perm)=$(length(perm))")
        isperm(perm) || error("input is not a permutation")
        dimsP = size(P)
        for i = 1:length(perm)
            dimsP[i] == dimsB[perm[i]] || throw(DimensionMismatch("destination tensor of incorrect size"))
        end

        #calculates all the strides
        strides_1 = 0
        @nexprs N d->(strides_{d+1} = stride(B, perm[d]))

        #Creates offset, because indexing starts at 1
        offset = 1 - sum(@ntuple N d->strides_{d+1})

        if isa(B, SubArray)
            offset += B.first_index - 1
            B = B.parent
        end

        ind = 1
        @nexprs 1 d->(counts_{N+1} = strides_{N+1}) # a trick to set counts_($N+1)
        @nloops(N, i, P,
            d->(counts_d = strides_d), # PRE
            d->(counts_{d+1} += strides_{d+1}), # POST
            begin # BODY
                sumc = sum(@ntuple N d->counts_{d+1})
                @inbounds P[ind] = B[sumc+offset]
                ind += 1
            end)

        return P
    end
end

## unique across dim

# TODO: this doesn't fit into the new hashing scheme in any obvious way

immutable Prehashed
    hash::Uint
end
hash(x::Prehashed) = x.hash

@ngenerate N typeof(A) function unique{T,N}(A::AbstractArray{T,N}, dim::Int)
    1 <= dim <= N || return copy(A)
    hashes = zeros(Uint, size(A, dim))

    # Compute hash for each row
    k = 0
    @nloops N i A d->(if d == dim; k = i_d; end) begin
       @inbounds hashes[k] = hash(hashes[k], hash((@nref N A i)))
    end

    # Collect index of first row for each hash
    uniquerow = Array(Int, size(A, dim))
    firstrow = Dict{Prehashed,Int}()
    for k = 1:size(A, dim)
        uniquerow[k] = get!(firstrow, Prehashed(hashes[k]), k)
    end
    uniquerows = collect(values(firstrow))

    # Check for collisions
    collided = falses(size(A, dim))
    @inbounds begin
        @nloops N i A d->(if d == dim
                              k = i_d
                              j_d = uniquerow[k]
                          else
                              j_d = i_d
                          end) begin
            if (@nref N A j) != (@nref N A i)
                collided[k] = true
            end
        end
    end

    if any(collided)
        nowcollided = BitArray(size(A, dim))
        while any(collided)
            # Collect index of first row for each collided hash
            empty!(firstrow)
            for j = 1:size(A, dim)
                collided[j] || continue
                uniquerow[j] = get!(firstrow, Prehashed(hashes[j]), j)
            end
            for v in values(firstrow)
                push!(uniquerows, v)
            end

            # Check for collisions
            fill!(nowcollided, false)
            @nloops N i A d->begin
                                 if d == dim
                                     k = i_d
                                     j_d = uniquerow[k]
                                     (!collided[k] || j_d == k) && continue
                                 else
                                     j_d = i_d
                                 end
                             end begin
                if (@nref N A j) != (@nref N A i)
                    nowcollided[k] = true
                end
            end
            (collided, nowcollided) = (nowcollided, collided)
        end
    end

    @nref N A d->d == dim ? sort!(uniquerows) : (1:size(A, d))
end

max_stars_repo_name: jakebolewski/julia

max_issues_repo_name: jakebolewski/julia

max_forks_repo_name: jakebolewski/julia


text:
#=
This is an integration test of a restarted calibration process.
=#

# Import modules to all processes
using Distributed
@everywhere using Pkg
@everywhere Pkg.activate(@__DIR__)
@everywhere begin
    using ArgParse
    using Glob
    import Random
    Random.seed!(1234)
    using CalibrateEDMF
    using CalibrateEDMF.DistributionUtils
    using CalibrateEDMF.Pipeline
    using CalibrateEDMF.HelperFuncs
    # Import EKP modules
    using EnsembleKalmanProcesses
    using EnsembleKalmanProcesses.ParameterDistributions
end
using JLD2
using Test

s = ArgParseSettings()
@add_arg_table s begin
    "--config"
    help = "Inverse problem config file path"
    arg_type = String
end
parsed_args = parse_args(ARGS, s)
config_rel_filepath = parsed_args["config"]
config_basename = first(split(config_rel_filepath, "."))

folder = joinpath(@__DIR__, "output", string(config_basename, "_julia_parallel"))
@assert isdir(folder)
outdir_path = open(f -> read(f, String), string(folder, config_basename, ".txt"))

include(joinpath(outdir_path, "config.jl"))
ekobjs = glob(joinpath(relpath(outdir_path), "ekobj_iter_*.jld2"))
iters = [parse(Int64, first(split(split(split(ekobj, "/")[end], "_")[end], "."))) for ekobj in ekobjs]
last_iteration = maximum(iters)

priors = deserialize_prior(load(joinpath(outdir_path, "prior.jld2")))
ekobj = load(ekobj_path(outdir_path, last_iteration))["ekp"]
config = get_config()
# continue calibration process for 5 iterations
config["process"]["N_iter"] = last_iteration + 5
restart_calibration(ekobj, priors, last_iteration, config, outdir_path)

# Dispatch SCM eval functions to workers
@everywhere begin
    scm_eval_train(version) = versioned_model_eval(version, $outdir_path, "train", $config)
    scm_eval_validation(version) = versioned_model_eval(version, $outdir_path, "validation", $config)
end
# Calibration process
N_iter = config["process"]["N_iter"]
@info "Continue EK updates for $(N_iter-last_iteration) iterations"
for iteration in (last_iteration + 1):N_iter
    @time begin
        @info "   iter = $iteration, ($(iteration - last_iteration) since restart)"
        versions = readlines(joinpath(outdir_path, "versions_$(iteration).txt"))
        ekp = load(ekobj_path(outdir_path, iteration))["ekp"]
        pmap(scm_eval_train, versions)
        pmap(scm_eval_validation, versions)
        ek_update(ekp, priors, iteration, config, versions, outdir_path)
    end
end

import NCDatasets
const NC = NCDatasets

nt = NC.Dataset(joinpath(outdir_path, "Diagnostics.nc"), "r") do ds
    (;
        mse_full_mean = Array(ds.group["metrics"]["mse_full_mean"]),
        mse_full_nn_mean = Array(ds.group["metrics"]["mse_full_nn_mean"]),
        mse_full_var = Array(ds.group["metrics"]["mse_full_var"]),
        val_mse_full_mean = Array(ds.group["metrics"]["val_mse_full_mean"]),
        val_mse_full_nn_mean = Array(ds.group["metrics"]["val_mse_full_nn_mean"]),
        val_mse_full_var = Array(ds.group["metrics"]["val_mse_full_var"]),
        phi_mean = Array(ds.group["ensemble_diags"]["phi_mean"]),
    )
end

mse_full_mean = nt.mse_full_mean[1:(end - 1)]
mse_full_nn_mean = nt.mse_full_nn_mean[1:(end - 1)]
mse_full_var = nt.mse_full_var[1:(end - 1)]
val_mse_full_mean = nt.val_mse_full_mean[1:(end - 1)]
val_mse_full_nn_mean = nt.val_mse_full_nn_mean[1:(end - 1)]
val_mse_full_var = nt.val_mse_full_var[1:(end - 1)]
phi_mean = nt.phi_mean

import Plots

# https://github.com/jheinen/GR.jl/issues/278#issuecomment-587090846
ENV["GKSwstype"] = "nul"

iters = 1:N_iter

@show iters
@show mse_full_mean
@show mse_full_nn_mean
@show mse_full_var
@show phi_mean

mse_full_nn_mean_std = mse_full_nn_mean .+ sqrt.(mse_full_var)
val_mse_full_nn_mean_std = val_mse_full_nn_mean .+ sqrt.(val_mse_full_var)

p1 = Plots.plot(iters, mse_full_nn_mean; label = "mean")
Plots.plot!(iters, mse_full_nn_mean_std; label = "")
Plots.plot!(iters, mse_full_nn_mean; fillrange = mse_full_nn_mean_std, fillalpha = 0.35, c = 1, label = "")
Plots.xlabel!("Iteration")
Plots.ylabel!("Train MSE (full)")
Plots.plot!(; left_margin = 40 * Plots.PlotMeasures.px)

p2 = Plots.plot(iters, val_mse_full_nn_mean; label = "mean")
Plots.plot!(iters, val_mse_full_nn_mean_std; label = "")
Plots.plot!(iters, val_mse_full_nn_mean; fillrange = val_mse_full_nn_mean_std, fillalpha = 0.35, c = 1, label = "")
Plots.xlabel!("Iteration")
Plots.ylabel!("Validation MSE (full)")
Plots.plot!(; left_margin = 40 * Plots.PlotMeasures.px)

Plots.plot(p1, p2; layout = (1, 2))
folder = joinpath(@__DIR__, "output", string(config_basename, "_julia_parallel"))
mkpath(folder)
Plots.savefig(joinpath(folder, string(config_basename, "_mse_continued.png")))

using Test
@testset "Julia Parallel Calibrate" begin
    # TODO: add better regression test (random seed)
    @test mse_full_mean[end] < mse_full_mean[1]
    @test mse_full_var[end] < mse_full_var[1]
end


text:
include(pwd() * "/test/ode_solution.jl")

# Impero Tests
include(pwd() * "/test/symbolic_tests.jl")
include(pwd() * "/test/symbolic_utils_tests.jl")
include(pwd() * "/test/test_equation.jl")

# Fourier Objects Tests
include(pwd() * "/test/fourier_tests.jl")
include(pwd() * "/test/fourier_field_tests.jl")

# Fourier Impero Interaction Test
include(pwd() * "/test/impero_fourier_test.jl")

# Timestepping Test
include(pwd() * "/test/time_step_test.jl")

# to_expr test
include(pwd() * "/test/to_expr_test.jl")

# Domain test
include(pwd() * "/test/domain_test.jl")

max_stars_repo_name: upiterbarg/QuasiGeostrophy

max_issues_repo_name: sandreza/QuasiGeostrophy

max_forks_repo_name: upiterbarg/QuasiGeostrophy


text:
using Luxor
import Random
Random.seed!(123);

function draw_kronecker(A, B)

    K = kron(A, B)'

    n, m = size(K)

    tiles = Tiler(12.5n, 12.5m, m, n, margin=0)

    for ((pos, n), k) in zip(tiles, K)
        if k==1
            sethue(Luxor.darker_green)
        elseif k==2
            sethue(Luxor.darker_red)
        elseif k==3
            sethue(Luxor.darker_purple)
        else
            randomhue()
        end
        if k > 0
            circle(pos, 5, :fill)
        end
    end
end

A = [1 1 0 0;
    0 1 0 1;
    1 0 1 1;
    0 1 0 1]

A = rand(Bool, 6, 6)
B = rand([0, 0, 0, 0, 1, 2, 3], 10, 15)

@pdf draw_kronecker(A', B) 1500 1000 "poster.pdf"

A = rand([0, 0, 1], 10, 15)
B = collect(0:3)

@pdf draw_kronecker(B, A) 240 600 "rectanle.pdf"

B = [1 2 3; 2 1 3; 3 2 1]'
A = rand([0, 0, 1], 10, 10)

@png draw_kronecker(B', A) 600 600 "logo.png"


text:
module CURAND

if Sys.iswindows()
    const libcurand = Libdl.find_library(["curand64_91","curand64_90","curand64_80","curand64_75"])
else
    const libcurand = Libdl.find_library(["libcurand"])
end
isempty(libcurand) && error("CURAND library cannot be found.")

include("define.jl")

macro apicall(f, args...)
    quote
        status = ccall(($f,libcurand), UInt32, $(map(esc,args)...))
        if status != CURAND_STATUS_SUCCESS
            Base.show_backtrace(STDOUT, backtrace())
            throw(errorstring(status))
        end
    end
end

function errorstring(status)
    status == CURAND_STATUS_SUCCESS && return "SUCCESS"
    status == CURAND_STATUS_VERSION_MISMATCH && return "VERSION_MISMATCH"
    status == CURAND_STATUS_NOT_INITIALIZED && return "NOT_INITIALIZED"
    status == CURAND_STATUS_ALLOCATION_FAILED && return "ALLOCATION_FAILED"
    status == CURAND_STATUS_TYPE_ERROR && return "TYPE_ERROR"
    status == CURAND_STATUS_OUT_OF_RANGE && return "OUT_OF_RANGE"
    status == CURAND_STATUS_LENGTH_NOT_MULTIPLE && return "LENGTH_NOT_MULTIPLE"
    status == CURAND_STATUS_DOUBLE_PRECISION_REQUIRED && return "DOUBLE_PRECISION_REQUIRED"
    status == CURAND_STATUS_LAUNCH_FAILURE && return "LAUNCH_FAILURE"
    status == CURAND_STATUS_PREEXISTING_FAILURE && return "PREEXISTING_FAILURE"
    status == CURAND_STATUS_INITIALIZATION_FAILED && return "INITIALIZATION_FAILED"
    status == CURAND_STATUS_ARCH_MISMATCH && return "ARCH_MISMATCH"
    status == CURAND_STATUS_INTERNAL_ERROR && return "INTERNAL_ERROR"
    throw("UNKNOWN ERROR")
end

const API_VERSION = begin
    ref = Ref{Cint}()
    @apicall :curandGetVersion (Ptr{Cint},) ref
    Int(ref[])
end
@info "CUDNN API $API_VERSION"

function _curand(rng)
    p = curandGenerator_t[0]
    curandCreateGenerator(p, rng)
    p[1]
end

###
function curng(rng_t)
  p = curandGenerator_t[0]
  curandCreateGenerator(p, rng_t)
  p[1]
end

rand_t = Float32
const _rng = curng(CURAND_RNG_PSEUDO_DEFAULT)
atexit(() -> curandDestroyGenerator(_rng))

cusrand(rng, seed) = curandSetPseudoRandomGeneratorSeed(rng, seed)
cusrand(seed) = curandSetPseudoRandomGeneratorSeed(_rng, seed)

cuoffset(rng, offset) = curandSetGeneratorOffset(rng, offset)
cuoffset(offset) = curandSetGeneratorOffset(_rng, offset)

cuorder(rng, order_t) = curandSetGeneratorOrdering(rng, order_t)
cuorder(order_t) = curandSetGeneratorOrdering(_rng, order_t)

function curng(rng_type, seed)
  rng = curng(rng_type)
  cusrand(rng, seed)
  rng
end

function curand(rng::curandGenerator_t, T::DataType, num::Int=1)
  arr = Main.CuArray(T, num)
  if  T == UInt32
    curandGenerate(rng, arr, num)
  elseif T == UInt64
    curandGenerateLongLong(rng, arr, num)
  elseif T == Float32
    curandGenerateUniform(rng, arr, num)
  elseif T == Float64
    curandGenerateUniformDouble(rng, arr, num)
  end
  arr
end

function curandn(rng::curandGenerator_t, T::DataType, num::Int=1; mean=T(0), stddev=T(1))
  arr = Main.CuArray(T, num)
  if T == Float32
    curandGenerateNormal(rng, arr, num, mean, stddev)
  elseif T == Float64
    curandGenerateNormalDouble(rng, arr, num, mean, stddev)
  end
  arr
end

function curandlogn(rng::curandGenerator_t, T::DataType, num::Int=1; mean=T(0), stddev=T(1))
  arr = Main.CuArray(T, num)
  if T == Float32
    curandGenerateLogNormal(rng, arr, num, mean, stddev)
  elseif T == Float64
    curandGenerateLogNormalDouble(rng, arr, num, mean, stddev)
    end
  arr
end

function curandpoisson(rng::curandGenerator_t, num::Int, lambda::Float64)
  arr = Main.CuArray(UInt32, num)
  curandGeneratePoisson(rng, arr, num, lambda)
  arr
end

curand(rng::curandGenerator_t, num::Int=1) = curand(rng, rand_t, num)
curand(T::DataType, num::Int=1) = curand(_rng, T, num)
curand(num::Int=1) = curand(_rng, rand_t, num)

function curandn(rng::curandGenerator_t, num::Int=1; mean=T(0), stddev=T(1))
  curandn(rng, rand_t, num, mean=mean, stdev=stdev)
end

function curandn(T::DataType, num::Int=1; mean=T(0), stddev=T(1))
  curandn(_rng, rand_t, num, mean=mean, stddev=stddev)
end

function curandn(num::Int=1; mean=0, stddev=1)
  curandn(_rng, rand_t, num, mean=rand_t(mean), stdev=rant_t(stdev))
end

function curandlogn(rng::curandGenerator_t, num::Int=1; mean=T(0), stddev=T(1))
  curandlogn(rng, rand_t, num, mean=mean, stdev=stdev)
end

function curandlogn(T::DataType, num::Int=1; mean=T(0), stddev=T(1))
  curandlogn(_rng, rand_t, num, mean=mean, stddev=stddev)
end

function curandlogn(num::Int=1; mean=0, stddev=1)
  curandlogn(_rng, rand_t, num, mean=rand_t(mean), stdev=rant_t(stdev))
end

curandpoisson(rng::curandGenerator_t, lambda::Float64) = curandpoisson(rng, 1, lambda)
curandpoisson(num::Int, lambda::Float64) = curandpoisson(_rng, num, lambda)
curandpoisson(lambda::Float64) = curandpoisson(_rng, 1, lambda)

function curandcreatedist(lambda::Float64)
  aptr = Ptr{Cvoid}[0]
  curandCreatePoissonDistribution(lambda, aptr)
  aptr[1]
end

function curandgetversion()
  aptr = Array(Int32, 1)
  curandGetVersion(aptr)
  aptr[1]
end

end


text:
export
    Quadrature,
    quadrature,
    quadrature_normal

"""
Lightweight wrapper for the nodes and weights of a quadrature. Also
include the domain.
"""
immutable Quadrature{T <: Real}
    domain::Interval{T}
    nodes::Vector{T}
    weights::Vector{T}
end

function Base.show(io::IO, q::Quadrature)
    print(io, "Quadrature of $(length(q.nodes)) points on $(q.domain)")
end

"""
Integrate a univariate function using the quadrature.
"""
function (q::Quadrature)(f)
    sum(weight*f(node) for (node, weight) in zip(q.nodes, q.weights))
end

@recipe function f{T <: Quadrature}(q::T)
    seriestype --> :scatter
    ylim --> (0, 1.05*maximum(q.weights))
    q.nodes, q.weights
end

"""
Return a quadrature rule for calculating an integral using the density
function of distribution `d`, truncated to `domain`.

Uses Legendre quadrature on a finite interval.
"""
function quadrature(n::Int, distribution::Distribution{Univariate,Continuous},
                    domain::Interval)
    @assert isfinite(domain)
    a, b = domain.lo, domain.hi
    y_nodes, weights = gausslegendre(n)
    x_nodes = fma.(y_nodes, (b-a)/2, middle(a,b))
    Quadrature(domain, x_nodes, # FIXME: remove [] around distribution in 0.6
               normalize!(pdf.([distribution], x_nodes) .* weights, 1))
end

function quadrature{T <: Distribution{Univariate,Continuous}}(n::Int, d::Truncated{T,Continuous})
    quadrature(n, d.untruncated, d.lower..d.upper)
end

"""
Return a quadrature that integrates a function under the `Normal(μ,σ)`
distribution using `N` nodes.
"""
function quadrature_normal(N::Int, μ, σ)
    nodes, weights = gausshermite(N)
    Quadrature(-Inf..Inf, nodes*√2*σ + μ, weights/√π)
end


text:
getnames(check_function::Function) = filter(check_function, uint32(0:65534))

# gets all the names currently boundo to programs
getProgramNames()      = getnames(glIsProgram)
getShaderNames()       = getnames(glIsShader)
getVertexArrayNames() = getnames(glIsVertexArray)

# display info for all active uniforms in a program
function getUniformsInfo(p::GLProgram)
    program = p.id
    # Get uniforms info (not in named blocks)
    @show activeUnif = glGetProgramiv(program, GL_ACTIVE_UNIFORMS)
    bufSize = 16
    name = zeros(UInt8, bufSize)
    buflen = Ref{GLsizei}(0)
    size = Ref{GLint}(0)
    type = Ref{GLenum}()

    for i=0:activeUnif-1
        glGetActiveUniform(program, i, bufSize, buflen, size, type, name)
        println(String(name),  " ", buflen[], " ", size[], " ", GLENUM(type[]).name)
    end
end

function uniform_name_type(p::GLProgram, location)
    bufSize = 32
    name = zeros(UInt8, bufSize)
    buflen = Ref{GLsizei}(0)
    size = Ref{GLint}(0)
    type = Ref{GLenum}()
    glGetActiveUniform(p.id, location, bufSize, buflen, size, type, name)
    println(String(name),  " ", buflen[], " ", size[], " ", GLENUM(type[]).name)
end

# display the values for uniforms in the default block
function getUniformInfo(p::GLProgram, uniName::Symbol)
    # is it a program ?
    @show program = p.id
    @show loc = glGetUniformLocation(program, uniName)
    @show name, typ, uniform_size = glGetActiveUniform(program, loc)
end

# display the values for a uniform in a named block
function getUniformInBlockInfo(p::GLProgram, blockName, uniName)
    program = p.id

    @show index = glGetUniformBlockIndex(program, blockName)
    if (index == GL_INVALID_INDEX)
        println("$uniName is not a valid uniform name in block $blockName")
    end
    @show bindIndex = glGetActiveUniformBlockiv(program, index, GL_UNIFORM_BLOCK_BINDING)
    @show bufferIndex = glGetIntegeri_v(GL_UNIFORM_BUFFER_BINDING, bindIndex)
    @show uniIndex = glGetUniformIndices(program, uniName)

    @show uniType = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_TYPE)
    @show uniOffset = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_OFFSET)
    @show uniSize = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_SIZE)
    @show uniArrayStride = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_ARRAY_STRIDE)
    @show uniMatStride = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_MATRIX_STRIDE)
end

# display information for a program's attributes
function getAttributesInfo(p::GLProgram)

    program = p.id
    # how many attribs?
    @show activeAttr = glGetProgramiv(program, GL_ACTIVE_ATTRIBUTES)
    # get location and type for each attrib
    for i=0:activeAttr-1
        @show name, typ, siz = glGetActiveAttrib(program,    i)
        @show loc = glGetAttribLocation(program, name)
    end
end

# display program's information
function getProgramInfo(p::GLProgram)
    # check if name is really a program
    @show program = p.id
    # Get the shader's name
    @show shaders = glGetAttachedShaders(program)
    for shader in shaders
        @show info = GLENUM(convert(GLenum, glGetShaderiv(shader, GL_SHADER_TYPE))).name
    end
    # Get program info
    @show info = glGetProgramiv(program, GL_PROGRAM_SEPARABLE)
    @show info = glGetProgramiv(program, GL_PROGRAM_BINARY_RETRIEVABLE_HINT)
    @show info = glGetProgramiv(program, GL_LINK_STATUS)
    @show info = glGetProgramiv(program, GL_VALIDATE_STATUS)
    @show info = glGetProgramiv(program, GL_DELETE_STATUS)
    @show info = glGetProgramiv(program, GL_ACTIVE_ATTRIBUTES)
    @show info = glGetProgramiv(program, GL_ACTIVE_UNIFORMS)
    @show info = glGetProgramiv(program, GL_ACTIVE_UNIFORM_BLOCKS)
    @show info = glGetProgramiv(program, GL_ACTIVE_ATOMIC_COUNTER_BUFFERS)
    @show info = glGetProgramiv(program, GL_TRANSFORM_FEEDBACK_BUFFER_MODE)
    @show info = glGetProgramiv(program, GL_TRANSFORM_FEEDBACK_VARYINGS)
end


text:
using ConstrainedDynamics

# Parameters
joint_axis = [1.0;0.0;0.0]

length1 = 0.5
width, depth = 0.5, 0.5
box1 = Box(width, depth, length1, 1., color = RGBA(1., 1., 0.))
b1 = Box(0.1, 0.1, .1, .1, color = RGBA(1., 0., 0.))

corners = [
    [[length1 / 2;length1 / 2;-length1 / 2]]
    [[length1 / 2;-length1 / 2;-length1 / 2]]
    [[-length1 / 2;length1 / 2;-length1 / 2]]
    [[-length1 / 2;-length1 / 2;-length1 / 2]]
    [[length1 / 2;length1 / 2;length1 / 2]]
    [[length1 / 2;-length1 / 2;length1 / 2]]
    [[-length1 / 2;length1 / 2;length1 / 2]]
    [[-length1 / 2;-length1 / 2;length1 / 2]]
]

# Links
origin = Origin{Float64}()

link1 = Body(box1)
links = [link1;[Body(b1) for i = 1:8]]

# Constraints
cf = 0.2
ineqcs = [InequalityConstraint(Friction(links[i + 1], [0;0;1.0], cf)) for i = 1:8]

joint0to1 = EqualityConstraint(Floating(origin, link1))
eqcs = [joint0to1;[EqualityConstraint(Fixed(link1, links[i + 1]; p1=corners[i])) for i = 1:8]]

shapes = [box1;b1]

mech = Mechanism(origin, links, eqcs, ineqcs, shapes = shapes)
setPosition!(link1,x = [0.;-2;1.5])
for i = 1:8
    setPosition!(link1, links[i + 1], p1 = corners[i])
end

ωtemp = [0.1;0.1;0.1]
setVelocity!(link1,v = [0;3;7.],ω = ωtemp)
for i = 1:8
    setVelocity!(link1, links[i + 1], p1 = corners[i])
end


text:
struct Positional
    position::Int
end

struct Keyword
    name::Symbol
end

const Argument = Union{Positional, Keyword, Expr}