text:
"Struct for an atmospheric profile"
struct AtmosphericProfile{FT}
    "Latitude of the AtmosphericProfile"
    lat
    "Longitude of the AtmosphericProfile"
    lon
    "Surface Pressure"
    psurf::FT
    "Temperature Profile"
    T::Array{FT,1}
    "Specific humidity profile"
    q::Array{FT,1}
    "Pressure Profile"
    p::Array{FT,1}
    "Pressure Levels"
    p_levels::Array{FT,1}
    "H2O Volume Mixing Ratio Profile"
    vmr_h2o::Array{FT,1}
    "Vertical Column Density (Dry)"
    vcd_dry::Array{FT,1}
    "Vertical Column Density (H2O)"
    vcd_h2o::Array{FT,1}
end

"Types for describing atmospheric parameters"
abstract type AbstractObsGeometry end

"Observation Geometry (basics)"
@with_kw struct ObsGeometry{FT} <: AbstractObsGeometry
    "Altitude of Observer `[Pa]`"
    obs_alt::FT
    "Solar Zenith Angle `[Degree]`"
    sza::FT
    "Viewing Zenith Angle(s) `[Degree]`"
    vza::Array{FT,1}
    "Viewing Azimuth Angle(s) `[Degree]`"
    vaz::Array{FT,1}
end

"Quadrature Types for RT streams"
abstract type AbstractQuadratureType end

"""
    struct RadauQuad

Use Gauss Radau quadrature scheme, which includes the SZA as point (see Sanghavi vSmartMOM)

"""
struct RadauQuad <:AbstractQuadratureType end

"""
    struct GaussQuadHemisphere

Use Gauss quadrature scheme, define interval [-1,1] within an hemisphere (90⁰), repeat for both

"""
struct GaussQuadHemisphere <:AbstractQuadratureType end

"""
    struct GaussQuadFullSphere

Use Gauss quadrature scheme, define interval [-1,1] for full sphere (180⁰), take half of it (less points near horizon compared to GaussQuadHemisphere)

"""
struct GaussQuadFullSphere <:AbstractQuadratureType end

"Abstract Type for Source Function Integration"
abstract type AbstractSourceType end

"Use Dummy Node Integration (DNI), SZA has to be a full node, Radau required"
struct DNI <:AbstractSourceType end

"Use Source Function Integration (SFI), Solar beam embedded in source term, can work with all quadrature schemes (faster)"
struct SFI <:AbstractSourceType end

"Abstract Type for Layer R,T and J matrices"
abstract type AbstractLayer end

"Composite Layer Matrices (`-/+` defined in τ coordinates, i.e. `-`=outgoing, `+`=incoming"
@with_kw mutable struct CompositeLayer{FT} <: AbstractLayer
    "Composite layer Reflectance matrix R (from + -> -)"
    R⁻⁺::AbstractArray{FT,3}
    "Composite layer Reflectance matrix R (from - -> +)"
    R⁺⁻::AbstractArray{FT,3}
    "Composite layer transmission matrix T (from + -> +)"
    T⁺⁺::AbstractArray{FT,3}
    "Composite layer transmission matrix T (from - -> -)"
    T⁻⁻::AbstractArray{FT,3}
    "Composite layer source matrix J (in + direction)"
    J₀⁺::AbstractArray{FT,3}
    "Composite layer source matrix J (in - direction)"
    J₀⁻::AbstractArray{FT,3}
end

"Added (Single) Layer Matrices (`-/+` defined in τ coordinates, i.e. `-`=outgoing, `+`=incoming"
@with_kw mutable struct AddedLayer{FT} <: AbstractLayer
    "Added layer Reflectance matrix R (from + -> -)"
    r⁻⁺::AbstractArray{FT,3}
    "Added layer transmission matrix T (from + -> +)"
    t⁺⁺::AbstractArray{FT,3}
    "Added layer Reflectance matrix R (from - -> +)"
    r⁺⁻::AbstractArray{FT,3}
    "Added layer transmission matrix T (from - -> -)"
    t⁻⁻::AbstractArray{FT,3}
    "Added layer source matrix J (in + direction)"
    J₀⁺::AbstractArray{FT,3}
    "Added layer source matrix J (in - direction)"
    J₀⁻::AbstractArray{FT,3}
end

"Abstract Type for Scattering Interfaces"
abstract type AbstractScatteringInterface end
"No scattering in either the added layer or the composite layer"
struct ScatteringInterface_00 <: AbstractScatteringInterface end
"No scattering in inhomogeneous composite layer; Scattering in homogeneous layer, added to bottom of the composite layer."
struct ScatteringInterface_01 <: AbstractScatteringInterface end
"Scattering in inhomogeneous composite layer; No scattering in homogeneous layer which is added to bottom of the composite layer."
struct ScatteringInterface_10 <: AbstractScatteringInterface end
"Scattering in inhomogeneous composite layer; Scattering in homogeneous layer, added to bottom of the composite layer."
struct ScatteringInterface_11 <: AbstractScatteringInterface end
"Scattering Interface between Surface and Composite Layer"
struct ScatteringInterface_AtmoSurf <: AbstractScatteringInterface end

"Abstract Type for Surface Types"
abstract type AbstractSurfaceType end

"Lambertian Surface (scalar per band)"
struct LambertianSurfaceScalar{FT} <: AbstractSurfaceType
    "Albedo (scalar)"
    albedo::FT
end

"Defined as Array (has to have the same length as the band!)"
struct LambertianSurfaceSpectrum{FT} <: AbstractSurfaceType
    "Albedo (vector)"
    albedo::AbstractArray{FT,1}
end

"Defined by polynomial terms as function of λ (degree derived from length of `a_coeff`)"
struct LambertianSurfacePolyFit{FT} <: AbstractSurfaceType
    "albedo(λ) = a_coeff[1] + a_coeff[2]*λ + a_coeff[3]*λ² + ... "
    a_coeff::AbstractArray{FT,1}
end

"""
    struct vSmartMOM_Parameters

A struct which holds all initial model parameters (before any computations)

# Fields
$(DocStringExtensions.FIELDS)
"""
mutable struct vSmartMOM_Parameters{FT<:Union{AbstractFloat, ForwardDiff.Dual}}

    "Float type to use in the RT (Float64/Float32)"
    float_type::DataType
    "band center wavelength (`nBand`)"
    λ_band::AbstractArray{FT}
    "reference wavelength"
    λ_ref::FT
    "Depolarization value"
    depol::FT
    "Truncation length for legendre terms (scalar for now, can do `nBand` later)"
    l_trunc::Integer
    "Exclusion angle for forward peak"
    Δ_angle::FT
    "Hard cutoff for maximum number of m iterations"
    max_m::Integer
    "Type of polarization (I/IQ/IQU/IQUV)"
    polarization_type::AbstractPolarizationType
    "Altitude of observer `[Pa]`"
    obs_alt::FT
    "Solar zenith angle [deg]"
    sza::FT
    "Viewing zenith angles [deg]"
    vza::AbstractArray{FT}
    "Viewing azimuthal angles [deg]"
    vaz::AbstractArray{FT}
    "Number of aerosol species"
    nAer::Integer
    "AOD at Reference wavelength (`nAer`)"
    τAer_ref::AbstractArray#{FT} #Suniti
    "Log mean radius (`nAer`)"
    μ::AbstractArray{FT} #Suniti
    "Log stddev of radius (`nAer`)"
    σ::AbstractArray{FT} #Suniti
    "Maximum radius"
    r_max::FT
    "Number of quadrature points for integration of size distribution"
    nquad_radius::Integer
    "Real part of refractive index (`nBand,nAer`)"
    nᵣ::AbstractArray#{FT}
    "Imag part of refractive index (`nBand,nAer`)"
    nᵢ::AbstractArray{FT}
    "Pressure peak [Pa] (`nAer`)"
    p₀::AbstractArray#{FT}
    "Pressure peak width [Pa] (`nAer`)"
    σp::AbstractArray{FT}
    "Path to atmospheric profile file"
    file::AbstractString
    "Time index to retrieve from file "
    timeIndex::Union{Integer, Nothing}
    "Latitude of atmospheric profile"
    lat::Union{Number, Nothing}
    "Longitude of atmospheric profile"
    lon::Union{Number, Nothing}
    "Length of profile reduction"
    profile_reduction_n::Integer
    "Starting wavelength for absorption grid (`nBand`)"
    spec_grid_start::AbstractArray{FT}
    "Ending wavelength for absorption grid (`nBand`)"
    spec_grid_end::AbstractArray{FT}
    "Number of spectral points in absorption grid (`nBand`)"
    spec_grid_n::AbstractArray{Integer}
    "Type of broadening function (Doppler/Lorentz/Voigt)"
    broadening_function::AbstractBroadeningFunction
    "Wing cutoff to use in cross-section calculation"
    wing_cutoff::Integer
    "Complex Error Function to use in Voigt calculations"
    CEF::AbstractComplexErrorFunction
    "Volume mixing ratio (needs to be per profile, gas species and band later)"
    vmr::FT
    "Surface reflectance type"
    brdf#::AbstractArray{AbstractSurfaceType}
    "Algorithm to use for fourier decomposition (NAI2/PCW)"
    decomp_type::AbstractFourierDecompositionType
    "Quadrature type for RT streams (RadauQuad/GaussQuadHemisphere/GaussQuadFullSphere)"
    quadrature_type::AbstractQuadratureType
    "Architecture to use for calculations (CPU/GPU)"
    architecture::AbstractArchitecture
end

"""
    struct QuadPoints{FT}

A struct which holds Quadrature points, weights, etc

# Fields
$(DocStringExtensions.FIELDS)
"""
struct QuadPoints{FT}
    "μ₀, cos(SZA)"
    μ₀::FT
    "Index in quadrature points with sun"
    iμ₀::Int
    "Index in quadrature points with sun (in qp_μN)"
    iμ₀Nstart::Int
    "Quadrature points"
    qp_μ::AbstractArray{FT,1}
    "Weights of quadrature points"
    wt_μ::AbstractArray{FT,1}
    "Quadrature points (repeated for polarizations)"
    qp_μN::AbstractArray{FT,1}
    "Weights of quadrature points (repeated for polarizations)"
    wt_μN::AbstractArray{FT,1}
    "Number of quadrature points"
    Nquad::Int
end

"""
    struct vSmartMOM_Model

A struct which holds all derived model parameters (including any computations)

# Fields
$(DocStringExtensions.FIELDS)
"""
mutable struct vSmartMOM_Model

    "Struct with all individual parameters"
    params::vSmartMOM_Parameters
    "Truncated aerosol optics"
    aerosol_optics::AbstractArray #Array{Array{AerosolOptics}}
    "Greek coefs in Rayleigh calculations"
    greek_rayleigh::GreekCoefs
    "Quadrature points/weights, etc"
    quad_points::QuadPoints
    "Array to hold cross-sections over entire atmospheric profile"
    τ_abs::AbstractArray
    "Rayleigh optical thickness"
    τRayl::AbstractArray
    "Aerosol optical thickness"
    τAer::AbstractArray
    "Observational Geometry (includes sza, vza, vaz)"
    obs_geom::ObsGeometry
    "Atmospheric profile to use"
    profile::AtmosphericProfile
    "Absorption model for computing absorption cross sections (will be an array later, for nGasSpecies,nBand)"
    #absorption_model::AbstractCrossSectionModel
    "Surface BRDF"
    brdf::AbstractSurfaceType
end

#"A struct to internally hold the computed atmosphere properties <<Suniti>>"
#abstract type ComputedProperties end

"""
    struct ComputedAtmosphereProperties

A struct which holds (for the entire atmosphere) all key layer optical properties required for the RT core solver

# Fields
$(DocStringExtensions.FIELDS)
"""
@with_kw struct ComputedAtmosphereProperties

    "Absorption optical depth vectors (wavelength dependent)"
    τ_λ_all
    "Albedo vectors (wavelength dependent)"
    ϖ_λ_all
    "Absorption optical depth scalars (not wavelength dependent)"
    τ_all
    "Albedo scalars (not wavelength dependent)"
    ϖ_all
    "Combined Z moments (forward)"
    Z⁺⁺_all
    "Combined Z moments (backward)"
    Z⁻⁺_all
    "Maximum dτs"
    dτ_max_all
    "dτs"
    dτ_all
    "Number of doublings (for all layers)"
    ndoubl_all
    "dτs (wavelength dependent)"
    dτ_λ_all
    "All expk"
    expk_all
    "Scattering flags"
    scatter_all
    "Sum of optical thicknesses of all layers above the current layer"
    τ_sum_all
    "Scattering interface type for each layer"
    scattering_interfaces_all
end

"""
    struct ComputedLayerProperties

A struct which holds all key layer optical properties required for the RT core solver

# Fields
$(DocStringExtensions.FIELDS)
"""
@with_kw struct ComputedLayerProperties

    "Absorption optical depth vector (wavelength dependent)"
    τ_λ
    "Albedo vector (wavelength dependent)"
    ϖ_λ
    "Absorption optical depth scalar (not wavelength dependent)"
    τ
    "Albedo scalar (not wavelength dependent)"
    ϖ
    "Combined Z moment (forward)"
    Z⁺⁺
    "Combined Z moment (backward)"
    Z⁻⁺
    "Maximum dτ"
    dτ_max
    "dτ"
    dτ
    "Number of doublings"
    ndoubl
    "dτ (wavelength dependent)"
    dτ_λ
    "expk"
    expk
    "Scattering flag"
    scatter
    "Sum of optical thicknesses of all layers above the current layer"
    τ_sum
    "Scattering interface type for current layer"
    scattering_interface
end


text:
mutable struct DEOptions{absType,relType,QT,tType,F1,F2,F3,F4,F5,F6,tstopsType,discType,ECType,SType,MI,tcache,savecache,disccache}
  maxiters::MI
  save_everystep::Bool
  adaptive::Bool
  abstol::absType
  reltol::relType
  gamma::QT
  qmax::QT
  qmin::QT
  qsteady_max::QT
  qsteady_min::QT
  failfactor::QT
  dtmax::tType
  dtmin::tType
  internalnorm::F1
  internalopnorm::F2
  save_idxs::SType
  tstops::tstopsType
  saveat::tstopsType
  d_discontinuities::discType
  tstops_cache::tcache
  saveat_cache::savecache
  d_discontinuities_cache::disccache
  userdata::ECType
  progress::Bool
  progress_steps::Int
  progress_name::String
  progress_message::F5
  timeseries_errors::Bool
  dense_errors::Bool
  beta1::QT
  beta2::QT
  qoldinit::QT
  dense::Bool
  save_on::Bool
  save_start::Bool
  save_end::Bool
  callback::F3
  isoutofdomain::F4
  unstable_check::F6
  verbose::Bool
  calck::Bool
  force_dtmin::Bool
  advance_to_tstop::Bool
  stop_at_next_tstop::Bool
end

"""
    ODEIntegrator
Fundamental `struct` allowing interactively stepping through the numerical solving of a differential equation.
The full documentation is hosted here: [http://docs.juliadiffeq.org/dev/basics/integrator.html](http://docs.juliadiffeq.org/dev/basics/integrator.html). This docstring
describes basic functionality only!

Initialize using `integrator = init(prob::ODEProblem, alg; kwargs...)`. The keyword args which are accepted are the same
[Common Solver Options](http://docs.juliadiffeq.org/dev/basics/common_solver_opts.html#Common-Solver-Options-1)
used by `solve`.

For reference, relevant fields of the `ODEIntegrator` are:

* `t` - time of the proposed step
* `u` - value at the proposed step
* `opts` - common solver options
* `alg` - the algorithm associated with the solution
* `f` - the function being solved
* `sol` - the current state of the solution
* `tprev` - the last timepoint
* `uprev` - the value at the last timepoint

`opts` holds all of the common solver options, and can be mutated to change the solver characteristics.
For example, to modify the absolute tolerance for the future timesteps, one can do:
```julia
integrator.opts.abstol = 1e-9
```
For more info see the linked documentation page.
"""
mutable struct ODEIntegrator{algType<:Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm},IIP,uType,duType,tType,pType,eigenType,QT,tdirType,ksEltype,SolType,F,CacheType,O,FSALType,EventErrorType,CallbackCacheType,IA} <: DiffEqBase.AbstractODEIntegrator{algType,IIP,uType,tType}
  sol::SolType
  u::uType
  du::duType
  k::ksEltype
  t::tType
  dt::tType
  f::F
  p::pType
  uprev::uType
  uprev2::uType
  duprev::duType
  tprev::tType
  alg::algType
  dtcache::tType
  dtchangeable::Bool
  dtpropose::tType
  tdir::tdirType
  eigen_est::eigenType
  EEst::QT
  qold::QT
  q11::QT
  erracc::QT
  dtacc::tType
  success_iter::Int
  iter::Int
  saveiter::Int
  saveiter_dense::Int
  cache::CacheType
  callback_cache::CallbackCacheType
  kshortsize::Int
  force_stepfail::Bool
  last_stepfail::Bool
  just_hit_tstop::Bool
  event_last_time::Int
  vector_event_last_time::Int
  last_event_error::EventErrorType
  accept_step::Bool
  isout::Bool
  reeval_fsal::Bool
  u_modified::Bool
  reinitialize::Bool
  isdae::Bool
  opts::O
  destats::DiffEqBase.DEStats
  initializealg::IA
  fsalfirst::FSALType
  fsallast::FSALType

  function ODEIntegrator{algType,IIP,uType,duType,tType,pType,eigenType,tTypeNoUnits,tdirType,ksEltype,SolType,
                F,CacheType,O,FSALType,EventErrorType,CallbackCacheType,InitializeAlgType}(
                sol,u,du,k,t,dt,f,p,uprev,uprev2,duprev,tprev,
      alg,dtcache,dtchangeable,dtpropose,tdir,
      eigen_est,EEst,qold,q11,erracc,dtacc,success_iter,
      iter,saveiter,saveiter_dense,cache,callback_cache,
      kshortsize,force_stepfail,last_stepfail,just_hit_tstop,
      event_last_time,vector_event_last_time,last_event_error,
      accept_step,isout,reeval_fsal,u_modified,reinitialize,isdae,
      opts,destats,initializealg) where {algType,IIP,uType,duType,tType,pType,
                                         eigenType,tTypeNoUnits,tdirType,
                                         ksEltype,SolType,F,CacheType,O,
                                         FSALType,EventErrorType,
                                         CallbackCacheType,InitializeAlgType}

      new{algType,IIP,uType,duType,tType,pType,eigenType,tTypeNoUnits,tdirType,ksEltype,SolType,
                  F,CacheType,O,FSALType,EventErrorType,CallbackCacheType,InitializeAlgType}(
                  sol,u,du,k,t,dt,f,p,uprev,uprev2,duprev,tprev,
      alg,dtcache,dtchangeable,dtpropose,tdir,
      eigen_est,EEst,qold,q11,erracc,dtacc,success_iter,
      iter,saveiter,saveiter_dense,cache,callback_cache,
      kshortsize,force_stepfail,last_stepfail,just_hit_tstop,
      event_last_time,vector_event_last_time,last_event_error,
      accept_step,isout,reeval_fsal,u_modified,reinitialize,isdae,
      opts,destats,initializealg) # Leave off fsalfirst and last
  end
end

# When this is changed, DelayDiffEq.jl must be changed as well!

ext: jl


text:
using Dates, Logging
using ShelfSpaceAllocation

# --- Arguments ---

time_limit = 3*60 # Seconds
mip_gap = 0.01
case = "medium"
product_path = joinpath(@__DIR__, "instances", case, "products.csv")
shelf_path = joinpath(@__DIR__, "instances", case, "shelves.csv")
output_dir = joinpath(@__DIR__, "output_exact", case, string(Dates.now()))

# ---

@info "Creating output directory"
mkpath(output_dir)

@info "Arguments" time_limit product_path shelf_path output_dir

@info "Loading parameters"
parameters = Params(product_path, shelf_path)
specs = Specs(height_placement=false, blocking=true)

@info "Creating the model"
model = ShelfSpaceAllocationModel(parameters, specs)

# Fix the block width for medium case
# using JuMP
# @constraint(model, [b = parameters.blocks, s = parameters.shelves],
#     model[:b_bs][b, s] == parameters.W_s[s]/2 * model[:z_bs][b, s]);

@info "Starting the optimization"
using JuMP, Gurobi
optimizer = optimizer_with_attributes(
    () -> Gurobi.Optimizer(Gurobi.Env()),
    "TimeLimit" => time_limit,
    "LogFile" => joinpath(output_dir, "gurobi.log"),
    "MIPFocus" => 3,
    "MIPGap" => mip_gap
)
set_optimizer(model, optimizer)
optimize!(model)

if termination_status(model) == MOI.INFEASIBLE
    exit()
end

@info "Variables"
variables = Variables(model)

@info "Objectives"
objectives = Objectives(model)

@info "Saving the results"
save_json(specs, joinpath(output_dir, "specs.json"))
save_json(parameters, joinpath(output_dir, "parameters.json"))
save_json(variables, joinpath(output_dir, "variables.json"))
save_json(objectives, joinpath(output_dir, "objectives.json"))

@info "Plotting"
using Plots, Parameters

ps = plot_planograms(parameters, variables)
for (i, p) in enumerate(ps)
    savefig(p, joinpath(output_dir, "planogram_$i.svg"))
end

p = plot_product_facings(parameters, variables)
savefig(p, joinpath(output_dir, "product_facings.svg"))

p = plot_demand_and_sales(parameters, variables)
savefig(p, joinpath(output_dir, "demand_and_sales.svg"))

p = plot_demand_sales_percentage(parameters, variables)
savefig(p, joinpath(output_dir, "demand_sales_percentage.svg"))

p = plot_allocation_amount(parameters, variables)
savefig(p, joinpath(output_dir, "allocation_amount.svg"))

p = plot_allocation_percentage(
    parameters, variables,
    with_optimizer(Gurobi.Optimizer, TimeLimit=60, LogToConsole=false))
savefig(p, joinpath(output_dir, "allocation_percentage.svg"))


text:
export discretize,timedirichlet

# Bivariate functions have BandedMatrix
op_eltype{T,D}(sp::Space{T,D,2})=BandedMatrix{promote_type(eltype(sp),eltype(domain(sp)))}
op_eltype_realdomain{T,D}(sp::Space{T,D,2})=BandedMatrix{promote_type(eltype(sp),real(eltype(domain(sp))))}

include("OperatorSchur.jl")
include("KroneckerOperator.jl")
include("dekron.jl")
include("diagop.jl")

include("factorizations.jl")
include("pdesolve.jl")
include("kron.jl")

## PDE

lap(d)=Laplacian(d)

function Laplacian(d::Union{ProductDomain,TensorSpace},k::Integer)
    @assert length(d)==2
    Dx2=Derivative(d,[2,0])
    Dy2=Derivative(d,[0,2])
    if k==1
        LaplacianWrapper(Dx2+Dy2,k)
    else
        LaplacianWrapper((Dx2+Dy2)^k,k)
    end
end

grad(d::ProductDomain)=[Derivative(d,[1,0]),Derivative(d,[0,1])]

for op in (:dirichlet,:neumann,:diffbcs)
    @eval begin
        function $op(d::Union{ProductDomain,TensorSpace},k...)
            @assert length(d)==2
            Bx=$op(d[1],k...)
            By=$op(d[2],k...)
            if isempty(Bx)
                I⊗By
            elseif isempty(By)
                Bx⊗I
            else
                [Bx⊗I;I⊗By]
            end
        end
    end
end

function timedirichlet(d::Union{ProductDomain,TensorSpace})
    @assert length(d.domains)==2
    Bx=dirichlet(d.domains[1])
    Bt=dirichlet(d.domains[2])[1]
    [I⊗Bt;Bx⊗I]
end

function *(B::Operator,f::ProductFun)
    if isafunctional(B)
        Fun(map(c->B*c,f.coefficients),space(f,2))
    else
        ProductFun(map(c->B*c,f.coefficients),space(f))
    end
end

*(f::ProductFun,B::Operator)=(B*(f.')).'


text:
using BinDeps
using Compat

@BinDeps.setup

libpq = library_dependency("libpq")

# package managers
provides(AptGet, "libpq5", libpq)

provides(Yum, "postgresql-libs", libpq)
provides(Yum, "postgresql96-libs", libpq)
provides(Yum, "postgresql95-libs", libpq)
provides(Yum, "postgresql94-libs", libpq)

if is_apple()
    if Pkg.installed("Homebrew") === nothing
        error("Homebrew package not installed, please run Pkg.add(\"Homebrew\")")
    end
    using Homebrew
    provides(Homebrew.HB, "libpq", libpq, os = :Darwin; installed_libpath="/usr/local/opt/libpq/lib")
    provides(Homebrew.HB, "postgresql", libpq, os = :Darwin)
    provides(Homebrew.HB, "postgresql@9.6", libpq, os = :Darwin; installed_libpath="/usr/local/opt/postgresql@9.6/lib")
    provides(Homebrew.HB, "postgresql@9.5", libpq, os = :Darwin; installed_libpath="/usr/local/opt/postgresql@9.5/lib")
    provides(Homebrew.HB, "postgresql@9.4", libpq, os = :Darwin; installed_libpath="/usr/local/opt/postgresql@9.4/lib")
end

@BinDeps.install


text:
abstract type PulseBlock end

struct Sequence{T} <: PulseBlock
    J::Matrix{T}
    fields::Vector{Vector{T}}
    durations::Vector{T}
    anisotropy::Vector{T}
    state::Matrix{T}
end

function Base.repeat(s::Sequence, n::Integer)
    fields = repeat(s.fields, n) #, outer=(1, n))
    durations = repeat(s.durations, n)
    Sequence(s.J, fields, durations, s.anisotropy, s.state)
end

function Base.iterate(s::Sequence, state=1)
    state > length(s.durations) ? nothing : s.durations[state], state + 1
end

Base.getindex(s::Sequence, i::Int) = (s.fields[i], s.durations[i])

function Base.setindex!(s::Sequence, v::Tuple{Vector{Float64}, Float64}, i::Int)
    s.fields[i] = v[1]
    s.durations[i] = v[2]
end

function Base.insert!(s::Sequence, i::Int, v::Tuple{Vector{Float64}, Float64})
    insert!(s.fields, i, v[1])
    insert!(s.durations, i, v[2])
end

absolute_times(s::Sequence) = [sum(s.durations[1:(i-1)]) for i in 1:length(s.durations)]

function clostest_index(s::Sequence, t)
    abs_times = absolute_times(s)
    argmin(abs.(t .- abs_times))
end

function create_time_choice(s::Sequence)
    f = function(integrator)
        if integrator.t == sum(s.durations[1:end-1])
            return nothing
        end
        index = clostest_index(s, integrator.t)
        integrator.t +  s.durations[index]
    end
end

function create_affect!(s::Sequence)
    f = function(integrator)
        J = s.J
        index = clostest_index(s, integrator.t)
        #h = s.fields[index]
        # integrator.p = ArrayPartition(J, h)
        system = XYZSystem(J, [0.0, 0.0, 0.0], [1.0, 1.0, 0])
        integrator.f = build_hamiltonian!(system)
    end
    f
end

using DifferentialEquations

DiffEqCallbacks.IterativeCallback(s::Sequence) = IterativeCallback(
    create_time_choice(s),
    create_affect!(s),
    save_positions=(false, false)
)

function wahuha(J, h, T, anisotropy, state)
    X = h .* [1., 0., 0.]
    Y = h .* [0., 1., 0.]
    Null = [0., 0., 0.]
    fields = [Null, X, Null, -Y, Null, Y, Null, -X, Null]

    t = π/2/h
    durations = [T, t, T, t, 2*T, t, T, t, T]
    return Sequence(J, fields, durations, anisotropy, state)
end

wahuha(J, h, T, anisotropy, state, n) = repeat(wahuha(J, h, T, anisotropy, state), n)

max_stars_repo_name: titusfranz/DTWA

max_issues_repo_name: titusfranz/DTWA

max_forks_repo_name: titusfranz/DTWA


text:
# # Gradient of a scalar field
#
# This example shows different methods to compute the gradient of a real-valued
# 3D scalar field ``θ(\bm{x})`` in Fourier space, where $\bm{x} = (x, y, z)$.
# It is assumed that the field is periodic with period $L = 2π$ along all
# dimensions.
#
# ## General procedure
#
# The discrete Fourier expansion of ``θ`` writes
# ```math
# θ(\bm{x}) = ∑_{\bm{k} ∈ \Z^3} \hat{θ}(\bm{k}) \, e^{i \bm{k} ⋅ \bm{x}},
# ```
# where $\bm{k} = (k_x, k_y, k_z)$ are the Fourier wave numbers and $\hat{θ}$ is
# the discrete Fourier transform of $θ$.
# Then, the spatial derivatives of $θ$ are given by
# ```math
# \frac{∂ θ(\bm{x})}{∂ x_i} =
# ∑_{\bm{k} ∈ \Z^3} i k_i \hat{θ}(\bm{k}) \, e^{i \bm{k} ⋅ \bm{x}},
# ```
# where the subscript $i$ denotes one of the spatial components $x$, $y$ or
# $z$.
#
# In other words, to compute $\bm{∇} θ = (∂_x θ, ∂_y θ, ∂_z θ)$, one has to:
# 1. transform $θ$ to Fourier space to obtain $\hat{θ}$,
# 2. multiply $\hat{θ}$ by $i \bm{k}$,
# 3. transform the result back to physical space to obtain $\bm{∇} θ$.

# ## Preparation

# In this section, we initialise a random real-valued scalar field $θ$ and compute
# its FFT.
# For more details see the [Tutorial](@ref).

using MPI
using PencilFFTs
using Random

MPI.Init()

## Input data dimensions (Nx × Ny × Nz)
dims = (64, 32, 64)

## Apply a 3D real-to-complex (r2c) FFT.
transform = Transforms.RFFT()

## Automatically create decomposition configuration
comm = MPI.COMM_WORLD
pen = Pencil(dims, comm)

## Create plan
plan = PencilFFTPlan(pen, transform)

## Allocate data and initialise field
θ = allocate_input(plan)
randn!(θ)

## Perform distributed FFT
θ_hat = plan * θ
nothing  # hide

# Finally, we initialise the output that will hold ∇θ in Fourier space.
# Noting that ∇θ is a vector field, we choose to store it as a tuple of
# 3 PencilArrays.

∇θ_hat = allocate_output(plan, Val(3))

## This is equivalent:
## ∇θ_hat = ntuple(d -> similar(θ_hat), Val(3))

summary(∇θ_hat)

# ## Fourier wave numbers

# In general, the Fourier wave numbers are of the form
# ``k_i = 0, ±\frac{2π}{L_i}, ±\frac{4π}{L_i}, ±\frac{6π}{L_i}, …``,
# where ``L_i`` is the period along dimension ``i``.
# When a real-to-complex Fourier transform is applied, roughly half of
# these wave numbers are redundant due to the Hermitian symmetry of the complex
# Fourier coefficients.
# In practice, this means that for the fastest dimension $x$ (along which
# a real-to-complex transform is performed), the negative wave numbers are
# dropped, i.e. ``k_x = 0, \frac{2π}{L_x}, \frac{4π}{L_x}, …``.

# The `AbstractFFTs` package provides a convenient way to generate the Fourier
# wave numbers, using the functions
# [`fftfreq`](https://juliamath.github.io/AbstractFFTs.jl/stable/api/#AbstractFFTs.fftfreq)
# and
# [`rfftfreq`](https://juliamath.github.io/AbstractFFTs.jl/stable/api/#AbstractFFTs.rfftfreq).
# We can use these functions to initialise a "grid" of wave numbers associated to
# our 3D real-to-complex transform:

using AbstractFFTs: fftfreq, rfftfreq

box_size = (2π, 2π, 2π)  # Lx, Ly, Lz
sample_rate = 2π .* dims ./ box_size

## In our case (Lx = 2π and Nx even), this gives kx = [0, 1, 2, ..., Nx/2].
kx = rfftfreq(dims[1], sample_rate[1])

## In our case (Ly = 2π and Ny even), this gives
## ky = [0, 1, 2, ..., Ny/2-1, -Ny/2, -Ny/2+1, ..., -1] (and similarly for kz).
ky = fftfreq(dims[2], sample_rate[2])
kz = fftfreq(dims[3], sample_rate[3])

kvec = (kx, ky, kz)

# Note that `kvec` now contains the wave numbers associated to the global domain.
# In the following, we will only need the wave numbers associated to the portion
# of the domain handled by the local MPI process.

# ## [Method 1: global views](@id gradient_method_global)

# [`PencilArray`](https://jipolanco.github.io/PencilArrays.jl/dev/PencilArrays/#PencilArrays.PencilArray)s, returned for instance by [`allocate_input`](@ref)
# and  [`allocate_output`](@ref), take indices that start at 1, regardless of the
# location of the subdomain associated to the local process on the global grid.
# (In other words, `PencilArray`s take *local* indices.)
# On the other hand, we have defined the wave number vector `kvec` which,
# for each MPI process, is defined over the global domain, and as such it takes
# *global* indices.

# One straightforward way of making data arrays compatible with wave numbers is
# to use global views, i.e. arrays that take global indices.
# These are generated from `PencilArray`s by calling the [`global_view`](https://jipolanco.github.io/PencilArrays.jl/dev/PencilArrays/#PencilArrays.global_view-Tuple{PencilArray})
# function.
# Note that, in general, global indices do *not* start at 1 for a given MPI
# process.
# A given process will own a range of data given by indices in `(i1:i2, j1:j2,
# k1:k2)`.

θ_glob = global_view(θ_hat)
∇θ_glob = global_view.(∇θ_hat)
summary(θ_glob)

# Once we have global views, we can combine data and wave numbers using the
# portion of global indices owned by the local MPI process, as shown below.
# We can use `CartesianIndices` to iterate over the global indices associated to
# the local process.

for I in CartesianIndices(θ_glob)
    i, j, k = Tuple(I)  # unpack indices

    ## Wave number vector associated to current Cartesian index.
    local kx, ky, kz  # hide
    kx = kvec[1][i]
    ky = kvec[2][j]
    kz = kvec[3][k]

    ## Compute gradient in Fourier space.
    ## Note that modifying ∇θ_glob also modifies the original PencilArray ∇θ_hat.
    ∇θ_glob[1][I] = im * kx * θ_glob[I]
    ∇θ_glob[2][I] = im * ky * θ_glob[I]
    ∇θ_glob[3][I] = im * kz * θ_glob[I]
end

# The above loop can be written in a slightly more efficient manner by precomputing
# `im * θ_glob[I]`:

@inbounds for I in CartesianIndices(θ_glob)
    i, j, k = Tuple(I)

    local kx, ky, kz  # hide
    kx = kvec[1][i]
    ky = kvec[2][j]
    kz = kvec[3][k]

    u = im * θ_glob[I]

    ∇θ_glob[1][I] = kx * u
    ∇θ_glob[2][I] = ky * u
    ∇θ_glob[3][I] = kz * u
end

# Also note that the above can be easily written in a more generic way, e.g. for
# arbitrary dimensions, thanks in part to the use of `CartesianIndices`.
# Moreover, in the above there is no notion of the dimension permutations
# discussed in [the tutorial](@ref tutorial:output_data_layout), as it is all
# hidden behind the implementation of `PencilArray`s.
# And as seen later in the [benchmarks](@ref gradient_benchmarks),
# these (hidden) permutations have zero cost, as the speed is identical
# to that of a function that explicitly takes into account these permutations.

# Finally, we can perform a backwards transform to obtain $\bm{∇} θ$ in physical
# space:

∇θ = plan \ ∇θ_hat;

# Note that the transform is automatically broadcast over the three fields
# of the `∇θ_hat` vector, and the result `∇θ` is also a tuple of
# three `PencilArray`s.

# ## [Method 2: explicit global indexing](@id gradient_method_global_explicit)

# Sometimes, one does not need to write generic code.
# In our case, one often knows the dimensionality of the problem and the
# memory layout of the data (i.e. the underlying index permutation).

# Below is a reimplementation of the above loop, using explicit indices instead
# of `CartesianIndices`, and assuming that the underlying index permutation is
# `(3, 2, 1)`, that is, data is stored in $(z, y, x)$ order.
# As discussed in [the tutorial](@ref tutorial:output_data_layout),
# this is the default for transformed arrays.
# This example also serves as a more explicit explanation for what is going on
# in the [first method](@ref gradient_method_global).

## Get local data range in the global grid.
rng = axes(θ_glob)  # = (i1:i2, j1:j2, k1:k2)

# For the loop below, we're assuming that the permutation is (3, 2, 1).
# In other words, the fastest index is the *last* one, and not the first one as
# it is usually in Julia.
# If the permutation is not (3, 2, 1), things will still work (well, except for
# the assertion below!), but the loop order will not be optimal.

@assert permutation(θ_hat) === Permutation(3, 2, 1)

@inbounds for i in rng[1], j in rng[2], k in rng[3]
    local kx, ky, kz  # hide
    kx = kvec[1][i]
    ky = kvec[2][j]
    kz = kvec[3][k]

    ## Note that we still access the arrays in (i, j, k) order.
    ## (The permutation happens behind the scenes!)
    u = im * θ_glob[i, j, k]

    ∇θ_glob[1][i, j, k] = kx * u
    ∇θ_glob[2][i, j, k] = ky * u
    ∇θ_glob[3][i, j, k] = kz * u
end

# ## [Method 3: using local indices](@id gradient_method_local)

# Alternatively, we can avoid global views and work directly on `PencilArray`s
# using local indices that start at 1.
# In this case, part of the strategy is to construct a "local" grid of wave
# numbers that can also be accessed with local indices.
# This can be conveniently done using the
# [`localgrid`](https://jipolanco.github.io/PencilArrays.jl/dev/LocalGrids/#PencilArrays.LocalGrids.localgrid)
# function of the PencilArrays.jl package, which accepts a `PencilArray` (or
# its associated `Pencil`) and the global coordinates (here `kvec`):

grid_fourier = localgrid(θ_hat, kvec)

# Note that one can directly iterate on the returned grid object:

@inbounds for I in CartesianIndices(grid_fourier)
    ## Wave number vector associated to current Cartesian index.
    local k⃗  # hide
    k⃗ = grid_fourier[I]
    u = im * θ_hat[I]
    ∇θ_hat[1][I] = k⃗[1] * u
    ∇θ_hat[2][I] = k⃗[2] * u
    ∇θ_hat[3][I] = k⃗[3] * u
end

# This implementation is as efficient as the other examples, while being
# slightly shorter to write.
# Moreover, it is quite generic, and can be made independent of the number of
# dimensions with little effort.

# ## [Method 4: using broadcasting](@id gradient_method_broadcast)

# Finally, note that the local grid object returned by `localgrid` makes it is
# possible to compute the gradient using broadcasting, thus fully avoiding scalar
# indexing.
# This can be quite convenient in some cases, and can also be very useful if
# one is working on GPUs (where scalar indexing is prohibitively expensive).
# Using broadcasting, the above examples simply become:

@. ∇θ_hat[1] = im * grid_fourier[1] * θ_hat
@. ∇θ_hat[2] = im * grid_fourier[2] * θ_hat
@. ∇θ_hat[3] = im * grid_fourier[3] * θ_hat
nothing  # hide

# Once again, as shown in the [benchmarks](@ref gradient_benchmarks) further
# below, this method performs quite similarly to the other ones.

# ## Summary

# The `PencilArrays` module provides different alternatives to deal with
# MPI-distributed data that may be subject to dimension permutations.
# In particular, one can choose to work with *global* indices (first two
# examples), with *local* indices (third example), or to avoid scalar indexing
# altogether (fourth example).

# If one wants to stay generic, making sure that the same code will work for
# arbitrary dimensions and will be efficient regardless of the underlying
# dimension permutation, methods [1](@ref gradient_method_global), [3](@ref
# gradient_method_local) or [4](@ref gradient_method_broadcast) should be
# preferred.
# These use `CartesianIndices` and make no assumptions on possible dimension
# permutations, which are by default enabled in the output of PencilFFTs
# transforms.
# In fact, such permutations are completely invisible in the implementations.

# The [second method](@ref gradient_method_global_explicit) uses explicit
# `(i, j, k)` indices.
# It assumes that the underlying permutation is `(3, 2, 1)` to loop with `i` as
# the *slowest* index and `k` as the *fastest*, which is the optimal order in
# this case given the permutation.
# As such, the implementation is less generic than the others, and
# differences in performance are negligible with respect to more generic variants.

# ## [Benchmark results](@id gradient_benchmarks)

# The following are the benchmark results obtained from running
# [`examples/gradient.jl`](https://github.com/jipolanco/PencilFFTs.jl/blob/master/examples/gradient.jl)
# on a laptop, using 2 MPI processes and Julia 1.7.2, with an input array of
# global dimensions ``64 × 32 × 64``.
# The different methods detailed above are marked on the right.
# The "lazy" marks indicate runs where the wave numbers were represented by
# lazy `Frequencies` objects (returned by `rfftfreq` and `fftfreq`). Otherwise,
# they were collected into `Vector`s.
# For some reason, plain `Vector`s are faster when working with grids generated
# by `localgrid`.

# In the script, additional implementations can be found which rely on a more
# advanced understanding of permutations and on the internals of the
# [`PencilArrays`](https://jipolanco.github.io/PencilArrays.jl/dev/) package.
# For instance, `gradient_local_parent!` directly works with the raw
# data stored in Julia `Array`s, while `gradient_local_linear!` completely
# avoids `CartesianIndices` while staying generic and efficient.
# Nevertheless, these display roughly the same performance as the above examples.
#
#         gradient_global_view!...                  89.900 μs
#         gradient_global_view! (lazy)...           92.060 μs  [Method 1]
#         gradient_global_view_explicit!...         88.958 μs
#         gradient_global_view_explicit! (lazy)...  81.055 μs  [Method 2]
#         gradient_local!...                        92.305 μs
#         gradient_grid!...                         92.770 μs
#         gradient_grid! (lazy)...                  101.388 μs  [Method 3]
#         gradient_grid_broadcast!...               88.606 μs
#         gradient_grid_broadcast! (lazy)...        151.020 μs  [Method 4]
#         gradient_local_parent!...                 92.248 μs
#         gradient_local_linear!...                 91.212 μs
#         gradient_local_linear_explicit!...        90.992 μs


text:
"""
Copyright (c) 2013 Dahua Lin

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
## -------------------------- AVL tree structure ---------------------------- ##
## From DataStructures v.0.18.8

# it has unique keys
# leftChild has keys which are less than the node
# rightChild has keys which are greater than the node
# height stores the height of the subtree.
mutable struct AVLTreeNode{K}
    height::Int8
    leftChild::Union{AVLTreeNode{K}, Nothing}
    rightChild::Union{AVLTreeNode{K}, Nothing}
    subsize::Int32
    data::K
    idx::Int64

    AVLTreeNode{K}(d::K, idx::Int64) where K = new{K}(1, nothing, nothing, 1, d, idx)
end

AVLTreeNode(d, idx) = AVLTreeNode{Any}(d, idx)

AVLTreeNode_or_null{T} = Union{AVLTreeNode{T}, Nothing}

"""
    mutable struct AVLTree{T}()
Creates an empty AVL tree. `tree[i]` returns `i`th smallest node.

## Supported operations
    Base.haskey(tree::AVLTree{K}, d::K)
Returns a bool which indicates presence of a node with key `d`

    Base.insert!(tree::AVLTree{K}, d::K)
Inserts a node with key `d` of type `K` (`Base.push!` is recommanded)

    Base.push!(tree::AVLTree{K}, key0)
Inserts a node with key `d` of arbitrary type that can be converted into type `K`

    Base.delete!(tree::AVLTree{K}, d::K)
Deletes a node with key `d` of type `K`

    sorted_rank(tree::AVLTree{K}, key::K)
Returns the rank of `key` present in the `tree`, if it present. A KeyError is thrown if `key` is not present.

    Base.getindex(tree::AVLTree{K}, ind::Integer)
Returns the `ind`th node
"""
mutable struct AVLTree{T}
    root::AVLTreeNode_or_null{T}
    count::Int

    AVLTree{T}() where T = new{T}(nothing, 0)
end

AVLTree() = AVLTree{Any}()

Base.length(tree::AVLTree) = tree.count

get_height(node::Union{AVLTreeNode, Nothing}) = (node == nothing) ? 0 : node.height

# balance is the difference of height between leftChild and rightChild of a node.
function get_balance(node::Union{AVLTreeNode, Nothing})
    if node == nothing
        return 0
    else
        return get_height(node.leftChild) - get_height(node.rightChild)
    end
end

# computes the height of the subtree, which basically is
# one added the maximum of the height of the left subtree and right subtree
compute_height(node::AVLTreeNode) = 1 + max(get_height(node.leftChild), get_height(node.rightChild))

get_subsize(node::AVLTreeNode_or_null) = (node == nothing) ? 0 : node.subsize

# compute the subtree size
function compute_subtree_size(node::AVLTreeNode_or_null)
    if node == nothing
        return 0
    else
        L = get_subsize(node.leftChild)
        R = get_subsize(node.rightChild)
        return (L + R + 1)
    end
end

"""
    left_rotate(node_x::AVLTreeNode)
Performs a left-rotation on `node_x`, updates height of the nodes, and returns the rotated node.
"""
function left_rotate(z::AVLTreeNode)
    y = z.rightChild
    α = y.leftChild
    y.leftChild = z
    z.rightChild = α
    z.height = compute_height(z)
    y.height = compute_height(y)
    z.subsize = compute_subtree_size(z)
    y.subsize = compute_subtree_size(y)
    return y
end

"""
    right_rotate(node_x::AVLTreeNode)
Performs a right-rotation on `node_x`, updates height of the nodes, and returns the rotated node.
"""
function right_rotate(z::AVLTreeNode)
    y = z.leftChild
    α = y.rightChild
    y.rightChild = z
    z.leftChild = α
    z.height = compute_height(z)
    y.height = compute_height(y)
    z.subsize = compute_subtree_size(z)
    y.subsize = compute_subtree_size(y)
    return y
end

"""
   minimum_node(tree::AVLTree, node::AVLTreeNode)
Returns the AVLTreeNode with minimum value in subtree of `node`.
"""
function minimum_node(node::Union{AVLTreeNode, Nothing})
    while node != nothing && node.leftChild != nothing
        node = node.leftChild
    end
    return node
end

function search_node(tree::AVLTree{K}, d::K) where K
    prev = nothing
    node = tree.root
    while node != nothing && node.data != nothing && node.data != d

        prev = node
        if d < node.data
            node = node.leftChild
        else
            node = node.rightChild
        end
    end

    return (node == nothing) ? prev : node
end

function Base.haskey(tree::AVLTree{K}, d::K) where K
    (tree.root == nothing) && return false
    node = search_node(tree, d)
    return (node.data == d)
end

Base.in(key, tree::AVLTree) = haskey(tree, key)

function Base.insert!(tree::AVLTree{K}, d::K, idx::Int64) where K

    function insert_node(node::Union{AVLTreeNode, Nothing}, key, idx)
        if node == nothing
            return AVLTreeNode{K}(key, idx)
        elseif key < node.data
            node.leftChild = insert_node(node.leftChild, key, idx)
        else
            node.rightChild = insert_node(node.rightChild, key, idx)
        end

        node.subsize = compute_subtree_size(node)
        node.height = compute_height(node)
        balance = get_balance(node)

        if balance > 1
            if key < node.leftChild.data
                return right_rotate(node)
            else
                node.leftChild = left_rotate(node.leftChild)
                return right_rotate(node)
            end
        end

        if balance < -1
            if key > node.rightChild.data
                return left_rotate(node)
            else
                node.rightChild = right_rotate(node.rightChild)
                return left_rotate(node)
            end
        end

        return node
    end

    haskey(tree, d) && return tree  #------------------------- key should be unique ------------------------------#

    tree.root = insert_node(tree.root, d, idx)
    tree.count += 1
    return tree
end

function Base.push!(tree::AVLTree{K}, key0, idx) where K
    key = convert(K, key0)
    insert!(tree, key, idx)
end

function Base.delete!(tree::AVLTree{K}, d::K) where K

    function delete_node!(node::Union{AVLTreeNode, Nothing}, key)
        if key < node.data
            node.leftChild = delete_node!(node.leftChild, key)
        elseif key > node.data
            node.rightChild = delete_node!(node.rightChild, key)
        else
            if node.leftChild == nothing
                result = node.rightChild
                return result
            elseif node.rightChild == nothing
                result = node.leftChild
                return result
            else
                result = minimum_node(node.rightChild)
                node.data = result.data
                node.idx = result.idx
                node.rightChild = delete_node!(node.rightChild, result.data)
            end
        end

        node.subsize = compute_subtree_size(node)
        node.height = compute_height(node)
        balance = get_balance(node)

        if balance > 1
            if get_balance(node.leftChild) >= 0
                return right_rotate(node)
            else
                node.leftChild = left_rotate(node.leftChild)
                return right_rotate(node)
            end
        end

        if balance < -1
            if get_balance(node.rightChild) <= 0
                return left_rotate(node)
            else
                node.rightChild = right_rotate(node.rightChild)
                return left_rotate(node)
            end
        end

        return node
    end

    # if the key is not in the tree, do nothing and return the tree
    !haskey(tree, d) && return tree

    # if the key is present, delete it from the tree
    tree.root = delete_node!(tree.root, d)
    tree.count -= 1
    return tree
end

"""
    sorted_rank(tree::AVLTree, key)
Returns the rank of `key` present in the `tree`, if it present. A KeyError is thrown if `key` is not present.
"""
function sorted_rank(tree::AVLTree{K}, key::K) where K
    !haskey(tree, key) && throw(KeyError(key))
    node = tree.root
    rank = 0
    while node.data != key
        if (node.data < key)
            rank += (1 + get_subsize(node.leftChild))
            node = node.rightChild
        else
            node = node.leftChild
        end
    end
    rank += (1 + get_subsize(node.leftChild))
    return rank
end

function Base.getindex(tree::AVLTree{K}, ind::Integer) where K
    @boundscheck (1 <= ind <= tree.count) || throw(BoundsError("$ind should be in between 1 and $(tree.count)"))
    function traverse_tree(node::AVLTreeNode_or_null, idxx)
        if (node != nothing)
            L = get_subsize(node.leftChild)
            if idxx <= L
                return traverse_tree(node.leftChild, idxx)
            elseif idxx == L + 1
                return node.data
            else
                return traverse_tree(node.rightChild, idxx - L - 1)
            end
        end
    end
    value = traverse_tree(tree.root, ind)
    return value
end

function inorder_traverse_node(node::AVLTreeNode_or_null)
    if node.leftChild == nothing && node.rightChild == nothing
        return [node]
    elseif node.leftChild == nothing
        return [node; inorder_traverse_node(node.rightChild)]
    elseif node.rightChild == nothing
        return [inorder_traverse_node(node.leftChild); node]
    else
        return [inorder_traverse_node(node.leftChild); node; inorder_traverse_node(node.rightChild)]
    end
end

"""
    inorder_traverse(tree::AVLTree)
Traverses tree in inorder way. Returns an Array{Any, 2} with elements of _node type.
"""
function inorder_traverse(tree::AVLTree)
    if tree.count == 0
        return []
    else
        return inorder_traverse_node(tree.root)
    end
end


text:
"""
abstract VariableState{F<:VariateForm}

Root of variable state type hierarchy
"""
abstract VariableState{F<:VariateForm}

typealias VariableStateVector{S<:VariableState} Vector{S}

variate_form{F<:VariateForm}(::Type{VariableState{F}}) = F

"""
Basic univariate variable state type

# Constructors

## BasicUnvVariableState{N<:Number}(value::N)

Construct a basic univariate variable state with some ``value``.

# Examples

```julia
julia> state = BasicUnvVariableState(1.)
Lora.BasicUnvVariableState{Float64}(1.0)

julia> state.value
1.0
```
"""
type BasicUnvVariableState{N<:Number} <: VariableState{Univariate}
  "Scalar value of basic univariate variable state"
  value::N
end

variate_form{N<:Number}(s::Type{BasicUnvVariableState{N}}) = Univariate
variate_form(::BasicUnvVariableState) = Univariate

Base.eltype{N<:Number}(::Type{BasicUnvVariableState{N}}) = N
Base.eltype{N<:Number}(s::BasicUnvVariableState{N}) = N

"""
Basic multivariate variable state type

# Constructors

## BasicMuvVariableState{N<:Number}(value::Vector{N})

Construct a basic multivariate variable state with some ``value``.

## BasicMuvVariableState{N<:Number}(size::Integer, ::Type{N}=Float64)

Construct a basic multivariate variable state with a ``value`` of specified ``size`` and element type.

# Examples

```julia
julia> state = BasicMuvVariableState([1, 2])
Lora.BasicMuvVariableState{Int64}([1,2],2)

julia> state.value
2-element Array{Int64,1}:
 1
 2

julia> state.size
2
```
"""
type BasicMuvVariableState{N<:Number} <: VariableState{Multivariate}
  "Vector value of basic multivariate variable state"
  value::Vector{N}
  "Integer representing the length of vector value of basic multivariate variable state"
  size::Integer
end

BasicMuvVariableState{N<:Number}(value::Vector{N}) = BasicMuvVariableState{N}(value, length(value))

BasicMuvVariableState{N<:Number}(size::Integer, ::Type{N}=Float64) = BasicMuvVariableState{N}(Array(N, size), size)

variate_form{N<:Number}(::Type{BasicMuvVariableState{N}}) = Multivariate
variate_form(::BasicMuvVariableState) = Multivariate

Base.eltype{N<:Number}(::Type{BasicMuvVariableState{N}}) = N
Base.eltype{N<:Number}(::BasicMuvVariableState{N}) = N

"""
Basic matrix-variate variable state type

# Constructors

## BasicMavVariableState{N<:Number}(value::Matrix{N})

Construct a basic matrix-variate variable state with some ``value``.

## BasicMavVariableState{N<:Number}(size::Tuple, ::Type{N}=Float64)

Construct a basic matrix-variate variable state with a ``value`` of specified ``size`` and element type.

# Examples

```julia
julia> state = BasicMavVariableState(eye(2))
Lora.BasicMavVariableState{Float64}(2x2 Array{Float64,2}:
 1.0  0.0
 0.0  1.0,(2,2))

julia> state.value
2x2 Array{Float64,2}:
 1.0  0.0
 0.0  1.0

julia> state.size
(2,2)
```
"""
type BasicMavVariableState{N<:Number} <: VariableState{Matrixvariate}
  "Matrix value of basic matrix-variate variable state"
  value::Matrix{N}
  "Tuple containing the dimensions of matrix value of basic matrix-variate variable state"
  size::Tuple{Integer, Integer}
end

BasicMavVariableState{N<:Number}(value::Matrix{N}) = BasicMavVariableState{N}(value, size(value))

BasicMavVariableState{N<:Number}(size::Tuple, ::Type{N}=Float64) = BasicMavVariableState{N}(Array(N, size...), size)

variate_form{N<:Number}(::Type{BasicMavVariableState{N}}) = Matrixvariate
variate_form(::BasicMavVariableState) = Matrixvariate

Base.eltype{N<:Number}(::Type{BasicMavVariableState{N}}) = N
Base.eltype{N<:Number}(::BasicMavVariableState{N}) = N


text:
using FastTransforms: cheb2leg, leg2cheb
using ApproxFun
#using Plots
Np=Int(1e5)
Nx=Int(25)
k0=0.5
dt=0.1
tmax=8
L=2*pi/k0
deg=0; #Smoothness of periodic boundary

function f(x,v)
  return exp(-0.5.*v.^2)/sqrt(2*pi).*(1+0.05.*cos(k0*x))
end

xk=rand(Np)*L
vk=randn(Np)
fk=f(xk,vk)
wk=f(xk,vk)*L

# Derivative of a chebyshev series
function chebTdx_series{T}(x::T, c::Array{T,1})
  N=length(c)-1;
  b1::T=0
  b2::T=0
  @inbounds for n=N:-1:2
     @inbounds b0::T=c[n+1]+ 2*x*(n+1)/n*b1 - (n+2)/n*b2
    b2=b1; b1=b0;
  end
  return c[2]+ 4*x.*b1 - 3*b2
end
# Derivative of a chebyshev series
function chebTdx_series{T}(x::Array{T,1}, c::Array{T,1})
  N=length(c)-1;
  alpha=Array(T,N-1)
  beta=Array(T,N-1)
  for n=N:-1:2
    alpha[n-1]=2*(n+1)/n
    beta[n-1]=-(n+2)/n
  end
  E=Array(T,length(x))
   @inbounds for pdx=1:length(x)
   b1::T=0
   b2::T=0
   @inbounds for n=N:-1:2
      @inbounds b0::T=c[n+1]+ x[pdx]*alpha[n-1]*b1 + beta[n-1]*b2
     b2=b1; b1=b0;
    end
    E[pdx]=c[2]+ 4*x[pdx].*b1 - 3*b2
  end
  return E
end
# Legendre polynomials
function LegendreP{T}(x::T,N::Int)
  Pn=Array(T, N)
  Pn[1]=1
  Pn[2]=x
  @inbounds for idx=3:N
    n=idx-2
    Pn[idx]= x.*(2*n+1)/(n+1).*Pn[idx-1]- n/(n+1).*Pn[idx-2]
  end
  return Pn
end

@inbounds function LegendreP_series{T}(x::Array{T,1}, u::Array{T,1})
  N::Int=length(u)-1;
  alpha=Array(T,N+1)
  beta=Array(T,N+1)
  #
  for n=0:N
     alpha[n+1]=(2*n+1)/(n+1)
     beta[n+1]=-(n+1)/(n+2)
  end
  E=Array(T,length(x))
   for pdx=1:length(x)
   b1::T=0
   b2::T=0
   for n=N:-1:0
     b0::T=u[n+1]+ x[pdx]*alpha[n+1]*b1 + beta[n+1]*b2
     b2=b1; b1=b0;
    end
    E[pdx]=b1
  end
  return E
end

# #@time LegendreP(xk[1],32)
# function  rhs_particle_legendre{T}(xk::Array{T,1},wk::Array{T,1}, Nx::Int)
#     Np=length(xk)
#     rho=zeros(T,Nx)
#
#     @inbounds for pdx=1:Np
#       P0::T=1.0; P1::T=xk[pdx];
#       rho[1]+=P0.*wk[pdx]
#       rho[2]+=P1.*wk[pdx]
#       @inbounds for idx=3:Nx
#         n=idx-2
#         P2::T=xk[idx].*(2*n+1)/(n+1).*P1- n/(n+1).*P0
#         P0=P1; P1=P2;
#         rho[idx]+=P2.*wk[pdx]
#       end
#     end
#   return rho
# end
function  rhs_particle_legendre{T}(xk::Array{T,1},wk::Array{T,1}, Nx::Int)
    const Np=length(xk)
    rho=zeros(T,Nx)
    # Precalculate coefficients
    alpha=Array(T,Nx-2)
    beta=Array(T,Nx-2)
    @inbounds for n=1:Nx-2
      alpha[n]=(2*n+1)/(n+1)
      beta[n]=-n/(n+1)
    end

    @inbounds for pdx=1:Np
      P0::T=1; P1::T=xk[pdx];
      rho[1]+=P0.*wk[pdx]
      rho[2]+=P1.*wk[pdx]
      @inbounds for idx=1:Nx-2
        P2::T=xk[pdx].*alpha[idx].*P1+ beta[idx].*P0
        P0=P1; P1=P2;
        @inbounds rho[idx+2]+=P2.*wk[pdx]
      end
    end
  return rho
end
function  rhs_particle{T}(xk::Array{T,1},wk::Array{T,1}, Nx::Int)
    Np=length(xk)
    rho=zeros(T,Nx)
    @inbounds for pdx=1:Np
      rho+=LegendreP(xk[pdx],Nx).*wk[pdx]
    end
  return rho
end

OM=Interval(0,L) # Domain
S=ApproxFun.Space(OM) #Default chebyshev space
B=periodic(S,deg) # Periodic boundary
B2=DefiniteIntegral(S) # Nullspace condition
D2=-Derivative(S,2) #Laplace operator
D1=Derivative(S,1) # First derivative for H1 norm
IN=DefiniteIntegral(S) # First derivative for L2 norm

#Phi = \([B;B2;D2],[0.0;0.0;rho];tolerance=1E-16,maxlength=Nx+2)
QR = qrfact([B;B2;D2])

# Coefficients for Runge Kutta
rksd=[2/3.0, -2/3.0, 1  ]
rksc=[ 7/24.0, 3/4.0, -1/24.0]

# Phi=QR \[zeros(Float64,deg); 0.0;0.0;rho]

DE, DV = eigs([B;B2;D2],Nx)
x=collect(linspace(0.0,L))
for idx=1:length(DV)
  DV[idx].coefficients=DV[idx].coefficients[1:Nx]
  DV[idx].coefficients[1]=0
  DV[idx].coefficients=DV[idx].coefficients./norm(DV[idx].coefficients)
  if DE[idx]==0
    DE[idx]=Inf
  end
  #plot(x,map(x->DV[idx](x),x))
end
Nf=2
p = sortperm(abs(DE));DV=DV[p[1:Nf]]; DE=DE[p[1:Nf]]

Nt=Int(ceil(tmax/dt))
fieldenergy=zeros(Float64,Nt)
kineticenergy=zeros(Float64,Nt)
momentum=zeros(Float64,Nt)
rho=Fun(S)
E=Fun(S)
Phi=Fun(S,zeros(Float64,Nx))

tic()
@inbounds for tdx=1:Nt

  @inbounds for rkdx=1:length(rksd)

    # Bin to right hand side
    #rhs=rhs_particle( xk/L*2-1,wk,Nx)/Np*(2/L)
    rhs=rhs_particle_legendre( xk/L*2-1,wk,Nx)/Np*(2/L)
    # Legendre L^2 projection
    rhs=rhs./(2./(2*(0:Nx-1)+1.0))
    # Remove ion background (quick'n dirty)
    rhs[1]=0
    #Transform to Chebyshev
    rhs=leg2cheb(rhs)

    #Eigenvalue filter
    Phi.coefficients=zeros(Float64,Nx)
    for idx=1:length(DV)
    Phi.coefficients.+=(dot(rhs,DV[idx].coefficients)/DE[idx]).*DV[idx].coefficients
    end

        # Solve periodic Poisson equation with Chebyshev
        # Chebyshev periodic Poisson
        #rho.coefficients=rhs
        # Phi = \([B;B2;D2],[0.0;0.0;rho];tolerance=1E-16)
        #Phi=QR \[zeros(Float64,deg);0.0;0.0;rho]

    #E=-rho./(-im*kx)
    if rkdx==1
      fieldenergy[tdx]=coefficients(IN*(D1*Phi)^2)[1]/2
      #*(L/2)
      #fieldenergy[tdx]=dot(Phi.coefficients[1:Nx],rhs)/2*(L/2)/(2/L)
      kineticenergy[tdx]=0.5.*sum(vk.^2.*wk )/Np
      momentum[tdx]=sum(vk.*wk)/Np
    end
    E=Fun(-D1*Phi,S)
    #E.coefficients[1]=0
    E_leg=cheb2leg(E.coefficients)
    vk.+=dt*rksc[rkdx]*LegendreP_series(xk/L*2-1,E_leg)
    xk.+=dt*rksd[rkdx].*vk
    xk=mod(xk,L)
  end

  #print(tdx)
end
toc()
ttime=(0:Nt-1)*dt

#print(ttime)
#plot(Phi)
using PyPlot
#: semilogy,plot,figure,ylabel,xlabel,grid,savefig

prefix="landau_strong"
prefix="legendre"

# figure()
# title("magnitude of coefficients")
# semilogy(abs(coefficients(rho)),label=L"density $\rho$")
# semilogy(abs(coefficients(Phi)),label=L"potential $\Phi$")
# semilogy(abs(coefficients(D1*Phi)),label=L"field $E$")
# xlabel("index")
# legend()
# grid()
# savefig("$prefix\_field_coeffs.png")
# figure()
# semilogy(ttime,kineticenergy)

figure()
semilogy(ttime,fieldenergy)
omega=1.415661888604536 - 0.153359466909605im;
semilogy(ttime[ttime.<8.],
0.5*fieldenergy[1]*abs(exp(-im*omega*(ttime[ttime.<8.]-0.4))).^2,
lw=3,"r")
grid();autoscale(tight="True")
ylabel("electrostatic energy");grid()
xlabel("time")
#tight_layout()
savefig("$prefix\_fieldenergy.png")

figure()
semilogy(ttime,kineticenergy)
grid();
#autoscale(tight="True")
ylabel("kinetic energy")
xlabel("time")
savefig("$prefix\_kineticenergy.png")

figure()
energy=kineticenergy+fieldenergy
semilogy(ttime,abs((energy-energy[1])./energy[1]))
grid()
ylabel("rel. energy error")
xlabel("time ")
savefig("$prefix\_energy_error.png")
figure()
semilogy(ttime,abs(momentum-momentum[1]))
xlabel("time "); grid()
ylabel("absolute momentum error")
savefig("$prefix\_momentum_error.png")

figure()
scatter(xk,vk,s=1, c=fk,lw=0)
#autoscale(tight="True")
axis([0,L,minimum(vk),maximum(vk)])
colorbar();tight_layout()
xlabel("x");ylabel("v")
savefig("$prefix\_particles.png")

# QR*Phi
#Phi
function  rhs_var_particle_legendre{T}(xk::Array{T,1},wk::Array{T,1},
               Nx::Int,   rhs::Array{T,1})
    const Np=length(xk)
    rho=zeros(T,Nx)
    # Precalculate coefficients
    alpha=Array(T,Nx-2)
    beta=Array(T,Nx-2)
    @inbounds for n=1:Nx-2
      alpha[n]=(2*n+1)/(n+1)
      beta[n]=-n/(n+1)
    end

    @inbounds for pdx=1:Np
      P0::T=1; P1::T=xk[pdx];
      rho[1]+=(P0.*wk[pdx]-rhs[1]).^2
      rho[2]+=(P1.*wk[pdx]-rhs[2]).^2
      @inbounds for idx=1:Nx-2
        P2::T=xk[pdx].*alpha[idx].*P1+ beta[idx].*P0
        P0=P1; P1=P2;
        @inbounds rho[idx+2]+=(P2.*wk[pdx]-rhs[idx+2]).^2
      end
    end
  return rho
end

rhs=rhs_particle_legendre( xk/L*2-1,wk,Nx)/Np
rhs_var=rhs_var_particle_legendre( xk/L*2-1,wk,Nx,rhs)/(Np-1)

# Legendre L^2 projection
rhs=rhs./(2./(2*(0:Nx-1)+1.0))*(2/L)
rhs_var=rhs_var./(2./(2*(0:Nx-1)+1.0)).^2*(2/L).^2
rhs[1]=0

figure()
title("Legendre coefficients")
semilogy(abs(rhs),label="magnitude")
plot(sqrt(rhs_var/Np),label="std. deviation")
grid();autoscale(tight="True")
xlabel("index")
legend()
savefig("$prefix\_legendre_coeffs_var.png")

#using JLD
#@save "$prefix\_results.jld"

#
 # figure()
 # errorbar(1:Nx, rhs, yerr=sqrt(rhs_var/Np) , fmt="-o")

#
# rhs=rho.coefficients
#
#
#
# Phi=QR \[0.0;0.0;rho]
#
#
#  coefficients(DV[1])
# #
#
# #
# # using Plots: plot
# #
#
# plot()
# plot(x,map(x->DV[1](x),x))

# coefficients(V[1:2])

# figure()
# semilogy(ttime,fieldenergy)
# grid()
# xlabel("time")
# ylabel("electrostatic energy")
#using Plots
#plot(Phi)

max_stars_repo_name: ameresj/VlasovPIF

max_issues_repo_name: ameresj/VlasovPIF

max_forks_repo_name: ameresj/VlasovPIF


text:
# unary
@i @inline function Base.:-(a!::GVar)
    -(a!.x)
    -(a!.g)
end

@i @inline function DEC(a!::GVar)
    DEC(a!.x)
end

# +-
@i @inline function ⊖(identity)(a!::GVar, b::GVar)
    a!.x -= b.x
    b.g += a!.g
end
@nograd ⊖(identity)(a!::Real, b::GVar)
@nograd ⊖(identity)(a!::GVar, b::Real)

# +- (triple)
@i @inline function ⊖(+)(out!::GVar, x::GVar, y::GVar)
    out!.x -= x.x + y.x
    x.g += out! |> grad
    y.g += out! |> grad
end

@i @inline function ⊖(+)(out!::GVar, x::GVar, y::Real)
    out!.x -= (x |> value) + (y |> value)
    x.g += out! |> grad
end

@i @inline function ⊖(+)(out!::GVar, x::Real, y::GVar)
    out!.x -= (x |> value) + (y |> value)
    y.g += out! |> grad
end

@i @inline function ⊖(-)(out!::GVar, x::GVar, y::GVar)
    out!.x -= x.x - y.x
    x.g += out! |> grad
    y.g -= out! |> grad
end

@i @inline function ⊖(-)(out!::GVar, x::Real, y::GVar)
    out!.x -= (x |> value) - y.x
    y.g -= out!.g
end

@i @inline function ⊖(-)(out!::GVar, x::GVar, y::Real)
    out!.x -= x.x - (y |> value)
    x.g += out! |> grad
end

# NOTE: it will error on `SWAP(a!::GVar, b)` or `SWAP(a!, b:GVar)`
@i @inline function SWAP(a!::GVar, b!::GVar)
    SWAP(a! |> value,  b! |> value)
    SWAP(a!.g, b!.g)
end

# */
@i @inline function ⊖(*)(out!::GVar, x::GVar, y::GVar)
    out!.x -= x.x * y.x
    x.g += out!.g * y.x
    y.g += x.x * out!.g
end

@i @inline function ⊖(*)(out!::GVar, x::Real, y::GVar)
    out!.x -= (x |> value) * y.x
    y.g += (x |> value) * out!.g
end

@i @inline function ⊖(*)(out!::GVar, x::GVar, y::Real)
    out!.x -= x.x * (y |> value)
    x.g += out!.g * (y |> value)
end

for DIV in [:/, :÷]
@eval @i @inline function ⊖($DIV)(out!::GVar{T}, x::GVar, y::GVar) where T
    out!.x -= $DIV(x.x, y.x)
    @routine @invcheckoff begin
        a1 ← zero(out! |> grad)
        a2 ← zero(out! |> grad)
        a1 += x.x * out!.g
        a2 += $DIV(a1, y.x)
    end
    x.g += $DIV(out!.g, y.x)
    y.g -= $DIV(a2, y.x)
    ~@routine
end

@eval @i @inline function ⊖($DIV)(out!::GVar{T}, x::Real, y::GVar) where T
    out!.x -= $DIV(x, y.x)
    @routine @invcheckoff begin
        a1 ← zero(out!.g)
        a2 ← zero(out!.g)
        a1 += x * out!.g
        a2 += $DIV(a1, y.x)
    end
    y.g -= $DIV(a2, y.x)
    ~@routine
end

@eval @i @inline function ⊖($DIV)(out!::GVar, x::GVar, y::Real)
    out!.x -= $DIV(x.x, y)
    x.g += $DIV(out!.g, y)
end
end

@i @inline function ⊖(^)(out!::GVar{T}, x::GVar, n::GVar) where T
    ⊖(^)(out!.x, x.x, n.x)

    # grad x
    @routine @invcheckoff begin
        @zeros T anc1 anc2 anc3 jac1 jac2
        DEC(n.x)
        anc1 += x.x^n.x
        INC(n.x)
        jac1 += anc1 * n.x

        # get grad of n
        anc2 += log(x.x)
        anc3 += x.x ^ n.x
        jac2 += anc3*anc2
    end
    x.g += out!.g * jac1
    n.g += out!.g * jac2
    ~@routine
end

@i @inline function ⊖(^)(out!::GVar{T}, x::GVar, n::Real) where T
    ⊖(^)(out!.x, x.x, n)
    @routine @invcheckoff begin
        anc1 ← zero(x.x)
        jac ← zero(x.x)

        DEC(n |> value)
        anc1 += x.x ^ n
        INC(n |> value)
        jac += anc1 * n
    end
    x.g += out!.g * jac
    ~@routine
end

@i @inline function ⊖(^)(out!::GVar{T}, x::Real, n::GVar) where T
    ⊖(^)(out!.x, x, n.x)
    # get jac of n
    @routine @invcheckoff begin
        anc1 ← zero(x)
        anc2 ← zero(x)
        jac ← zero(x)

        anc1 += log(x)
        anc2 += x ^ n.x
        jac += anc1*anc2
    end
    n.g += out!.g * jac
    ~@routine
end

@i @inline function ⊖(atan)(out!::GVar{T}, y::GVar, x::GVar) where T
    ⊖(atan)(out!.x, y.x, x.x)
    @routine @invcheckoff begin
        @zeros T xy2 jac_x jac_y
        xy2 += abs2(x.x)
        xy2 += abs2(y.x)
        jac_y += x.x / xy2
        jac_x += (-y.x) / xy2
    end
    y.g += out!.g * jac_y
    x.g += out!.g * jac_x
    ~@routine
end

@i @inline function ⊖(atan)(out!::GVar{T}, x::GVar) where T
    ⊖(atan)(out!.x, x.x)
    @routine @invcheckoff begin
        xy2 ← one(T)
        xy2 += abs2(x.x)
    end
    x.g += out!.g / xy2
    ~@routine
end

@i @inline function ⊖(abs)(out!::GVar, x::GVar{T}) where T
    out!.x -= abs(x.x)
    if (x > 0, ~)
        x.g += out!.g
    else
        x.g -= out!.g
    end
end

@i @inline function ⊖(abs2)(out!::GVar, x::GVar{T}) where T
    out!.x -= abs2(x.x)
    x.g += out!.g * x.x
    x.g += out!.g * x.x
end
@nograd ⊖(abs2)(a!::GVar, b::Real)
@nograd ⊖(abs2)(a!::Real, b::GVar)

for op in [:*, :/, :^, :+, :-]
    @eval @nograd ⊖($op)(out!::GVar, x::Real, y::Real)
    @eval @nograd ⊖($op)(out!::Real, x::Real, y::GVar)
    @eval @nograd ⊖($op)(out!::Real, x::GVar, y::GVar)
    @eval @nograd ⊖($op)(out!::Real, x::GVar, y::Real)
end

@i @inline function ⊖(sqrt)(out!::GVar, x::GVar{T}) where T
    @routine @invcheckoff begin
        @zeros T anc1 anc2
        anc1 += sqrt(x.x)
        anc2 += 2 * anc1
    end
    out!.x -= anc1
    x.g += out!.g / anc2
    ~@routine
end

@i @inline function ⊖(exp)(out!::GVar, x::GVar{T}) where T
    @routine @invcheckoff begin
        anc1 ← zero(T)
        anc1 += exp(x.x)
    end
    out!.x -= anc1
    x.g += out!.g * anc1
    ~@routine
end

@i @inline function ⊖(log)(out!::GVar, x::GVar{T}) where T
    out!.x -= log(x.x)
    x.g += out!.g / x.x
end

@i @inline function ⊖(sin)(out!::GVar, x::GVar{T}) where T
    out!.x -= sin(x.x)
    @routine @invcheckoff begin
        anc1 ← zero(x.x)
        anc1 += cos(x.x)
    end
    x.g += out!.g * anc1
    ~@routine
end

@i @inline function ⊖(cos)(out!::GVar, x::GVar{T}) where T
    out!.x -= cos(x.x)
    @routine @invcheckoff begin
        anc1 ← zero(x.x)
        anc1 -= sin(x.x)
    end
    x.g += out!.g * anc1
    ~@routine
end

@i @inline function ⊖(tanh)(out!::GVar, x::GVar{T}) where T
    @routine @invcheckoff begin
        anc1 ← zero(x.x)
        anc2 ← one(x.x)
        anc1 += tanh(x.x)
        anc2 -= anc1^2
    end
    out!.x -= anc1
    x.g += out!.g * anc2
    ~@routine
end

@i @inline function ⊖(sincos)(out!::Tuple{T1,T1}, x::GVar{T}) where {T1<:GVar, T}
    @routine @invcheckoff begin
        s ← zero(T)
        c ← zero(T)
        (s, c) += sincos(x.x)
    end
    (out! .|> value) -= (s, c)
    x.g += (out! |> tget(1) |> grad) * c
    x.g -= (out! |> tget(2) |> grad) * s
    ~@routine
end

for op in [:sqrt, :exp, :log, :sin, :cos, :tanh]
    @eval @nograd ⊖($op)(out!::Real, x::GVar)
    @eval @nograd ⊖($op)(out!::GVar, x::Real)
end

@nograd ⊖(sincos)(out!::Tuple{<:Real,<:Real}, x::GVar)
@nograd ⊖(sincos)(out!::Tuple{<:GVar,<:GVar}, x::Real)

@i @inline function IROT(a!::GVar, b!::GVar, θ::GVar)
    IROT(a!.x, b!.x, θ.x)
    -(θ |> value)
    θ.x -= π/2
    ROT(a!.g, b!.g, θ.x)
    θ.g += a!.x * a!.g
    θ.g += b!.x * b!.g
    θ.x += π/2
    -(θ |> value)
    ROT(a!.g, b!.g, π/2)
end

@i @inline function IROT(a!::GVar, b!::GVar, θ::Real)
    IROT(a!.x, b!.x, θ)
    -(θ)
    θ -= π/2
    ROT(a!.g, b!.g, θ)
    θ += π/2
    -(θ)
    ROT(a!.g, b!.g, π/2)
end

@nograd IROT(a!::Real, b!::Real, θ::GVar)

export primitive_grad
function primitive_grad end

@i function (mf::MinusEq)(out!::GVar, args...; kwargs...)
    out!.x -= mf.f((args .|> value)...; kwargs...)
    (args .|> grad) .+= (@skip! out!.g) .* primitive_grad(mf.f, (args .|> value)...; kwargs...)
end

@i function (mf::MinusEq)(out!::GVar, x::GVar; kwargs...)
    out!.x -= mf.f(x .|> value; kwargs...)
    x.g += (@skip! out!.g) * primitive_grad(mf.f, x.x; kwargs...)
end

@i function :(-=)(convert)(out!::GVar{Tx, Tg}, y::GVar) where {Tx, Tg}
    out!.x -= convert(y.x)
    y.g += convert(out!.g)
end


text:
export draw_basin!, basins_of_attraction

"""
    basins_of_attraction(grid::Tuple, ds::DynamicalSystem; kwargs...) -> basins, attractors
Compute an estimate of the basins of attraction of a dynamical system `ds` on
a partitioning of the state space given by `grid`. The method has been
inspired by the 2D grid approach developed by Nusse & Yorke [^Yorke1997].
It works _without_ knowledge of where attractors are; it identifies them automatically.

The dynamical system can be:
* An actual `DiscreteDynamicalSystem` or `ContinuousDynamicalSystem`.
* A Poincaré map of `ContinuousDynamicalSystem`: [`poincaremap`](@ref).
* A stroboscopic map, i.e. a periodically forced `ContinuousDynamicalSystem` (see examples
  for this particular application).

`grid` is a tuple of ranges defining the grid of initial conditions, for example
`grid = (xg, yg)` where `xg = yg = range(-5, 5; length = 100)`. The grid is not necessarilly
of the same dimension as the state space, attractors can be found in lower dimensional
projections.

The output `basins` is an integer-valued array on the `grid`, with its entries labelling
which basin of attraction the given grid point belongs to.
The output `attractors` is a dictionary whose keys correspond to the attractor number and
the values contains the points of the attractors found.
Notice that for some attractors this list may be incomplete.

See also [`match_attractors!`](@ref), [`basin_fractions`](@ref), [`tipping_probabilities`](@ref).

[^Yorke1997]: H. E. Nusse and J. A. Yorke, Dynamics: numerical explorations Ch. 7, Springer, New York, 1997

## Keyword Arguments
* `Δt = 1`: Approximate time step of the integrator. It is recommended to use values such
  that one step will typically make the integrator move to a different cell of the
  state space partitioning.
* `T` : Period of the stroboscopic map, in case of a continuous dynamical system with periodic
  time forcing. This argument is incompatible with `Δt`.
* `idxs = 1:length(grid)`: This vector selects the variables of the system that will define the
  subspace the dynamics will be projected into.
* `complete_state = zeros(D-Dg)`: This argument allows setting the _remaining_ variables
  of the dynamical system state on each initial condition `u`, with `Dg` the dimension
  of the grid. It can be either a vector of length `D-Dg`, or a function `f(y)` that
  returns a vector of length `D-Dg` given the _projected_ initial condition on the grid `y`.
* `diffeq...`: Keyword arguments propagated to [`integrator`](@ref).
* `mx_chk_att = 2`: A parameter that sets the maximum checks of consecutives hits of an attractor
  before deciding the basin of the initial condition.
* `mx_chk_hit_bas = 10` : Maximum check of consecutive visits of the same basin of attraction.
  This number can be increased for higher accuracy.
* `mx_chk_fnd_att = 100` : Maximum check of unnumbered cell before considering we have an attractor.
  This number can be increased for higher accuracy.
* `mx_chk_loc_att = 60` : Maximum check of consecutive cells marked as an attractor before considering
  that we have all the available pieces of the attractor.
* `mx_chk_lost` : Maximum check of iterations outside the defined grid before we consider the orbit
  lost outside. This number can be increased for higher accuracy. It defaults to `20` if no
  attractors are given (see discussion on refining basins), and to `1000` if attractors are given.
* `horizon_limit = 1e6` : If the norm of the integrator state reaches this limit we consider that the
  orbit diverges.
* `show_progress = true` : By default a progress bar is shown using ProgressMeter.jl.
* `attractors, ε`: See discussion on refining basins below.

## Description
`basins` has the following organization:
* The basins are coded in sequential order from 1 up to the number of attractors.
* If the trajectory diverges or converges to an attractor outside the defined grid it is
  numbered `-1`

`attractors` has the following organization:
* The keys of the dictionary correspond to the number of the attractor.
* The value associated to a key is a [`Dataset`](@ref) with the *guessed* location of the
  attractor on the state space.

The method starts by picking the first available initial condition on the grid not yet
numbered. The dynamical system is then iterated until one of the following conditions happens:
1. The trajectory hits a known attractor already numbered `mx_chk_att` consecutive times: the
   initial condition is numbered with the corresponding number.
1. The trajectory spends `mx_chk_lost` steps outside the defiend grid: the initial
   condition is set to -1.
1. The trajectory hits a known basin `mx_chk_hit_bas` times in a row: the initial condition
   belongs to that basin and is numbered accordingly.
1. The trajectory hits `mx_chk_fnd_att` times in a row an unnumbered cell: it is considered
   an attractor and is labelled with a new number.

Regarding performace, this method is at worst as fast as tracking the attractors.
In most cases there is a signicative improvement in speed.

Notice that in the case we have to project the dynamics on a lower dimensional space,
there are edge cases where the system may have two attractors
that are close on the projected space but are far apart in another dimension. They could
be collapsed or confused into the same attractor. This is a drawback of this method.

## Refining basins of attraction
Sometimes one would like to be able to refine the found basins of attraction by recomputing
`basins_of_attraction` on a smaller, and more fine-grained, `grid`. If however this
new `grid` does not contain the attractors, `basins_of_attraction` would (by default)
attribute the value `-1` to all grid points. For these cases, an extra search clause can
be provided by setting the keywords `attractors, ε`. The `attractors` is a dictionary
mapping attractor IDs to `Dataset`s (i.e., the same as the return value of
`basins_of_attraction`). The algorithm checks at each step whether the system state is
`ε`-close (Euclidean norm) to any of the given attractors, and if so it attributes the stating grid point
to the basin of the close attractor. By default `ε` is equal to the mean grid spacing.

A word of advice while using this method: in order to work properly, `ε` should be
about the size of a grid cell that has been used to compute the given `attractors`. It is
recomended to keep the same step size (i.e., use the same integrator) since it may have an
influence in some cases. This algorithm is usually slower than the method with the
attractors on the grid.
"""
function basins_of_attraction(grid::Tuple, ds;
        Δt=1, T=0, idxs = 1:length(grid),
        complete_state = zeros(eltype(get_state(ds)), length(get_state(ds)) - length(grid)),
        diffeq = NamedTuple(), kwargs...
        # `kwargs` tunes the basin finding algorithm, e.g. `mx_chk_att`.
        # these keywords are actually expanded in `_identify_basin_of_cell!`
    )
    @assert length(idxs) == length(grid)
    integ = ds isa PoincareMap ? ds : integrator(ds; diffeq...)
    idxs = SVector(idxs...)
    return basins_of_attraction(grid, integ, Δt, T, idxs, complete_state; kwargs...)
end

function basins_of_attraction(grid, integ, Δt, T, idxs::SVector, complete_state; kwargs...)
    D = length(get_state(integ))
    if complete_state isa AbstractVector && (length(complete_state) ≠ D-length(idxs))
        error("Vector `complete_state` must have length D-Dg!")
    end
    if T > 0
        iter_f! = (integ) -> step!(integ, T, true)
    elseif integ isa PoincareMap
        iter_f! = step!
    else # generic case
        iter_f! = (integ) -> step!(integ, Δt) # we don't have to step _exactly_ `Δt` here
    end
    complete_and_reinit! = CompleteAndReinit(complete_state, idxs, length(get_state(integ)))
    get_projected_state = (integ) -> view(get_state(integ), idxs)
    bsn_nfo = draw_basin!(
        grid, integ, iter_f!, complete_and_reinit!, get_projected_state; kwargs...
    )
    return bsn_nfo.basin, bsn_nfo.attractors
end

"""
    CompleteAndReinit(complete_state, idxs, D)
Helper struct that completes a state and reinitializes the integrator once called
as a function with arguments `f(integ, y)` with `integ` the initialized dynamical
system integrator and `y` the projected initial condition on the grid.
"""
struct CompleteAndReinit{C, Y, R}
    complete_state::C
    u::Vector{Float64} # dummy variable for a state in full state space
    idxs::SVector{Y, Int}
    remidxs::R
end
function CompleteAndReinit(complete_state, idxs, D::Int)
    remidxs = setdiff(1:D, idxs)
    remidxs = isempty(remidxs) ? nothing : SVector(remidxs...)
    u = zeros(D)
    if complete_state isa AbstractVector
        @assert eltype(complete_state) <: Number
    end
    return CompleteAndReinit(complete_state, u, idxs, remidxs)
end
function (c::CompleteAndReinit{<: AbstractVector})(integ, y)
    c.u[c.idxs] .= y
    if !isnothing(c.remidxs)
        c.u[c.remidxs] .= c.complete_state
    end
    reinit!(integ, c.u)
end
function (c::CompleteAndReinit)(integ, y) # case where `complete_state` is a function
    c.u[c.idxs] .= y
    if !isnothing(c.remidxs)
        c.u[c.remidxs] .= c.complete_state(y)
    end
    reinit!(integ, c.u)
end


text:
using FastRationals, LinearAlgebra, BenchmarkTools

BenchmarkTools.DEFAULT_PARAMETERS.evals = 1;
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100;
BenchmarkTools.DEFAULT_PARAMETERS.time_tolerance = 1.0e-8;
BenchmarkTools.DEFAULT_PARAMETERS.overhead = BenchmarkTools.estimate_overhead();

function Base.rand(::Type{Rational{I}}, n::Int=1) where {I<:Base.BitInteger}
    n = max(1, n)
    numers = rand(I, n)
    denoms = map(x->max(1,x), abs.(rand(I, n)))
    return numers .// denoms
end;

function Base.rand(::Type{Rational{BigInt}}, n::Int=1)
    n = max(1, n)
    qs = rand(Rational{Int128}, n)
    return Rational{BigInt}.(qs)
end;

q200 = rand(Rational{BigInt}, 200);
fq200 = FastRational{BigInt}.(q200);

q500 = rand(Rational{BigInt}, 500);
fq500 = FastRational{BigInt}.(q500);

q200sum = @belapsed sum($q200);
fq200sum = @belapsed sum($fq200);
q200prod = @belapsed prod($q200);
fq200prod = @belapsed prod($fq200);

q500sum = @belapsed sum($q500);
fq500sum = @belapsed sum($fq500);
q500prod = @belapsed prod($q500);
fq500prod = @belapsed prod($fq500);

relspeed_200sum =
  round( (q200sum / fq200sum), digits=1);
relspeed_200prod =
  round( (q200prod / fq200prod), digits=1);

relspeed_500sum =
  round( (q500sum / fq500sum), digits=1);
relspeed_500prod =
  round( (q500prod / fq500prod), digits=1);

q8x8 = reshape(q200[1:64], (8,8));
fq8x8 = reshape(fq200[1:64], (8,8));
q15x15 = reshape(q500[1:225], (15,15));
fq15x15 = reshape(fq500[1:225], (15,15));

q8x8mul = @belapsed $q8x8 * $q8x8;
fq8x8mul = @belapsed $fq8x8 * $fq8x8;
q15x15mul = @belapsed $q15x15 * $q15x15;
fq15x15mul = @belapsed $fq15x15 * $fq15x15;

q8x8tr = @belapsed tr($q8x8);
fq8x8tr = @belapsed tr($fq8x8);
q15x15tr = @belapsed tr($q15x15);
fq15x15tr = @belapsed tr($fq15x15);

relspeed_8x8mul =
  round( (q8x8mul / fq8x8mul), digits=1);
relspeed_8x8tr =
  round( (q8x8tr / fq8x8tr), digits=1);

relspeed_15x15mul =
  round( (q15x15mul / fq15x15mul), digits=1);
relspeed_15x15tr =
  round( (q15x15tr / fq15x15tr), digits=1);

relspeeds = string(
"\n\n\trelative speeds",
"\n\t (200)\t (500)\n\n",
"sum:   \t $relspeed_200sum \t $relspeed_500sum \n",
"prod:  \t $relspeed_200prod \t $relspeed_500prod \n");

print(relspeeds);

relspeeds = string(
"\n\n\trelative speeds",
"\n\t (8x8)\t (15x15)\n\n",
"matmul:\t $relspeed_8x8mul \t $relspeed_15x15mul\n",
"mat tr:\t $relspeed_8x8tr \t $relspeed_15x15tr \n");

print(relspeeds);
;


text:
function vec_to_vech{S<:Int}(n::S)

  x = zeros(Int,round(Int,(n*(n+1)/2)),n*n)

  c = 1
  r = 1
  counter = 1

  for i = 1:n*n

    if r >= c
      x[counter,i] = 1
      counter += 1
    end

    r += 1
    if r > n
      c += 1
      r = 1
    end

  end

  return x

end


text:
nt = (a=5, b=[(a=20, b=2), (a=33, b=0), (a=44, b=3.3)], c=(a=2, b=[1, 2.0]))
ms = MStruct(a=5, b=[(a=20, b=2), (a=33, b=0), (a=44, b=3.3)], c=(a=2, b=[1, 2.0]))

@testset "Construction" begin
    @test MStruct(nt) == ms
    @test MStruct{Float64}(nt) == ms
    @test eltype(MStruct{ForwardDiff.Dual}(nt)) == ForwardDiff.Dual
end

@testset "Attributes" begin
    @test length(ms) == 10
    @test size(ms) == (10,)
    @test propertynames(ms) == (:a, :b, :c)

    # Make sure MStruct internal fields aren't accessed when keys have same name
    @testset "With reserved keywords" begin
        with_reserved = MStruct{Int64}(a=1, data=[2,3])
        with_reserved.data[2] = 4
        @test with_reserved.data == [2,4]
        @test propertynames(with_reserved) == (:a, :data)
    end
end

@testset "Similarity" begin
    @test typeof(similar(ms)) == typeof(ms)
    @test typeof(similar(ms, Float32)) == typeof(MStruct{Float32}(nt))
    @test eltype(similar(ms, Float32)) == Float32
    @test eltype(similar(ms, ForwardDiff.Dual)) == ForwardDiff.Dual
end

@testset "Conversion" begin
    nt_homogeneous = (a=5., b=[(a=20., b=2.), (a=33., b=0.), (a=44., b=3.3)], c=(a=2., b=[1., 2.]))
    @test NamedTuple(ms) == nt_homogeneous
end

@testset "Set/Get" begin
    ms2 = deepcopy(ms)

    ms2.a = 20
    @test ms2.a == 20

    ms2.c.a = 50
    @test ms2.c.a == 50

    ms2.b[1].b = 30
    @test ms2.b[1].b == 30

    @test ms[:b][1]["a"] == 20
    @test collect(ms) == Float64[5, 20, 2, 33, 0, 44, 3.3, 2, 1, 2]
    @test ms[:] == collect(ms)
    @test ms[2:5] == Float64[20, 2, 33, 0]
    @test ms[end] == 2
end

max_stars_repo_name: jonniedie/NamedViewVectors


text:
module ThreadPools

export @bthreads, @qthreads, @qbthreads
export @logthreads, @logbthreads, @logqthreads, @logqbthreads
export pwith, @pthreads, results, isactive, @tspawnat
export dumplog, readlog, showactivity, showstats
export tmap, bmap, qmap, qbmap
export logtmap, logbmap, logqmap, logqbmap
export tforeach, bforeach, qforeach, qbforeach
export logtforeach, logbforeach, logqforeach, logqbforeach

include("interface.jl")
include("logs.jl")
include("staticpool.jl")
include("qpool.jl")
include("logstaticpool.jl")
include("logqpool.jl")
include("macros.jl")
include("simplefuncs.jl")

end # module


text:
using Aqua

Aqua.test_all(GeometricAlgebra, ambiguities=false, unbound_args=false)


text:
using NDVoronoi
using Documenter

makedocs(;
    modules=[NDVoronoi],
    authors="Lukas Botsch <lukas.botsch@gmail.com>",
    repo="https://github.com/lbotsch/NDVoronoi.jl/blob/{commit}{path}#L{line}",
    sitename="NDVoronoi.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://lbotsch.github.io/NDVoronoi.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/lbotsch/NDVoronoi.jl",
)